1use crate::error::{Error, Result};
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub struct Budget {
13 pub max_calls: u64,
15 pub max_bytes: u64,
17 pub max_wall_ms: u64,
19 pub max_nodes: u64,
21 pub max_depth: u64,
23 pub max_array_len: u64,
25 pub max_value_size: u64,
27}
28
29impl Default for Budget {
30 fn default() -> Self {
31 Self {
32 max_calls: 10_000,
33 max_bytes: 64 << 20,
34 max_wall_ms: 5_000,
35 max_nodes: 100_000,
36 max_depth: 256,
37 max_array_len: 100_000,
38 max_value_size: 10_000_000,
39 }
40 }
41}
42
43impl Budget {
44 pub fn strict() -> Self {
46 Self {
47 max_calls: 64,
48 max_bytes: 1 << 20,
49 max_wall_ms: 100,
50 max_nodes: 4_096,
51 max_depth: 64,
52 max_array_len: 1_024,
53 max_value_size: 1 << 20,
54 }
55 }
56
57 pub fn calls(limit: u64) -> Self {
59 Self {
60 max_calls: limit,
61 ..Budget::default()
62 }
63 }
64
65 pub fn check(&self) -> Result<()> {
67 for (name, value) in [
68 ("max_calls", self.max_calls),
69 ("max_bytes", self.max_bytes),
70 ("max_wall_ms", self.max_wall_ms),
71 ("max_nodes", self.max_nodes),
72 ("max_depth", self.max_depth),
73 ("max_array_len", self.max_array_len),
74 ("max_value_size", self.max_value_size),
75 ] {
76 if value == 0 {
77 return Err(Error::InvalidArgument(match name {
78 "max_calls" => "max_calls must be non-zero",
79 "max_bytes" => "max_bytes must be non-zero",
80 "max_wall_ms" => "max_wall_ms must be non-zero",
81 "max_nodes" => "max_nodes must be non-zero",
82 "max_depth" => "max_depth must be non-zero",
83 "max_array_len" => "max_array_len must be non-zero",
84 _ => "max_value_size must be non-zero",
85 }));
86 }
87 }
88 Ok(())
89 }
90
91 pub fn is_attenuation_of(&self, parent: &Budget) -> bool {
93 self.max_calls <= parent.max_calls
94 && self.max_bytes <= parent.max_bytes
95 && self.max_wall_ms <= parent.max_wall_ms
96 && self.max_nodes <= parent.max_nodes
97 && self.max_depth <= parent.max_depth
98 && self.max_array_len <= parent.max_array_len
99 && self.max_value_size <= parent.max_value_size
100 }
101
102 pub fn intersect(&self, other: &Budget) -> Budget {
104 Budget {
105 max_calls: self.max_calls.min(other.max_calls),
106 max_bytes: self.max_bytes.min(other.max_bytes),
107 max_wall_ms: self.max_wall_ms.min(other.max_wall_ms),
108 max_nodes: self.max_nodes.min(other.max_nodes),
109 max_depth: self.max_depth.min(other.max_depth),
110 max_array_len: self.max_array_len.min(other.max_array_len),
111 max_value_size: self.max_value_size.min(other.max_value_size),
112 }
113 }
114
115 pub fn to_active_budget(&self) -> active_core::Budget {
118 active_core::Budget {
119 max_nodes: self.max_nodes,
120 max_depth: self.max_depth,
121 max_array_len: self.max_array_len,
122 max_keys: self.max_depth.saturating_mul(64),
123 max_wall_ms: self.max_wall_ms,
124 max_value_size: self.max_value_size,
125 max_total_value_size: self.max_value_size.saturating_mul(4),
126 }
127 }
128}
129
130#[derive(Debug, Clone, Copy, PartialEq, Eq)]
132pub struct BudgetMeter {
133 budget: Budget,
134 calls: u64,
135 bytes: u64,
136}
137
138impl BudgetMeter {
139 pub fn new(budget: Budget) -> Self {
141 Self {
142 budget,
143 calls: 0,
144 bytes: 0,
145 }
146 }
147
148 pub fn budget(&self) -> &Budget {
150 &self.budget
151 }
152
153 pub fn calls(&self) -> u64 {
155 self.calls
156 }
157
158 pub fn bytes(&self) -> u64 {
160 self.bytes
161 }
162
163 pub fn remaining_calls(&self) -> u64 {
165 self.budget.max_calls.saturating_sub(self.calls)
166 }
167
168 pub fn charge(&mut self, bytes: u64) -> Result<()> {
170 if self.calls >= self.budget.max_calls {
171 return Err(Error::BudgetExceeded("call"));
172 }
173 if self.bytes.saturating_add(bytes) > self.budget.max_bytes {
174 return Err(Error::BudgetExceeded("byte"));
175 }
176 self.calls += 1;
177 self.bytes += bytes;
178 Ok(())
179 }
180}