sark_core/http/codec/decode/chunked/
mod.rs

1mod parse;
2
3use http::{HeaderName, HeaderValue};
4use parse::Framing;
5
6use crate::error::{Error, Result};
7
8const MAX_BODY_SIZE: usize = 100 * 1024 * 1024;
9
10fn parse_chunk_size(bytes: &[u8]) -> Result<usize> {
11    if bytes.is_empty() {
12        return Err(Error::BadRequest("Empty chunk size".into()));
13    }
14    let mut size: usize = 0;
15    for &b in bytes {
16        let digit = match b {
17            b'0'..=b'9' => (b - b'0') as usize,
18            b'a'..=b'f' => (b - b'a' + 10) as usize,
19            b'A'..=b'F' => (b - b'A' + 10) as usize,
20            _ => return Err(Error::BadRequest("Invalid chunk size".into())),
21        };
22        size = size
23            .checked_mul(16)
24            .and_then(|s| s.checked_add(digit))
25            .ok_or_else(|| Error::BadRequest("Chunk size overflow".into()))?;
26    }
27    Ok(size)
28}
29
30pub(crate) struct DecodeResult {
31    pub(crate) body: Vec<u8>,
32    pub(crate) trailers: Vec<(HeaderName, HeaderValue)>,
33}
34
35pub enum DecodeEvent<'a> {
36    NeedMore,
37    Chunk(&'a [u8]),
38    Done(Vec<(HeaderName, HeaderValue)>),
39}
40
41enum State {
42    SizeLine,
43    Data(usize),
44    Trailers,
45    Done,
46}
47
48pub struct BodyDecoder {
49    state: State,
50    max_body: usize,
51    body_len: usize,
52}
53
54impl Default for BodyDecoder {
55    fn default() -> Self {
56        Self::new()
57    }
58}
59
60impl BodyDecoder {
61    pub fn new() -> Self {
62        Self::with_limit(MAX_BODY_SIZE)
63    }
64
65    pub fn with_limit(max_body: usize) -> Self {
66        Self {
67            state: State::SizeLine,
68            max_body,
69            body_len: 0,
70        }
71    }
72
73    pub fn decode<'a>(&mut self, buf: &'a [u8]) -> Result<(usize, DecodeEvent<'a>)> {
74        let mut pos = 0;
75
76        loop {
77            match self.state {
78                State::SizeLine => {
79                    let crlf = match Framing::find_crlf(&buf[pos..]) {
80                        Some(offset) => pos + offset,
81                        None => return Ok((pos, DecodeEvent::NeedMore)),
82                    };
83
84                    let line = &buf[pos..crlf];
85                    let size_bytes = match line.iter().position(|&b| b == b';') {
86                        Some(semi) => &line[..semi],
87                        None => line,
88                    };
89                    let chunk_size = parse_chunk_size(size_bytes)?;
90
91                    if chunk_size > self.max_body || self.body_len + chunk_size > self.max_body {
92                        return Err(Error::PayloadTooLarge(
93                            "Chunked body exceeds size limit".into(),
94                        ));
95                    }
96
97                    pos = crlf + 2;
98                    if chunk_size == 0 {
99                        self.state = State::Trailers;
100                    } else {
101                        self.state = State::Data(chunk_size);
102                    }
103                }
104                State::Data(chunk_size) => {
105                    if buf.len() < pos + chunk_size + 2 {
106                        return Ok((pos, DecodeEvent::NeedMore));
107                    }
108
109                    let chunk = &buf[pos..pos + chunk_size];
110                    if &buf[pos + chunk_size..pos + chunk_size + 2] != b"\r\n" {
111                        return Err(Error::BadRequest("Invalid chunk terminator".into()));
112                    }
113
114                    pos += chunk_size + 2;
115                    self.body_len += chunk_size;
116                    self.state = State::SizeLine;
117                    return Ok((pos, DecodeEvent::Chunk(chunk)));
118                }
119                State::Trailers => match Framing::parse_trailers(&buf[pos..])? {
120                    Some((trailers, consumed)) => {
121                        pos += consumed;
122                        self.state = State::Done;
123                        return Ok((pos, DecodeEvent::Done(trailers)));
124                    }
125                    None => return Ok((pos, DecodeEvent::NeedMore)),
126                },
127                State::Done => return Ok((0, DecodeEvent::Done(Vec::new()))),
128            }
129        }
130    }
131}
132
133impl crate::http::codec::Parse {
134    pub(crate) fn try_decode_chunked(buf: &[u8]) -> Result<Option<DecodeResult>> {
135        Self::try_decode_chunked_limited(buf, MAX_BODY_SIZE)
136    }
137
138    pub(crate) fn try_decode_chunked_limited(
139        buf: &[u8],
140        max_body: usize,
141    ) -> Result<Option<DecodeResult>> {
142        let mut decoder = BodyDecoder::with_limit(max_body);
143        let mut body = Vec::new();
144        let mut input = buf;
145
146        loop {
147            let (consumed, event) = decoder.decode(input)?;
148            input = &input[consumed..];
149
150            match event {
151                DecodeEvent::NeedMore => return Ok(None),
152                DecodeEvent::Chunk(chunk) => body.extend_from_slice(chunk),
153                DecodeEvent::Done(trailers) => {
154                    return Ok(Some(DecodeResult { body, trailers }));
155                }
156            }
157        }
158    }
159
160    pub fn chunked_body(buf: &[u8]) -> Result<Option<Vec<u8>>> {
161        Ok(Self::try_decode_chunked(buf)?.map(|r| r.body))
162    }
163
164    pub fn chunked_body_consumed(
165        buf: &[u8],
166        max_body: usize,
167    ) -> Result<Option<(usize, o3::buffer::Shared)>> {
168        let mut decoder = BodyDecoder::with_limit(max_body);
169        let mut body = Vec::new();
170        let mut input = buf;
171        let total = buf.len();
172
173        loop {
174            let (consumed, event) = decoder.decode(input)?;
175            input = &input[consumed..];
176
177            match event {
178                DecodeEvent::NeedMore => return Ok(None),
179                DecodeEvent::Chunk(chunk) => body.extend_from_slice(chunk),
180                DecodeEvent::Done(_) => {
181                    let consumed_total = total - input.len();
182                    return Ok(Some((consumed_total, o3::buffer::Shared::from(body))));
183                }
184            }
185        }
186    }
187}
sark_core/http/codec/decode/chunked/mod.rs

sark_core/http/codec/decode/chunked/
mod.rs
