Skip to main content

sark_json/
depth.rs

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}