1use std::time::Duration;
12
13use qubit_budget::BudgetError;
14use qubit_budget::MeasuredBudgetError;
15use qubit_budget::Observation;
16use qubit_budget::QuantityConversionError;
17use qubit_budget::QuantityMeasurement;
18use qubit_budget::json::JsonEncodeLimits;
19use qubit_budget::json::JsonEncodeLimits as JsonEncodeLimitsU64;
20use qubit_budget::json::JsonMeasurement;
21use qubit_budget::json::JsonResource;
22use qubit_budget::json::JsonValueLimits;
23
24use super::internal::WireDataTypeV1;
25use super::internal::WirePreflightStrategy;
26use crate::MultiValues;
27use crate::Value;
28use crate::ValueContainer;
29use crate::ValueRef;
30
31mod internal;
32
33use self::internal::json_length_writer::JsonLengthWriter;
34use self::internal::wire_scalar_measurement::decimal_len;
35use self::internal::wire_scalar_measurement::json_float32_len;
36use self::internal::wire_scalar_measurement::json_float64_len;
37use self::internal::wire_scalar_measurement::unsigned_decimal_len;
38
39macro_rules! define_value_ref_preflight_strategy {
40 (
41 $($arg:expr),*;
42 $(
43 (
44 [$($cfg:meta),*],
45 $variant:ident,
46 $type:ty,
47 $_data_type:expr,
48 $_materialization:ident,
49 $_json_class:ident,
50 $_number_projection:ident,
51 $_value_doc:literal,
52 $_multi_doc:literal,
53 [$($scalar_attr:meta),*],
54 [$($collection_attr:meta),*],
55 $_tag:literal,
56 $_wire_preflight:ident
57 )
58 ),+ $(,)?
59 ) => {
60 fn value_ref_preflight_strategy(value: ValueRef<'_>) -> WirePreflightStrategy {
61 match value {
62 ValueRef::Unset(_) => WirePreflightStrategy::UnsetText,
63 $(
64 $(#[$cfg])*
65 ValueRef::$variant(_) => WireDataTypeV1::$variant.preflight_strategy(),
66 )+
67 }
68 }
69 };
70}
71
72for_each_value_type!(define_value_ref_preflight_strategy);
73
74#[derive(Debug, Clone)]
100pub struct ValueWireEncodePreflight {
101 limits: JsonEncodeLimits,
102 nodes: usize,
103 payload_bytes: usize,
104 output_bytes: usize,
105}
106
107fn saturating_u64_to_usize(value: u64) -> usize {
111 usize::try_from(value).unwrap_or(usize::MAX)
112}
113
114impl ValueWireEncodePreflight {
115 #[must_use]
117 pub fn new(limits: JsonEncodeLimits) -> Self {
118 Self {
119 limits,
120 nodes: 0,
121 payload_bytes: 0,
122 output_bytes: 0,
123 }
124 }
125
126 #[must_use]
131 pub fn new_value_limits(limits: JsonValueLimits) -> Self {
132 let mut checker = Self::new(JsonEncodeLimits::new());
133 checker.limits = JsonEncodeLimits::builder().value_limits(limits).build();
134 checker
135 }
136
137 #[must_use]
143 pub fn new_u64_limits(limits: JsonEncodeLimitsU64<JsonResource, u64>) -> Self {
144 let value = limits.value_limits();
145 let mut builder = JsonEncodeLimits::builder();
146 if let Some(limit) = limits.max_output_bytes() {
147 builder = builder.max_output_bytes(saturating_u64_to_usize(limit));
148 }
149 if let Some(limit) = value.max_depth() {
150 builder = builder.max_depth(saturating_u64_to_usize(limit));
151 }
152 if let Some(limit) = value.max_nodes() {
153 builder = builder.max_nodes(saturating_u64_to_usize(limit));
154 }
155 if let Some(limit) = value.max_sequence_items() {
156 builder = builder.max_sequence_items(saturating_u64_to_usize(limit));
157 }
158 if let Some(limit) = value.max_map_entries() {
159 builder = builder.max_map_entries(saturating_u64_to_usize(limit));
160 }
161 if let Some(limit) = value.max_key_bytes() {
162 builder = builder.max_key_bytes(saturating_u64_to_usize(limit));
163 }
164 if let Some(limit) = value.max_string_bytes() {
165 builder = builder.max_string_bytes(saturating_u64_to_usize(limit));
166 }
167 if let Some(limit) = value.max_number_bytes() {
168 builder = builder.max_number_bytes(saturating_u64_to_usize(limit));
169 }
170 if let Some(limit) = value.max_payload_bytes() {
171 builder = builder.max_payload_bytes(saturating_u64_to_usize(limit));
172 }
173 Self::new(builder.build())
174 }
175
176 pub fn check_value(&mut self, value: &Value) -> Result<(), MeasuredBudgetError<JsonResource, usize>> {
181 self.transaction(|checker| checker.check_value_at(value, 1))
182 }
183
184 pub fn check_values(&mut self, values: &MultiValues) -> Result<(), MeasuredBudgetError<JsonResource, usize>> {
189 self.transaction(|checker| checker.check_values_at(values, 1))
190 }
191
192 pub fn check_container(&mut self, value: &ValueContainer) -> Result<(), MeasuredBudgetError<JsonResource, usize>> {
197 self.transaction(|checker| match value {
198 ValueContainer::Scalar(value) => checker.check_value_at(value, 1),
199 ValueContainer::Collection(values) => checker.check_values_at(values, 1),
200 })
201 }
202
203 fn transaction<F>(&mut self, check: F) -> Result<(), MeasuredBudgetError<JsonResource, usize>>
208 where
209 F: FnOnce(&mut Self) -> Result<(), MeasuredBudgetError<JsonResource, usize>>,
210 {
211 let mut candidate = self.clone();
212 match check(&mut candidate) {
213 Ok(()) => {
214 *self = candidate;
215 Ok(())
216 }
217 Err(error) => Err(error),
218 }
219 }
220
221 fn check_value_at(&mut self, value: &Value, depth: usize) -> Result<(), MeasuredBudgetError<JsonResource, usize>> {
222 self.check_view_at(value.view(), depth)
223 }
224
225 fn check_view_at(
226 &mut self,
227 value: ValueRef<'_>,
228 depth: usize,
229 ) -> Result<(), MeasuredBudgetError<JsonResource, usize>> {
230 match value_ref_preflight_strategy(value) {
231 WirePreflightStrategy::Boolean => match value {
232 ValueRef::Bool(_) => self.admit(JsonMeasurement::Boolean { depth }, 1, 0),
233 _ => unreachable!("table strategy and ValueRef variant diverged"),
234 },
235 WirePreflightStrategy::BorrowedText => match value {
236 ValueRef::Char(value) => self.admit_string(depth, value.len_utf8()),
237 ValueRef::String(value) => self.admit_string(depth, value.len()),
238 #[cfg(feature = "url")]
239 ValueRef::Url(value) => self.admit_string(depth, value.as_str().len()),
240 _ => unreachable!("table strategy and ValueRef variant diverged"),
241 },
242 WirePreflightStrategy::JsonInteger => match value {
243 ValueRef::Int8(value) => self.admit_number(depth, decimal_len(value as i128)),
244 ValueRef::Int16(value) => self.admit_number(depth, decimal_len(value as i128)),
245 ValueRef::Int32(value) => self.admit_number(depth, decimal_len(value as i128)),
246 ValueRef::Int64(value) => self.admit_number(depth, decimal_len(value as i128)),
247 ValueRef::UInt8(value) => self.admit_number(depth, unsigned_decimal_len(value as u128)),
248 ValueRef::UInt16(value) => self.admit_number(depth, unsigned_decimal_len(value as u128)),
249 ValueRef::UInt32(value) => self.admit_number(depth, unsigned_decimal_len(value as u128)),
250 ValueRef::UInt64(value) => self.admit_number(depth, unsigned_decimal_len(value as u128)),
251 _ => unreachable!("table strategy and ValueRef variant diverged"),
252 },
253 WirePreflightStrategy::DecimalText => match value {
254 ValueRef::Int128(value) => self.admit_string(depth, decimal_len(value)),
255 ValueRef::UInt128(value) => self.admit_string(depth, unsigned_decimal_len(value)),
256 _ => unreachable!("table strategy and ValueRef variant diverged"),
257 },
258 WirePreflightStrategy::UnsetText => match value {
259 ValueRef::Unset(data_type) => {
260 let tag = WireDataTypeV1::from(data_type).tag();
261 self.admit_string(depth, tag.len())
262 }
263 _ => unreachable!("table strategy and ValueRef variant diverged"),
264 },
265 WirePreflightStrategy::JsonFloat => match value {
266 ValueRef::Float32(value) => {
267 self.admit_number(depth, if value.is_finite() { json_float32_len(value) } else { 1 })
268 }
269 ValueRef::Float64(value) => {
270 self.admit_number(depth, if value.is_finite() { json_float64_len(value) } else { 1 })
271 }
272 _ => unreachable!("table strategy and ValueRef variant diverged"),
273 },
274 WirePreflightStrategy::BigIntegerText => {
275 #[cfg(feature = "big-integer")]
276 {
277 match value {
278 ValueRef::BigInteger(value) => self.admit_string(depth, bigint_digits(value)),
279 _ => unreachable!("table strategy and ValueRef variant diverged"),
280 }
281 }
282 #[cfg(not(feature = "big-integer"))]
283 {
284 unreachable!("big-integer strategy is unavailable without its feature")
285 }
286 }
287 WirePreflightStrategy::DecimalObject => {
288 #[cfg(feature = "big-decimal")]
289 {
290 match value {
291 ValueRef::BigDecimal(value) => {
292 let (coefficient, scale) = value.as_bigint_and_scale();
293 drop(coefficient);
294 self.admit_object_with_keys(
295 depth,
296 [("coefficient", 0, true), ("scale", decimal_len(scale as i128), false)],
297 )
298 }
299 _ => unreachable!("table strategy and ValueRef variant diverged"),
300 }
301 }
302 #[cfg(not(feature = "big-decimal"))]
303 {
304 unreachable!("big-decimal strategy is unavailable without its feature")
305 }
306 }
307 WirePreflightStrategy::TemporalText => match value {
308 #[cfg(feature = "chrono")]
309 ValueRef::Date(_) | ValueRef::Time(_) | ValueRef::DateTime(_) | ValueRef::Instant(_) => {
310 self.admit_string(depth, 0)
311 }
312 _ => unreachable!("table strategy and ValueRef variant diverged"),
313 },
314 WirePreflightStrategy::DurationObject => match value {
315 ValueRef::Duration(value) => self.admit_duration(depth, *value),
316 _ => unreachable!("table strategy and ValueRef variant diverged"),
317 },
318 WirePreflightStrategy::StringMap => match value {
319 ValueRef::StringMap(value) => self.check_string_map(value, depth),
320 _ => unreachable!("table strategy and ValueRef variant diverged"),
321 },
322 #[cfg(feature = "json")]
323 WirePreflightStrategy::JsonTree => match value {
324 ValueRef::Json(value) => self.check_json(value, depth),
325 _ => unreachable!("table strategy and ValueRef variant diverged"),
326 },
327 }
328 }
329
330 fn check_values_at(
331 &mut self,
332 values: &MultiValues,
333 depth: usize,
334 ) -> Result<(), MeasuredBudgetError<JsonResource, usize>> {
335 self.admit(
336 JsonMeasurement::Array {
337 depth,
338 items: values.len(),
339 },
340 values.len().saturating_add(2),
341 2,
342 )?;
343 let view = values.view();
344 for index in 0..view.len() {
345 if let Some(value) = view.get(index) {
346 self.check_view_at(value, depth + 1)?;
347 }
348 }
349 Ok(())
350 }
351
352 fn check_string_map(
353 &mut self,
354 map: &std::collections::HashMap<String, String>,
355 depth: usize,
356 ) -> Result<(), MeasuredBudgetError<JsonResource, usize>> {
357 self.admit(
358 JsonMeasurement::Object {
359 depth,
360 entries: map.len(),
361 },
362 map.len().saturating_mul(2).saturating_add(2),
363 2,
364 )?;
365 for (key, value) in map {
366 self.check_point(JsonMeasurement::Key { bytes: key.len() })?;
367 self.admit_string(depth + 1, value.len())?;
368 }
369 Ok(())
370 }
371
372 #[cfg(feature = "json")]
373 fn check_json(
374 &mut self,
375 value: &serde_json::Value,
376 depth: usize,
377 ) -> Result<(), MeasuredBudgetError<JsonResource, usize>> {
378 enum Frame<'a> {
379 Visit(&'a serde_json::Value, usize),
380 Array(std::slice::Iter<'a, serde_json::Value>, usize),
381 Object(serde_json::map::Iter<'a>, usize),
382 }
383
384 let mut frames = vec![Frame::Visit(value, depth)];
385 while let Some(frame) = frames.pop() {
386 match frame {
387 Frame::Visit(value, depth) => match value {
388 serde_json::Value::Null => self.admit(JsonMeasurement::Null { depth }, 1, 0)?,
389 serde_json::Value::Bool(_) => self.admit(JsonMeasurement::Boolean { depth }, 1, 0)?,
390 serde_json::Value::Number(value) => {
391 let mut writer = JsonLengthWriter::default();
392 serde_json::to_writer(&mut writer, value).expect("JSON Number serialization cannot fail");
393 self.admit_number(depth, writer.len)?;
394 }
395 serde_json::Value::String(value) => self.admit_string(depth, value.len())?,
396 serde_json::Value::Array(values) => {
397 self.admit(
398 JsonMeasurement::Array {
399 depth,
400 items: values.len(),
401 },
402 values.len().saturating_add(2),
403 2,
404 )?;
405 frames.push(Frame::Array(values.iter(), depth + 1));
406 }
407 serde_json::Value::Object(values) => {
408 self.admit(
409 JsonMeasurement::Object {
410 depth,
411 entries: values.len(),
412 },
413 values.len().saturating_mul(2).saturating_add(2),
414 2,
415 )?;
416 frames.push(Frame::Object(values.iter(), depth + 1));
417 }
418 },
419 Frame::Array(mut values, depth) => {
420 if let Some(value) = values.next() {
421 frames.push(Frame::Array(values, depth));
422 frames.push(Frame::Visit(value, depth));
423 }
424 }
425 Frame::Object(mut values, depth) => {
426 if let Some((key, value)) = values.next() {
427 self.check_point(JsonMeasurement::Key { bytes: key.len() })?;
428 frames.push(Frame::Object(values, depth));
429 frames.push(Frame::Visit(value, depth));
430 }
431 }
432 }
433 }
434 Ok(())
435 }
436
437 fn admit_string(&mut self, depth: usize, bytes: usize) -> Result<(), MeasuredBudgetError<JsonResource, usize>> {
438 self.admit(JsonMeasurement::String { depth, bytes }, bytes.saturating_add(2), 1)
439 }
440
441 fn admit_number(&mut self, depth: usize, bytes: usize) -> Result<(), MeasuredBudgetError<JsonResource, usize>> {
442 self.admit(JsonMeasurement::Number { depth, bytes }, bytes, 1)
443 }
444
445 fn admit_duration(
446 &mut self,
447 depth: usize,
448 value: Duration,
449 ) -> Result<(), MeasuredBudgetError<JsonResource, usize>> {
450 self.admit_object_with_keys(
451 depth,
452 [
453 ("secs", unsigned_decimal_len(value.as_secs() as u128), false),
454 ("nanos", unsigned_decimal_len(value.subsec_nanos() as u128), false),
455 ],
456 )
457 }
458
459 fn admit_object_with_keys<const N: usize>(
460 &mut self,
461 depth: usize,
462 fields: [(&str, usize, bool); N],
463 ) -> Result<(), MeasuredBudgetError<JsonResource, usize>> {
464 let output = fields
465 .iter()
466 .try_fold(2_usize, |total, (key, _, _)| {
467 total
468 .checked_add(key.len())
469 .and_then(|total| total.checked_add(3))
470 .ok_or(())
471 })
472 .ok()
473 .and_then(|total| total.checked_add(N.saturating_sub(1)))
474 .ok_or_else(|| self.quantity_error(JsonResource::OutputBytes))?;
475 self.admit(JsonMeasurement::Object { depth, entries: N }, output, 0)?;
476 for (key, value, is_string) in fields {
477 self.check_point(JsonMeasurement::Key { bytes: key.len() })?;
478 if is_string {
479 self.admit_string(depth + 1, value)?;
480 } else {
481 self.admit_number(depth + 1, value)?;
482 }
483 }
484 Ok(())
485 }
486
487 fn admit(
488 &mut self,
489 measurement: JsonMeasurement,
490 output: usize,
491 payload: usize,
492 ) -> Result<(), MeasuredBudgetError<JsonResource, usize>> {
493 self.check_point(measurement)?;
494 self.nodes = self
495 .nodes
496 .checked_add(1)
497 .ok_or_else(|| self.quantity_error(JsonResource::Nodes))?;
498 self.payload_bytes = self
499 .payload_bytes
500 .checked_add(payload)
501 .ok_or_else(|| self.quantity_error(JsonResource::PayloadBytes))?;
502 self.output_bytes = self
503 .output_bytes
504 .checked_add(output)
505 .ok_or_else(|| self.quantity_error(JsonResource::OutputBytes))?;
506 self.check_cumulative(JsonResource::Nodes, self.nodes, self.limits.value_limits().max_nodes())?;
507 self.check_cumulative(
508 JsonResource::PayloadBytes,
509 self.payload_bytes,
510 self.limits.value_limits().max_payload_bytes(),
511 )?;
512 self.check_cumulative(
513 JsonResource::OutputBytes,
514 self.output_bytes,
515 self.limits.max_output_bytes(),
516 )
517 }
518
519 fn check_point(&self, measurement: JsonMeasurement) -> Result<(), MeasuredBudgetError<JsonResource, usize>> {
520 self.limits.value_limits().check_point(measurement)
521 }
522
523 fn check_cumulative(
524 &self,
525 resource: JsonResource,
526 observed: usize,
527 maximum: Option<usize>,
528 ) -> Result<(), MeasuredBudgetError<JsonResource, usize>> {
529 match maximum {
530 Some(maximum) if observed > maximum => Err(self.lower_bound_error(resource, observed, maximum)),
531 _ => Ok(()),
532 }
533 }
534
535 fn lower_bound_error(
536 &self,
537 resource: JsonResource,
538 observed: usize,
539 maximum: usize,
540 ) -> MeasuredBudgetError<JsonResource, usize> {
541 MeasuredBudgetError::Budget(BudgetError::LimitExceeded {
542 resource,
543 observed: Observation::AtLeast(observed),
544 maximum,
545 })
546 }
547
548 fn quantity_error(&self, resource: JsonResource) -> MeasuredBudgetError<JsonResource, usize> {
549 MeasuredBudgetError::quantity(
550 resource,
551 QuantityConversionError::new(QuantityMeasurement::Usize(usize::MAX), "usize"),
552 )
553 }
554}
555
556#[cfg(any(feature = "big-integer", feature = "big-decimal"))]
557fn bigint_digits(value: &num_bigint::BigInt) -> usize {
558 if value.sign() == num_bigint::Sign::NoSign {
559 1
560 } else {
561 ((value.bits().saturating_sub(1)) / 4 + 1) as usize + usize::from(value.sign() == num_bigint::Sign::Minus)
562 }
563}
564
565#[cfg(test)]
566mod tests {
567 use std::panic::AssertUnwindSafe;
568
569 use qubit_budget::MeasuredBudgetError;
570 use qubit_budget::json::JsonEncodeLimits;
571 use qubit_budget::json::JsonResource;
572
573 use super::ValueWireEncodePreflight;
574 use crate::Value;
575
576 fn assert_overflow(error: MeasuredBudgetError<JsonResource, usize>) {
578 assert!(
579 matches!(error, MeasuredBudgetError::Quantity { .. }),
580 "expected a quantity overflow, got {error:?}",
581 );
582 }
583
584 #[test]
585 fn test_check_value_reports_node_counter_overflow_and_rolls_back() {
586 let mut checker = ValueWireEncodePreflight::new(JsonEncodeLimits::new());
587 checker.nodes = usize::MAX;
588
589 let error = checker
590 .check_value(&Value::Bool(true))
591 .expect_err("the node counter cannot exceed usize::MAX");
592
593 assert_overflow(error);
594 assert_eq!(checker.nodes, usize::MAX);
595 assert_eq!(checker.payload_bytes, 0);
596 assert_eq!(checker.output_bytes, 0);
597 }
598
599 #[test]
600 fn test_check_value_reports_payload_counter_overflow_and_rolls_back() {
601 let mut checker = ValueWireEncodePreflight::new(JsonEncodeLimits::new());
602 checker.payload_bytes = usize::MAX;
603
604 let error = checker
605 .check_value(&Value::String("x".to_owned()))
606 .expect_err("the payload counter cannot exceed usize::MAX");
607
608 assert_overflow(error);
609 assert_eq!(checker.nodes, 0);
610 assert_eq!(checker.payload_bytes, usize::MAX);
611 assert_eq!(checker.output_bytes, 0);
612 }
613
614 #[test]
615 fn test_check_value_reports_output_counter_overflow_and_rolls_back() {
616 let mut checker = ValueWireEncodePreflight::new(JsonEncodeLimits::new());
617 checker.output_bytes = usize::MAX;
618
619 let error = checker
620 .check_value(&Value::Bool(true))
621 .expect_err("the output counter cannot exceed usize::MAX");
622
623 assert_overflow(error);
624 assert_eq!(checker.nodes, 0);
625 assert_eq!(checker.payload_bytes, 0);
626 assert_eq!(checker.output_bytes, usize::MAX);
627 }
628
629 #[test]
630 fn test_transaction_does_not_commit_counters_when_check_panics() {
631 let mut checker = ValueWireEncodePreflight::new(JsonEncodeLimits::new());
632 let result = std::panic::catch_unwind(AssertUnwindSafe(|| {
633 checker.transaction(|candidate| {
634 candidate.nodes = 17;
635 candidate.payload_bytes = 23;
636 panic!("simulated preflight panic");
637 })
638 }));
639
640 assert!(result.is_err());
641 assert_eq!(checker.nodes, 0);
642 assert_eq!(checker.payload_bytes, 0);
643 assert_eq!(checker.output_bytes, 0);
644 }
645}