Skip to main content

sidereon_core/ntrip/
chunk.rs

1use crate::{Error, Result};
2
3const MAX_CHUNK_SIZE: usize = 16 * 1024 * 1024;
4const MAX_LINE: usize = 8 * 1024;
5const MAX_HEX_DIGITS: usize = 8;
6
7#[derive(Clone, Debug, PartialEq, Eq)]
8pub struct ChunkedDecoder {
9    buf: Vec<u8>,
10    state: ChunkState,
11    finished: bool,
12    poison: Option<String>,
13    trailer_len: usize,
14}
15
16#[derive(Clone, Debug, PartialEq, Eq)]
17enum ChunkState {
18    Size,
19    Data(usize),
20    DataLineEnd,
21    Trailer,
22}
23
24impl Default for ChunkedDecoder {
25    fn default() -> Self {
26        Self::new()
27    }
28}
29
30impl ChunkedDecoder {
31    pub fn new() -> Self {
32        Self {
33            buf: Vec::new(),
34            state: ChunkState::Size,
35            finished: false,
36            poison: None,
37            trailer_len: 0,
38        }
39    }
40
41    pub fn push(&mut self, bytes: &[u8]) -> Result<Vec<u8>> {
42        if let Some(message) = &self.poison {
43            return Err(Error::Parse(message.clone()));
44        }
45        if self.finished {
46            return Ok(Vec::new());
47        }
48        self.buf.extend_from_slice(bytes);
49        let mut out = Vec::new();
50
51        loop {
52            match self.state {
53                ChunkState::Size => {
54                    let Some(line) = self.read_line()? else {
55                        break;
56                    };
57                    let size = match parse_size_line(&line) {
58                        Ok(size) => size,
59                        Err(err) => return self.poison_from(err),
60                    };
61                    if size == 0 {
62                        self.state = ChunkState::Trailer;
63                    } else {
64                        self.state = ChunkState::Data(size);
65                    }
66                }
67                ChunkState::Data(remaining) => {
68                    if self.buf.is_empty() {
69                        break;
70                    }
71                    let n = remaining.min(self.buf.len());
72                    out.extend_from_slice(&self.buf[..n]);
73                    self.buf.drain(..n);
74                    let left = remaining - n;
75                    self.state = if left == 0 {
76                        ChunkState::DataLineEnd
77                    } else {
78                        ChunkState::Data(left)
79                    };
80                }
81                ChunkState::DataLineEnd => {
82                    if self.buf.is_empty() {
83                        break;
84                    }
85                    if self.buf[0] == b'\n' {
86                        self.buf.drain(..1);
87                        self.state = ChunkState::Size;
88                    } else if self.buf[0] == b'\r' {
89                        if self.buf.len() < 2 {
90                            break;
91                        }
92                        if self.buf[1] != b'\n' {
93                            return self.poison("malformed chunk data terminator");
94                        }
95                        self.buf.drain(..2);
96                        self.state = ChunkState::Size;
97                    } else {
98                        return self.poison("malformed chunk data terminator");
99                    }
100                }
101                ChunkState::Trailer => {
102                    let Some(line) = self.read_line()? else {
103                        break;
104                    };
105                    self.trailer_len += line.len();
106                    if self.trailer_len > MAX_LINE {
107                        return self.poison("chunk trailer section too long");
108                    }
109                    if line.is_empty() {
110                        self.finished = true;
111                        break;
112                    }
113                }
114            }
115        }
116
117        Ok(out)
118    }
119
120    pub fn finished(&self) -> bool {
121        self.finished
122    }
123
124    pub fn reset(&mut self) {
125        *self = Self::new();
126    }
127
128    fn read_line(&mut self) -> Result<Option<Vec<u8>>> {
129        if let Some(pos) = self.buf.iter().position(|&b| b == b'\n') {
130            let mut line: Vec<u8> = self.buf.drain(..=pos).collect();
131            if line.ends_with(b"\n") {
132                line.pop();
133            }
134            if line.ends_with(b"\r") {
135                line.pop();
136            }
137            return Ok(Some(line));
138        }
139        if self.buf.len() > MAX_LINE {
140            return self.poison("chunk line too long");
141        }
142        Ok(None)
143    }
144
145    fn poison<T>(&mut self, message: &str) -> Result<T> {
146        self.poison = Some(message.to_string());
147        Err(Error::Parse(message.to_string()))
148    }
149
150    fn poison_from<T>(&mut self, err: Error) -> Result<T> {
151        let message = err.to_string();
152        self.poison = Some(message.clone());
153        Err(Error::Parse(message))
154    }
155}
156
157fn parse_size_line(line: &[u8]) -> Result<usize> {
158    if line.len() > MAX_LINE {
159        return Err(Error::Parse("chunk size line too long".into()));
160    }
161    let size_part = line.split(|&b| b == b';').next().unwrap_or(line);
162    let size_part = trim_ascii(size_part);
163    if size_part.is_empty() {
164        return Err(Error::Parse("empty chunk size".into()));
165    }
166    if size_part.len() > MAX_HEX_DIGITS {
167        return Err(Error::Parse("chunk size has too many hex digits".into()));
168    }
169    let mut size = 0usize;
170    for &b in size_part {
171        let digit = match b {
172            b'0'..=b'9' => usize::from(b - b'0'),
173            b'a'..=b'f' => usize::from(b - b'a' + 10),
174            b'A'..=b'F' => usize::from(b - b'A' + 10),
175            _ => return Err(Error::Parse("invalid chunk size digit".into())),
176        };
177        size = (size << 4) | digit;
178    }
179    if size > MAX_CHUNK_SIZE {
180        return Err(Error::Parse("chunk size exceeds limit".into()));
181    }
182    Ok(size)
183}
184
185fn trim_ascii(bytes: &[u8]) -> &[u8] {
186    let mut start = 0;
187    let mut end = bytes.len();
188    while start < end && bytes[start].is_ascii_whitespace() {
189        start += 1;
190    }
191    while end > start && bytes[end - 1].is_ascii_whitespace() {
192        end -= 1;
193    }
194    &bytes[start..end]
195}