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 work_with(
127 self,
128 sources: [&Budget; 2],
129 amount: u64,
130 execution: &Memory,
131 ) -> Result<(), Limit> {
132 let mut extras = [None; 2];
133 for (i, extra) in sources.into_iter().enumerate() {
134 if core::ptr::eq(extra, self.origin) {
135 continue;
136 }
137 if let Some(operation) = self.operation
138 && core::ptr::eq(extra, operation)
139 {
140 continue;
141 }
142 if i == 1 && core::ptr::eq(extra, sources[0]) {
143 continue;
144 }
145 if amount > extra.limits.max_work - extra.work.get() {
146 return Err(Limit::Work);
147 }
148 extras[i] = Some(extra);
149 }
150 self.work_in(amount, execution)?;
151 for extra in extras.into_iter().flatten() {
152 extra.work.set(extra.work.get() + amount);
153 }
154 Ok(())
155 }
156 pub(crate) fn work_remaining(self) -> u64 {
158 let origin = self.origin.limits.max_work - self.origin.work.get();
159 match self.operation {
160 Some(operation) => origin.min(operation.limits.max_work - operation.work.get()),
161 None => origin,
162 }
163 }
164 pub(crate) fn local_work(self, amount: u64) -> Result<(), Limit> {
165 if amount > self.origin.limits.max_work - self.origin.work.get() {
166 return Err(Limit::Work);
167 }
168 if let Some(operation) = self.operation
169 && amount > operation.limits.max_work - operation.work.get()
170 {
171 return Err(Limit::Work);
172 }
173 self.origin.work.set(self.origin.work.get() + amount);
174 if let Some(operation) = self.operation {
175 operation.work.set(operation.work.get() + amount);
176 }
177 Ok(())
178 }
179 pub(crate) fn element(self) -> Result<(), Limit> {
180 if self.origin.elements.get() >= self.origin.limits.max_elements {
181 return Err(Limit::Elements);
182 }
183 if let Some(operation) = self.operation
184 && operation.elements.get() >= operation.limits.max_elements
185 {
186 return Err(Limit::Elements);
187 }
188 self.origin.elements.set(self.origin.elements.get() + 1);
189 if let Some(operation) = self.operation {
190 operation.elements.set(operation.elements.get() + 1);
191 }
192 Ok(())
193 }
194 pub(crate) fn depth(self, depth: u32) -> Result<(), Limit> {
195 self.origin.depth(depth)?;
196 if let Some(operation) = self.operation {
197 operation.depth(depth)?;
198 }
199 Ok(())
200 }
201 pub(crate) fn length(self, length: u64) -> Result<(), Limit> {
202 self.origin.length(length)?;
203 if let Some(operation) = self.operation {
204 operation.length(length)?;
205 }
206 Ok(())
207 }
208}
209
210impl Budget {
211 pub(crate) fn new(limits: Limits, memory: Memory) -> Self {
212 Self {
213 limits,
214 elements: Cell::new(0),
215 work: Cell::new(0),
216 memory,
217 }
218 }
219
220 pub(crate) fn limits(&self) -> Limits {
221 self.limits
222 }
223
224 pub(crate) fn memory(&self) -> &Memory {
225 &self.memory
226 }
227
228 #[cfg(test)]
233 pub(crate) fn element(&self) -> Result<(), Limit> {
234 BudgetView::root(self).element()
235 }
236
237 #[cfg(test)]
239 pub(crate) fn work(&self, amount: u64) -> Result<(), Limit> {
240 BudgetView::root(self).work_in(amount, &self.memory)
241 }
242
243 pub(crate) fn depth(&self, depth: u32) -> Result<(), Limit> {
245 if depth > self.limits.max_depth {
246 return Err(Limit::Depth);
247 }
248 Ok(())
249 }
250
251 pub(crate) fn length(&self, length: u64) -> Result<(), Limit> {
253 if length > self.limits.max_element_length {
254 return Err(Limit::ElementLength);
255 }
256 Ok(())
257 }
258}
259
260#[cfg(test)]
261mod tests;