structured_asn1/
limits.rs1use core::cell::Cell;
8
9use crate::error::Limit;
10use crate::memory::Memory;
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub struct Limits {
19 pub max_depth: u32,
24
25 pub max_elements: u64,
27
28 pub max_element_length: u64,
30
31 pub max_work: u64,
39}
40
41impl Default for Limits {
42 fn default() -> Self {
43 Self {
44 max_depth: 1_000,
45 max_elements: 1_000_000,
46 max_element_length: 1 << 32,
47 max_work: 16_000_000,
48 }
49 }
50}
51
52impl Limits {
53 #[must_use]
56 pub fn none() -> Self {
57 Self {
58 max_depth: 0,
59 max_elements: 0,
60 max_element_length: 0,
61 max_work: 0,
62 }
63 }
64}
65
66#[derive(Debug)]
73pub(crate) struct Budget {
74 limits: Limits,
75 elements: Cell<u64>,
76 work: Cell<u64>,
77 memory: Memory,
78}
79
80#[derive(Debug, Clone, Copy)]
83pub(crate) struct BudgetView<'a> {
84 origin: &'a Budget,
85 operation: Option<&'a Budget>,
86}
87
88impl<'a> BudgetView<'a> {
89 pub(crate) const fn root(origin: &'a Budget) -> Self {
90 Self {
91 origin,
92 operation: None,
93 }
94 }
95 pub(crate) fn query(origin: &'a Budget, operation: &'a Budget) -> Self {
96 Self {
97 origin,
98 operation: if core::ptr::eq(origin, operation) {
99 None
100 } else {
101 Some(operation)
102 },
103 }
104 }
105 pub(crate) fn operation(self) -> &'a Budget {
106 self.operation.unwrap_or(self.origin)
107 }
108 pub(crate) fn work_in(self, amount: u64, execution: &Memory) -> Result<(), Limit> {
109 if amount > self.origin.limits.max_work - self.origin.work.get() {
110 return Err(Limit::Work);
111 }
112 if let Some(operation) = self.operation
113 && amount > operation.limits.max_work - operation.work.get()
114 {
115 return Err(Limit::Work);
116 }
117 execution.charge_work(amount).map_err(|_| Limit::Work)?;
118 self.origin.work.set(self.origin.work.get() + amount);
119 if let Some(operation) = self.operation {
120 operation.work.set(operation.work.get() + amount);
121 }
122 Ok(())
123 }
124 pub(crate) fn local_work(self, amount: u64) -> Result<(), Limit> {
125 if amount > self.origin.limits.max_work - self.origin.work.get() {
126 return Err(Limit::Work);
127 }
128 if let Some(operation) = self.operation
129 && amount > operation.limits.max_work - operation.work.get()
130 {
131 return Err(Limit::Work);
132 }
133 self.origin.work.set(self.origin.work.get() + amount);
134 if let Some(operation) = self.operation {
135 operation.work.set(operation.work.get() + amount);
136 }
137 Ok(())
138 }
139 pub(crate) fn element(self) -> Result<(), Limit> {
140 if self.origin.elements.get() >= self.origin.limits.max_elements {
141 return Err(Limit::Elements);
142 }
143 if let Some(operation) = self.operation
144 && operation.elements.get() >= operation.limits.max_elements
145 {
146 return Err(Limit::Elements);
147 }
148 self.origin.elements.set(self.origin.elements.get() + 1);
149 if let Some(operation) = self.operation {
150 operation.elements.set(operation.elements.get() + 1);
151 }
152 Ok(())
153 }
154 pub(crate) fn depth(self, depth: u32) -> Result<(), Limit> {
155 self.origin.depth(depth)?;
156 if let Some(operation) = self.operation {
157 operation.depth(depth)?;
158 }
159 Ok(())
160 }
161 pub(crate) fn length(self, length: u64) -> Result<(), Limit> {
162 self.origin.length(length)?;
163 if let Some(operation) = self.operation {
164 operation.length(length)?;
165 }
166 Ok(())
167 }
168}
169
170impl Budget {
171 pub(crate) fn new(limits: Limits, memory: Memory) -> Self {
172 Self {
173 limits,
174 elements: Cell::new(0),
175 work: Cell::new(0),
176 memory,
177 }
178 }
179
180 pub(crate) fn limits(&self) -> Limits {
181 self.limits
182 }
183
184 pub(crate) fn memory(&self) -> &Memory {
185 &self.memory
186 }
187
188 #[cfg(test)]
193 pub(crate) fn element(&self) -> Result<(), Limit> {
194 BudgetView::root(self).element()
195 }
196
197 #[cfg(test)]
199 pub(crate) fn work(&self, amount: u64) -> Result<(), Limit> {
200 BudgetView::root(self).work_in(amount, &self.memory)
201 }
202
203 pub(crate) fn depth(&self, depth: u32) -> Result<(), Limit> {
205 if depth > self.limits.max_depth {
206 return Err(Limit::Depth);
207 }
208 Ok(())
209 }
210
211 pub(crate) fn length(&self, length: u64) -> Result<(), Limit> {
213 if length > self.limits.max_element_length {
214 return Err(Limit::ElementLength);
215 }
216 Ok(())
217 }
218}
219
220#[cfg(test)]
221mod tests;