sark_core/http/codec/decode/chunked/

mod.rs

mod parse;

use http::{HeaderName, HeaderValue};
use parse::Framing;

use crate::error::{Error, Result};

const MAX_BODY_SIZE: usize = 100 * 1024 * 1024;

fn parse_chunk_size(bytes: &[u8]) -> Result<usize> {
    if bytes.is_empty() {
        return Err(Error::BadRequest("Empty chunk size".into()));
    }
    let mut size: usize = 0;
    for &b in bytes {
        let digit = match b {
            b'0'..=b'9' => (b - b'0') as usize,
            b'a'..=b'f' => (b - b'a' + 10) as usize,
            b'A'..=b'F' => (b - b'A' + 10) as usize,
            _ => return Err(Error::BadRequest("Invalid chunk size".into())),
        };
        size = size
            .checked_mul(16)
            .and_then(|s| s.checked_add(digit))
            .ok_or_else(|| Error::BadRequest("Chunk size overflow".into()))?;
    }
    Ok(size)
}

pub(crate) struct DecodeResult {
    pub(crate) body: Vec<u8>,
    pub(crate) trailers: Vec<(HeaderName, HeaderValue)>,
}

pub enum DecodeEvent {
    NeedMore,
    Chunk(Vec<u8>),
    Done(Vec<(HeaderName, HeaderValue)>),
}

enum State {
    SizeLine,
    Data(usize),
    Trailers,
    Done,
}

pub struct BodyDecoder {
    state: State,
    max_body: usize,
    body_len: usize,
}

impl Default for BodyDecoder {
    fn default() -> Self {
        Self::new()
    }
}

impl BodyDecoder {
    pub fn new() -> Self {
        Self::with_limit(MAX_BODY_SIZE)
    }

    pub fn with_limit(max_body: usize) -> Self {
        Self {
            state: State::SizeLine,
            max_body,
            body_len: 0,
        }
    }

    pub fn decode(&mut self, buf: &[u8]) -> Result<(usize, DecodeEvent)> {
        let mut pos = 0;

        loop {
            match self.state {
                State::SizeLine => {
                    let crlf = match Framing::find_crlf(&buf[pos..]) {
                        Some(offset) => pos + offset,
                        None => return Ok((pos, DecodeEvent::NeedMore)),
                    };

                    let line = &buf[pos..crlf];
                    let size_bytes = match line.iter().position(|&b| b == b';') {
                        Some(semi) => &line[..semi],
                        None => line,
                    };
                    let chunk_size = parse_chunk_size(size_bytes)?;

                    if chunk_size > self.max_body || self.body_len + chunk_size > self.max_body {
                        return Err(Error::PayloadTooLarge(
                            "Chunked body exceeds size limit".into(),
                        ));
                    }

                    pos = crlf + 2;
                    if chunk_size == 0 {
                        self.state = State::Trailers;
                    } else {
                        self.state = State::Data(chunk_size);
                    }
                }
                State::Data(chunk_size) => {
                    if buf.len() < pos + chunk_size + 2 {
                        return Ok((pos, DecodeEvent::NeedMore));
                    }

                    let chunk = buf[pos..pos + chunk_size].to_vec();
                    if &buf[pos + chunk_size..pos + chunk_size + 2] != b"\r\n" {
                        return Err(Error::BadRequest("Invalid chunk terminator".into()));
                    }

                    pos += chunk_size + 2;
                    self.body_len += chunk_size;
                    self.state = State::SizeLine;
                    return Ok((pos, DecodeEvent::Chunk(chunk)));
                }
                State::Trailers => match Framing::parse_trailers(&buf[pos..])? {
                    Some((trailers, consumed)) => {
                        pos += consumed;
                        self.state = State::Done;
                        return Ok((pos, DecodeEvent::Done(trailers)));
                    }
                    None => return Ok((pos, DecodeEvent::NeedMore)),
                },
                State::Done => return Ok((0, DecodeEvent::Done(Vec::new()))),
            }
        }
    }
}

impl crate::http::codec::Parse {
    pub(crate) fn try_decode_chunked(buf: &[u8]) -> Result<Option<DecodeResult>> {
        Self::try_decode_chunked_limited(buf, MAX_BODY_SIZE)
    }

    pub(crate) fn try_decode_chunked_limited(
        buf: &[u8],
        max_body: usize,
    ) -> Result<Option<DecodeResult>> {
        let mut decoder = BodyDecoder::with_limit(max_body);
        let mut body = Vec::new();
        let mut input = buf;

        loop {
            let (consumed, event) = decoder.decode(input)?;
            input = &input[consumed..];

            match event {
                DecodeEvent::NeedMore => return Ok(None),
                DecodeEvent::Chunk(chunk) => body.extend_from_slice(&chunk),
                DecodeEvent::Done(trailers) => {
                    return Ok(Some(DecodeResult { body, trailers }));
                }
            }
        }
    }

    pub fn chunked_body(buf: &[u8]) -> Result<Option<Vec<u8>>> {
        Ok(Self::try_decode_chunked(buf)?.map(|r| r.body))
    }

    pub fn chunked_body_consumed(buf: &[u8], max_body: usize) -> Result<Option<(usize, Vec<u8>)>> {
        let mut decoder = BodyDecoder::with_limit(max_body);
        let mut body = Vec::new();
        let mut input = buf;
        let total = buf.len();

        loop {
            let (consumed, event) = decoder.decode(input)?;
            input = &input[consumed..];

            match event {
                DecodeEvent::NeedMore => return Ok(None),
                DecodeEvent::Chunk(chunk) => body.extend_from_slice(&chunk),
                DecodeEvent::Done(_) => {
                    let consumed_total = total - input.len();
                    return Ok(Some((consumed_total, body)));
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn chunk_size_plus_prefix_rejected() {
        assert!(parse_chunk_size(b"+a").is_err());
    }

    #[test]
    fn chunk_size_minus_prefix_rejected() {
        assert!(parse_chunk_size(b"-a").is_err());
    }

    #[test]
    fn chunk_size_surrounding_whitespace_rejected() {
        assert!(parse_chunk_size(b" 5 ").is_err());
        assert!(parse_chunk_size(b"\t5").is_err());
        assert!(parse_chunk_size(b"5 ").is_err());
    }

    #[test]
    fn chunk_size_empty_rejected() {
        assert!(parse_chunk_size(b"").is_err());
    }

    #[test]
    fn chunk_size_valid_hex_accepted() {
        assert_eq!(parse_chunk_size(b"a").unwrap(), 10);
        assert_eq!(parse_chunk_size(b"1F").unwrap(), 31);
        assert_eq!(parse_chunk_size(b"0").unwrap(), 0);
    }

    #[test]
    fn chunk_size_huge_hex_overflows() {
        assert!(parse_chunk_size(b"ffffffffffffffffff").is_err());
    }

    #[test]
    fn chunk_extension_after_size_accepted() {
        let raw = b"5;name=value\r\nhello\r\n0\r\n\r\n";
        let body = crate::http::codec::Parse::chunked_body(raw)
            .unwrap()
            .unwrap();
        assert_eq!(body, b"hello");
    }

    #[test]
    fn normal_chunked_body_accepted() {
        let raw = b"5\r\nhello\r\n0\r\n\r\n";
        let body = crate::http::codec::Parse::chunked_body(raw)
            .unwrap()
            .unwrap();
        assert_eq!(body, b"hello");
    }

    #[test]
    fn chunk_size_whitespace_in_stream_rejected() {
        let raw = b" 5 \r\nhello\r\n0\r\n\r\n";
        assert!(crate::http::codec::Parse::chunked_body(raw).is_err());
    }
}