pipa/http/
chunked.rs

1#[derive(Debug)]
2pub struct ChunkedDecoder {
3    state: ChunkState,
4    chunk_size: usize,
5    buffer: Vec<u8>,
6    done: bool,
7}
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10enum ChunkState {
11    Size,
12    Data,
13    Trailer,
14    Done,
15}
16
17impl ChunkedDecoder {
18    pub fn new() -> Self {
19        ChunkedDecoder {
20            state: ChunkState::Size,
21            chunk_size: 0,
22            buffer: Vec::new(),
23            done: false,
24        }
25    }
26
27    pub fn feed(&mut self, data: &[u8]) -> Result<Vec<u8>, String> {
28        self.buffer.extend_from_slice(data);
29        let mut output = Vec::new();
30
31        loop {
32            match self.state {
33                ChunkState::Size => {
34                    if let Some(end) = self.buffer.iter().position(|&b| b == b'\r') {
35                        if end + 1 >= self.buffer.len() {
36                            break;
37                        }
38                        let line = &self.buffer[..end];
39                        let size_str = core::str::from_utf8(line)
40                            .map_err(|e| format!("chunk size utf8: {e}"))?;
41                        let size_end = size_str.find(|c: char| c == ';' || c == ' ');
42                        let size_hex = match size_end {
43                            Some(pos) => &size_str[..pos],
44                            None => size_str,
45                        };
46                        let size = usize::from_str_radix(size_hex.trim(), 16)
47                            .map_err(|e| format!("chunk size parse: {e}"))?;
48                        self.chunk_size = size;
49                        self.buffer.drain(..end + 2);
50                        if size == 0 {
51                            self.state = ChunkState::Trailer;
52                        } else {
53                            self.state = ChunkState::Data;
54                        }
55                    } else {
56                        break;
57                    }
58                }
59                ChunkState::Data => {
60                    if self.buffer.len() >= self.chunk_size + 2 {
61                        let chunk_data = self.buffer.drain(..self.chunk_size).collect::<Vec<_>>();
62                        self.buffer.drain(..2);
63                        output.extend(chunk_data);
64                        self.state = ChunkState::Size;
65                    } else {
66                        break;
67                    }
68                }
69                ChunkState::Done => {
70                    break;
71                }
72                ChunkState::Trailer => {
73                    if self.buffer.len() >= 2 {
74                        if self.buffer[0] == b'\r' && self.buffer[1] == b'\n' {
75                            self.buffer.drain(..2);
76                            self.state = ChunkState::Done;
77                            self.done = true;
78                            break;
79                        }
80                        if let Some(end) = self.buffer.windows(2).position(|w| w == b"\r\n") {
81                            self.buffer.drain(..end + 2);
82                        } else {
83                            break;
84                        }
85                    } else {
86                        break;
87                    }
88                }
89            }
90        }
91
92        Ok(output)
93    }
94
95    pub fn is_done(&self) -> bool {
96        self.done
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    #[test]
105    fn test_single_chunk() {
106        let mut dec = ChunkedDecoder::new();
107        let out = dec.feed(b"5\r\nHello\r\n0\r\n\r\n").unwrap();
108        assert_eq!(out, b"Hello");
109        assert!(dec.is_done());
110    }
111
112    #[test]
113    fn test_multi_chunk() {
114        let mut dec = ChunkedDecoder::new();
115        let out = dec
116            .feed(b"6\r\nHello \r\n6\r\nWorld!\r\n0\r\n\r\n")
117            .unwrap();
118        assert_eq!(out, b"Hello World!");
119        assert!(dec.is_done());
120    }
121
122    #[test]
123    fn test_partial_feed() {
124        let mut dec = ChunkedDecoder::new();
125        let out1 = dec.feed(b"5\r\nHel").unwrap();
126        assert!(out1.is_empty());
127        assert!(!dec.is_done());
128
129        let out2 = dec.feed(b"lo\r\n0\r\n\r\n").unwrap();
130        assert_eq!(out2, b"Hello");
131        assert!(dec.is_done());
132    }
133}
pipa/http/chunked.rs

pipa/http/
chunked.rs
