pipa/http/
body_reader.rs

1use crate::http::chunked::ChunkedDecoder;
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4pub enum BodyMode {
5    ContentLength(usize),
6    Chunked,
7    ConnectionClose,
8    None,
9}
10
11#[derive(Debug)]
12pub struct BodyReader {
13    mode: BodyMode,
14    buffer: Vec<u8>,
15    total_read: usize,
16    done: bool,
17    chunked_decoder: Option<ChunkedDecoder>,
18}
19
20impl BodyReader {
21    pub fn new(mode: BodyMode) -> Self {
22        let chunked_decoder = if matches!(mode, BodyMode::Chunked) {
23            Some(ChunkedDecoder::new())
24        } else {
25            None
26        };
27        BodyReader {
28            mode,
29            buffer: Vec::new(),
30            total_read: 0,
31            done: false,
32            chunked_decoder,
33        }
34    }
35
36    pub fn feed(&mut self, data: &[u8]) -> Result<Vec<u8>, String> {
37        if self.done {
38            return Ok(Vec::new());
39        }
40
41        match self.mode {
42            BodyMode::ContentLength(len) => {
43                let remaining = len.saturating_sub(self.total_read);
44                let take = data.len().min(remaining);
45                let chunk = data[..take].to_vec();
46                self.buffer.extend_from_slice(&chunk);
47                self.total_read += take;
48                if self.total_read >= len {
49                    self.done = true;
50                }
51                Ok(chunk)
52            }
53            BodyMode::Chunked => {
54                let decoder = self
55                    .chunked_decoder
56                    .as_mut()
57                    .ok_or("chunked decoder not initialized")?;
58                let decoded = decoder.feed(data)?;
59                if !decoded.is_empty() {
60                    self.buffer.extend_from_slice(&decoded);
61                    self.total_read += decoded.len();
62                }
63                if decoder.is_done() {
64                    self.done = true;
65                }
66                Ok(decoded)
67            }
68            BodyMode::ConnectionClose => {
69                self.buffer.extend_from_slice(data);
70                self.total_read += data.len();
71                Ok(data.to_vec())
72            }
73            BodyMode::None => {
74                self.done = true;
75                Ok(Vec::new())
76            }
77        }
78    }
79
80    pub fn finish(&mut self) {
81        self.done = true;
82    }
83
84    pub fn is_done(&self) -> bool {
85        self.done
86    }
87
88    pub fn buffer(&self) -> &[u8] {
89        &self.buffer
90    }
91
92    pub fn total_read(&self) -> usize {
93        self.total_read
94    }
95
96    pub fn mode(&self) -> BodyMode {
97        self.mode
98    }
99
100    pub fn take_body(&mut self) -> Vec<u8> {
101        self.done = true;
102        core::mem::take(&mut self.buffer)
103    }
104
105    pub fn body_text(&self) -> Result<String, String> {
106        String::from_utf8(self.buffer.clone()).map_err(|e| format!("body utf8 decode: {e}"))
107    }
108}
pipa/http/body_reader.rs

pipa/http/
body_reader.rs
