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
10pub(crate) struct DecodeResult {
11    pub(crate) body: Vec<u8>,
12    pub(crate) trailers: Vec<(HeaderName, HeaderValue)>,
13}
14
15pub enum DecodeEvent {
16    NeedMore,
17    Chunk(Vec<u8>),
18    Done(Vec<(HeaderName, HeaderValue)>),
19}
20
21enum State {
22    SizeLine,
23    Data(usize),
24    Trailers,
25    Done,
26}
27
28pub struct BodyDecoder {
29    state: State,
30    max_body: usize,
31    body_len: usize,
32}
33
34impl Default for BodyDecoder {
35    fn default() -> Self {
36        Self::new()
37    }
38}
39
40impl BodyDecoder {
41    pub fn new() -> Self {
42        Self::with_limit(MAX_BODY_SIZE)
43    }
44
45    pub fn with_limit(max_body: usize) -> Self {
46        Self {
47            state: State::SizeLine,
48            max_body,
49            body_len: 0,
50        }
51    }
52
53    pub fn decode(&mut self, buf: &[u8]) -> Result<(usize, DecodeEvent)> {
54        let mut pos = 0;
55
56        loop {
57            match self.state {
58                State::SizeLine => {
59                    let crlf = match Framing::find_crlf(&buf[pos..]) {
60                        Some(offset) => pos + offset,
61                        None => return Ok((pos, DecodeEvent::NeedMore)),
62                    };
63
64                    let size_str = std::str::from_utf8(&buf[pos..crlf])
65                        .map_err(|_| Error::BadRequest("Invalid chunk size encoding".into()))?;
66                    let size_str = size_str.split(';').next().unwrap_or("").trim();
67                    let chunk_size = usize::from_str_radix(size_str, 16).map_err(|_| {
68                        Error::BadRequest(format!("Invalid chunk size: {:?}", size_str).into())
69                    })?;
70
71                    if chunk_size > self.max_body || self.body_len + chunk_size > self.max_body {
72                        return Err(Error::PayloadTooLarge(
73                            "Chunked body exceeds size limit".into(),
74                        ));
75                    }
76
77                    pos = crlf + 2;
78                    if chunk_size == 0 {
79                        self.state = State::Trailers;
80                    } else {
81                        self.state = State::Data(chunk_size);
82                    }
83                }
84                State::Data(chunk_size) => {
85                    if buf.len() < pos + chunk_size + 2 {
86                        return Ok((pos, DecodeEvent::NeedMore));
87                    }
88
89                    let chunk = buf[pos..pos + chunk_size].to_vec();
90                    if &buf[pos + chunk_size..pos + chunk_size + 2] != b"\r\n" {
91                        return Err(Error::BadRequest("Invalid chunk terminator".into()));
92                    }
93
94                    pos += chunk_size + 2;
95                    self.body_len += chunk_size;
96                    self.state = State::SizeLine;
97                    return Ok((pos, DecodeEvent::Chunk(chunk)));
98                }
99                State::Trailers => match Framing::parse_trailers(&buf[pos..])? {
100                    Some((trailers, consumed)) => {
101                        pos += consumed;
102                        self.state = State::Done;
103                        return Ok((pos, DecodeEvent::Done(trailers)));
104                    }
105                    None => return Ok((pos, DecodeEvent::NeedMore)),
106                },
107                State::Done => return Ok((0, DecodeEvent::Done(Vec::new()))),
108            }
109        }
110    }
111}
112
113impl crate::http::codec::Parse {
114    pub(crate) fn try_decode_chunked(buf: &[u8]) -> Result<Option<DecodeResult>> {
115        Self::try_decode_chunked_limited(buf, MAX_BODY_SIZE)
116    }
117
118    pub(crate) fn try_decode_chunked_limited(
119        buf: &[u8],
120        max_body: usize,
121    ) -> Result<Option<DecodeResult>> {
122        let mut decoder = BodyDecoder::with_limit(max_body);
123        let mut body = Vec::new();
124        let mut input = buf;
125
126        loop {
127            let (consumed, event) = decoder.decode(input)?;
128            input = &input[consumed..];
129
130            match event {
131                DecodeEvent::NeedMore => return Ok(None),
132                DecodeEvent::Chunk(chunk) => body.extend_from_slice(&chunk),
133                DecodeEvent::Done(trailers) => {
134                    return Ok(Some(DecodeResult { body, trailers }));
135                }
136            }
137        }
138    }
139
140    pub fn chunked_body(buf: &[u8]) -> Result<Option<Vec<u8>>> {
141        Ok(Self::try_decode_chunked(buf)?.map(|r| r.body))
142    }
143
144    pub fn chunked_body_consumed(buf: &[u8], max_body: usize) -> Result<Option<(usize, Vec<u8>)>> {
145        let mut decoder = BodyDecoder::with_limit(max_body);
146        let mut body = Vec::new();
147        let mut input = buf;
148        let total = buf.len();
149
150        loop {
151            let (consumed, event) = decoder.decode(input)?;
152            input = &input[consumed..];
153
154            match event {
155                DecodeEvent::NeedMore => return Ok(None),
156                DecodeEvent::Chunk(chunk) => body.extend_from_slice(&chunk),
157                DecodeEvent::Done(_) => {
158                    let consumed_total = total - input.len();
159                    return Ok(Some((consumed_total, body)));
160                }
161            }
162        }
163    }
164}
sark_core/http/codec/decode/chunked/mod.rs

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