http1_spec/
chunked_body_parser.rs

1use std::io::{BufRead, Read as _};
2
3use crate::{
4    body_parser::{BodyParseError, BodyParseOutput, BodyParser},
5    CR, CRLF, LF,
6};
7
8//
9//
10//
11const LENGTH_MAX_LEN: usize = 4; // b"FFFF"
12const DATA_DEFAULT_LEN: usize = 512;
13
14//
15//
16//
17#[derive(Default)]
18pub struct ChunkedBodyParser {
19    //
20    state: State,
21    length_buf: Vec<u8>,
22    length: u16,
23    data_buf: Vec<u8>,
24}
25
26#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
27enum State {
28    Idle,
29    WaitLengthParse,
30    WaitDataParse,
31    WaitDataParsing,
32    WaitCRLFParse(ActionAfterCRLFParsed),
33}
34impl Default for State {
35    fn default() -> Self {
36        Self::Idle
37    }
38}
39
40#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
41enum ActionAfterCRLFParsed {
42    Continue,
43    Break,
44}
45
46impl ChunkedBodyParser {
47    pub fn new() -> Self {
48        Self {
49            length_buf: Vec::with_capacity(LENGTH_MAX_LEN),
50            data_buf: vec![0u8; DATA_DEFAULT_LEN],
51            ..Default::default()
52        }
53    }
54
55    pub fn with_data_buf(data_buf: Vec<u8>) -> Self {
56        Self {
57            length_buf: Vec::with_capacity(LENGTH_MAX_LEN),
58            data_buf,
59            ..Default::default()
60        }
61    }
62}
63
64//
65//
66//
67impl BodyParser for ChunkedBodyParser {
68    fn parse<R: BufRead>(
69        &mut self,
70        r: &mut R,
71        body_buf: &mut Vec<u8>,
72    ) -> Result<BodyParseOutput, BodyParseError> {
73        let mut take = r.take(0);
74        let mut parsed_num_bytes = 0_usize;
75
76        loop {
77            if self.state <= State::WaitLengthParse {
78                let end_bytes_len = 2_usize;
79                take.set_limit(LENGTH_MAX_LEN as u64 + end_bytes_len as u64);
80
81                self.length_buf.clear();
82                let n = take
83                    .read_until(LF, &mut self.length_buf)
84                    .map_err(BodyParseError::ReadError)?;
85
86                if n < end_bytes_len {
87                    return Ok(BodyParseOutput::Partial(parsed_num_bytes));
88                }
89                if !self.length_buf[..n].ends_with(&[LF]) {
90                    if n >= LENGTH_MAX_LEN {
91                        return Err(BodyParseError::TooLongChunksOfLength);
92                    } else {
93                        return Ok(BodyParseOutput::Partial(parsed_num_bytes));
94                    }
95                }
96                if !self.length_buf[..n - 1].ends_with(&[CR]) {
97                    return Err(BodyParseError::InvalidCRLF);
98                }
99                let length_bytes = &self.length_buf[..n - end_bytes_len];
100                let length_str = core::str::from_utf8(length_bytes)
101                    .map_err(|_| BodyParseError::InvalidChunksOfLength(None))?;
102                let length = u16::from_str_radix(length_str, 16)
103                    .map_err(|err| BodyParseError::InvalidChunksOfLength(Some(err)))?;
104
105                self.length = length;
106                parsed_num_bytes += n;
107
108                if length == 0 {
109                    self.state = State::WaitCRLFParse(ActionAfterCRLFParsed::Break);
110                } else {
111                    self.state = State::WaitDataParse;
112                }
113            }
114
115            if self.state <= State::WaitDataParsing {
116                take.set_limit(self.length as u64);
117
118                let n = take
119                    .read(&mut self.data_buf)
120                    .map_err(BodyParseError::ReadError)?;
121                body_buf.extend_from_slice(&self.data_buf[..n]);
122
123                self.length -= n as u16;
124                parsed_num_bytes += n;
125
126                if self.length == 0 {
127                    self.state = State::WaitCRLFParse(ActionAfterCRLFParsed::Continue);
128                } else {
129                    self.state = State::WaitDataParsing;
130
131                    return Ok(BodyParseOutput::Partial(parsed_num_bytes));
132                }
133            }
134
135            if let State::WaitCRLFParse(action) = &self.state {
136                let end_bytes_len = 2_usize;
137                take.set_limit(end_bytes_len as u64);
138
139                self.length_buf.clear();
140                let n = take
141                    .read_until(LF, &mut self.length_buf)
142                    .map_err(BodyParseError::ReadError)?;
143                if n < end_bytes_len {
144                    return Ok(BodyParseOutput::Partial(parsed_num_bytes));
145                }
146                if &self.length_buf[..n] != CRLF {
147                    return Err(BodyParseError::InvalidCRLF);
148                }
149                parsed_num_bytes += n;
150
151                match action {
152                    ActionAfterCRLFParsed::Continue => {
153                        self.state = State::WaitLengthParse;
154
155                        continue;
156                    }
157                    ActionAfterCRLFParsed::Break => {
158                        self.state = State::Idle;
159
160                        break Ok(BodyParseOutput::Completed(parsed_num_bytes));
161                    }
162                }
163            }
164
165            unreachable!()
166        }
167    }
168}