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_remaining(self) -> u64 {
126 let origin = self.origin.limits.max_work - self.origin.work.get();
127 match self.operation {
128 Some(operation) => origin.min(operation.limits.max_work - operation.work.get()),
129 None => origin,
130 }
131 }
132 pub(crate) fn local_work(self, amount: u64) -> Result<(), Limit> {
133 if amount > self.origin.limits.max_work - self.origin.work.get() {
134 return Err(Limit::Work);
135 }
136 if let Some(operation) = self.operation
137 && amount > operation.limits.max_work - operation.work.get()
138 {
139 return Err(Limit::Work);
140 }
141 self.origin.work.set(self.origin.work.get() + amount);
142 if let Some(operation) = self.operation {
143 operation.work.set(operation.work.get() + amount);
144 }
145 Ok(())
146 }
147 pub(crate) fn element(self) -> Result<(), Limit> {
148 if self.origin.elements.get() >= self.origin.limits.max_elements {
149 return Err(Limit::Elements);
150 }
151 if let Some(operation) = self.operation
152 && operation.elements.get() >= operation.limits.max_elements
153 {
154 return Err(Limit::Elements);
155 }
156 self.origin.elements.set(self.origin.elements.get() + 1);
157 if let Some(operation) = self.operation {
158 operation.elements.set(operation.elements.get() + 1);
159 }
160 Ok(())
161 }
162 pub(crate) fn depth(self, depth: u32) -> Result<(), Limit> {
163 self.origin.depth(depth)?;
164 if let Some(operation) = self.operation {
165 operation.depth(depth)?;
166 }
167 Ok(())
168 }
169 pub(crate) fn length(self, length: u64) -> Result<(), Limit> {
170 self.origin.length(length)?;
171 if let Some(operation) = self.operation {
172 operation.length(length)?;
173 }
174 Ok(())
175 }
176}
177
178impl Budget {
179 pub(crate) fn new(limits: Limits, memory: Memory) -> Self {
180 Self {
181 limits,
182 elements: Cell::new(0),
183 work: Cell::new(0),
184 memory,
185 }
186 }
187
188 pub(crate) fn limits(&self) -> Limits {
189 self.limits
190 }
191
192 pub(crate) fn memory(&self) -> &Memory {
193 &self.memory
194 }
195
196 #[cfg(test)]
201 pub(crate) fn element(&self) -> Result<(), Limit> {
202 BudgetView::root(self).element()
203 }
204
205 #[cfg(test)]
207 pub(crate) fn work(&self, amount: u64) -> Result<(), Limit> {
208 BudgetView::root(self).work_in(amount, &self.memory)
209 }
210
211 pub(crate) fn depth(&self, depth: u32) -> Result<(), Limit> {
213 if depth > self.limits.max_depth {
214 return Err(Limit::Depth);
215 }
216 Ok(())
217 }
218
219 pub(crate) fn length(&self, length: u64) -> Result<(), Limit> {
221 if length > self.limits.max_element_length {
222 return Err(Limit::ElementLength);
223 }
224 Ok(())
225 }
226}
227
228#[cfg(test)]
229mod tests;