1use core::cell::Cell;
2
3use crate::Result;
4use crate::error::Fail;
5
6pub const MAX_DEPTH: u32 = 128;
7
8thread_local! {
9 static DEPTH: Cell<u32> = const { Cell::new(0) };
10}
11
12pub struct DepthGuard {
13 _seal: (),
14}
15
16impl DepthGuard {
17 pub fn enter() -> Result<Self> {
18 DEPTH.with(|depth| {
19 let level = depth.get();
20 if level >= MAX_DEPTH {
21 return Err(Fail::with("JSON nesting too deep"));
22 }
23 depth.set(level + 1);
24 Ok(Self { _seal: () })
25 })
26 }
27
28 pub fn level() -> u32 {
29 DEPTH.with(Cell::get)
30 }
31}
32
33impl Drop for DepthGuard {
34 fn drop(&mut self) {
35 DEPTH.with(|depth| depth.set(depth.get().saturating_sub(1)));
36 }
37}
38
39#[cfg(test)]
40mod tests {
41 use super::*;
42
43 #[test]
44 fn guard_increments_and_releases() {
45 assert_eq!(DepthGuard::level(), 0);
46 {
47 let _a = DepthGuard::enter().unwrap();
48 assert_eq!(DepthGuard::level(), 1);
49 {
50 let _b = DepthGuard::enter().unwrap();
51 assert_eq!(DepthGuard::level(), 2);
52 }
53 assert_eq!(DepthGuard::level(), 1);
54 }
55 assert_eq!(DepthGuard::level(), 0);
56 }
57
58 #[test]
59 fn guard_caps_at_max_depth() {
60 let mut held = Vec::new();
61 for _ in 0..MAX_DEPTH {
62 held.push(DepthGuard::enter().expect("under cap"));
63 }
64 assert_eq!(DepthGuard::level(), MAX_DEPTH);
65 assert!(DepthGuard::enter().is_err());
66 drop(held);
67 assert_eq!(DepthGuard::level(), 0);
68 DepthGuard::enter().expect("recovered after release");
69 }
70
71 fn decode_recursive(remaining: u32) -> Result<u32> {
72 let _guard = DepthGuard::enter()?;
73 if remaining == 0 {
74 return Ok(DepthGuard::level());
75 }
76 decode_recursive(remaining - 1)
77 }
78
79 #[test]
80 fn recursive_decode_errors_before_overflow() {
81 assert!(decode_recursive(10).is_ok());
82 let err = decode_recursive(MAX_DEPTH + 5_000);
83 assert!(err.is_err(), "deep recursion must error, not overflow");
84 assert_eq!(DepthGuard::level(), 0);
85 }
86}