1use std::sync::Arc;
4use std::sync::atomic::{AtomicBool, Ordering};
5#[cfg(not(target_arch = "wasm32"))]
6use std::time::Duration;
7use std::time::Instant;
8
9use serde::{Deserialize, Serialize};
10
11use crate::error::{EngineError, ErrorCode};
12
13#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
16#[serde(default)]
17pub struct Limits {
18 pub max_request_bytes: usize,
19 pub max_expression_length: usize,
20 pub max_expression_tokens: usize,
21 pub max_ast_depth: usize,
22 pub max_digits: usize,
23 pub max_integer_bits: u32,
24 pub max_decimal_scale: u32,
25 pub max_array_len: usize,
26 pub max_matrix_elements: usize,
27 pub max_batch_nodes: usize,
28 pub max_output_bytes: usize,
29 pub max_trace_bytes: usize,
30 pub max_exponent: u32,
31 pub max_factorial: u32,
32 pub max_iterations: u64,
33 pub max_operations: u64,
34 pub max_recursion_depth: usize,
35 pub max_string_len: usize,
36 pub max_series_terms: u64,
37 pub request_timeout_ms: u64,
38}
39
40impl Default for Limits {
41 fn default() -> Self {
42 Limits::conservative()
43 }
44}
45
46impl Limits {
47 pub fn conservative() -> Limits {
49 Limits {
50 max_request_bytes: 1 << 20,
51 max_expression_length: 16 << 10,
52 max_expression_tokens: 4096,
53 max_ast_depth: 64,
54 max_digits: 4096,
55 max_integer_bits: 8192,
56 max_decimal_scale: 4096,
57 max_array_len: 100_000,
58 max_matrix_elements: 250_000,
59 max_batch_nodes: 256,
60 max_output_bytes: 4 << 20,
61 max_trace_bytes: 256 << 10,
62 max_exponent: 100_000,
63 max_factorial: 10_000,
64 max_iterations: 1_000_000,
65 max_operations: 10_000_000,
66 max_recursion_depth: 64,
67 max_string_len: 1 << 20,
68 max_series_terms: 1_000_000,
69 request_timeout_ms: 30_000,
70 }
71 }
72
73 pub fn lowered_by(&self, overrides: &LimitsOverride) -> Result<Limits, EngineError> {
77 let mut out = self.clone();
78 macro_rules! lower {
79 ($field:ident) => {
80 if let Some(value) = overrides.$field {
81 if value > self.$field {
82 return Err(EngineError::new(
83 ErrorCode::ResourceLimit,
84 format!(
85 "requested {} ({}) exceeds server policy ({})",
86 stringify!($field),
87 value,
88 self.$field
89 ),
90 ));
91 }
92 out.$field = value;
93 }
94 };
95 }
96 lower!(max_request_bytes);
97 lower!(max_expression_length);
98 lower!(max_expression_tokens);
99 lower!(max_ast_depth);
100 lower!(max_digits);
101 lower!(max_integer_bits);
102 lower!(max_decimal_scale);
103 lower!(max_array_len);
104 lower!(max_matrix_elements);
105 lower!(max_batch_nodes);
106 lower!(max_output_bytes);
107 lower!(max_trace_bytes);
108 lower!(max_exponent);
109 lower!(max_factorial);
110 lower!(max_iterations);
111 lower!(max_operations);
112 lower!(max_recursion_depth);
113 lower!(max_string_len);
114 lower!(max_series_terms);
115 lower!(request_timeout_ms);
116 Ok(out)
117 }
118
119 #[cfg(not(target_arch = "wasm32"))]
125 pub fn deadline(&self) -> Option<Instant> {
126 if self.request_timeout_ms == 0 {
127 None
128 } else {
129 Some(Instant::now() + Duration::from_millis(self.request_timeout_ms))
130 }
131 }
132
133 #[cfg(target_arch = "wasm32")]
134 pub fn deadline(&self) -> Option<Instant> {
135 let _ = self.request_timeout_ms;
136 None
137 }
138}
139
140#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
143#[serde(default, deny_unknown_fields)]
144pub struct LimitsOverride {
145 pub max_request_bytes: Option<usize>,
146 pub max_expression_length: Option<usize>,
147 pub max_expression_tokens: Option<usize>,
148 pub max_ast_depth: Option<usize>,
149 pub max_digits: Option<usize>,
150 pub max_integer_bits: Option<u32>,
151 pub max_decimal_scale: Option<u32>,
152 pub max_array_len: Option<usize>,
153 pub max_matrix_elements: Option<usize>,
154 pub max_batch_nodes: Option<usize>,
155 pub max_output_bytes: Option<usize>,
156 pub max_trace_bytes: Option<usize>,
157 pub max_exponent: Option<u32>,
158 pub max_factorial: Option<u32>,
159 pub max_iterations: Option<u64>,
160 pub max_operations: Option<u64>,
161 pub max_recursion_depth: Option<usize>,
162 pub max_string_len: Option<usize>,
163 pub max_series_terms: Option<u64>,
164 pub request_timeout_ms: Option<u64>,
165}
166
167#[derive(Clone, Debug, Default)]
169pub struct CancellationToken {
170 flag: Arc<AtomicBool>,
171}
172
173impl CancellationToken {
174 pub fn new() -> CancellationToken {
175 CancellationToken {
176 flag: Arc::new(AtomicBool::new(false)),
177 }
178 }
179
180 pub fn cancel(&self) {
181 self.flag.store(true, Ordering::SeqCst);
182 }
183
184 pub fn is_cancelled(&self) -> bool {
185 self.flag.load(Ordering::SeqCst)
186 }
187
188 pub fn check(&self) -> Result<(), EngineError> {
190 if self.is_cancelled() {
191 Err(EngineError::cancelled())
192 } else {
193 Ok(())
194 }
195 }
196}
197
198#[derive(Debug)]
200pub struct Budget {
201 limit: u64,
202 used: u64,
203 token: CancellationToken,
204}
205
206impl Budget {
207 pub fn new(limit: u64, token: CancellationToken) -> Budget {
208 Budget {
209 limit,
210 used: 0,
211 token,
212 }
213 }
214
215 pub fn tick(&mut self) -> Result<(), EngineError> {
217 self.tick_by(1)
218 }
219
220 pub fn tick_by(&mut self, amount: u64) -> Result<(), EngineError> {
221 self.token.check()?;
222 self.used = self.used.saturating_add(amount);
223 if self.used > self.limit {
224 return Err(EngineError::new(
225 ErrorCode::ResourceLimit,
226 format!("operation budget of {} exceeded", self.limit),
227 ));
228 }
229 Ok(())
230 }
231
232 pub fn used(&self) -> u64 {
233 self.used
234 }
235}
236
237#[cfg(test)]
238mod tests {
239 use super::*;
240
241 #[test]
242 fn overrides_cannot_raise_limits() {
243 let policy = Limits::conservative();
244 let overrides = LimitsOverride {
245 max_array_len: Some(policy.max_array_len + 1),
246 ..Default::default()
247 };
248 assert!(policy.lowered_by(&overrides).is_err());
249 let overrides = LimitsOverride {
250 max_array_len: Some(10),
251 ..Default::default()
252 };
253 assert_eq!(policy.lowered_by(&overrides).unwrap().max_array_len, 10);
254 }
255
256 #[test]
257 fn cancellation_is_cooperative() {
258 let token = CancellationToken::new();
259 assert!(token.check().is_ok());
260 token.cancel();
261 assert_eq!(token.check().unwrap_err().code, ErrorCode::Cancelled);
262 }
263}