sark-core 0.9.0

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<'a> {
    NeedMore,
    Chunk(&'a [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<'a>(&mut self, buf: &'a [u8]) -> Result<(usize, DecodeEvent<'a>)> {
        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];
                    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, o3::buffer::Shared)>> {
        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, o3::buffer::Shared::from(body))));
                }
            }
        }
    }
}