1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
use crate::{
element::Element,
error::{Error, Result},
MAX_DEPTH,
};
#[derive(Clone, Debug)]
pub struct DepthTracker {
tracking: Vec<u32>,
}
impl DepthTracker {
/// Create a new depth tracker
pub fn new() -> Self {
Self {
tracking: Vec::new(),
}
}
/// Update the depth tracker on each new element to serialize.
#[inline]
pub fn update_elem(&mut self, elem: &Element) -> Result<()> {
// Subtract from count for next element
if let Some(v) = self.tracking.last_mut() {
*v -= 1;
}
// Increase nest depth if this is a nesting element
match elem {
Element::Map(len) => self.tracking.push(2 * (*len as u32)), // 2 elements per map item
Element::Array(len) => self.tracking.push(*len as u32),
_ => (),
}
// Check to see if we hit the nesting limit
if self.tracking.len() > MAX_DEPTH {
return Err(Error::ParseLimit("Depth limit exceeded".to_string()));
}
self.purge_zeros();
Ok(())
}
/// Drop any depth tracking elements that have hit zero
#[inline]
pub fn purge_zeros(&mut self) {
loop {
match self.tracking.last() {
Some(v) if *v == 0 => {
self.tracking.pop();
}
_ => break,
}
}
}
/// Drop a depth before we've seen enough elements. This can be used by map/seq serializers
/// that didn't know their total length ahead of time. This way, they can put in a
/// maximally-sized map/array element, then run through the depth tracker as normal, calling
/// this when done.
#[inline]
pub fn early_end(&mut self) {
self.tracking.pop();
self.purge_zeros();
}
}