zenith_http1/

chunked.rs

//! HTTP/1.1 Chunked 传输编码解析
//!
//! 严格遵循 RFC 7230 §4.1：
//! ```text
//! chunked-body = *chunk
//!                last-chunk
//!                trailer-part
//!                CRLF
//! chunk = chunk-size [ chunk-extension ] CRLF
//!         chunk-data CRLF
//! chunk-size = 1*HEXDIG
//! ```
//!
//! 防护点：
//! - chunk-size 必须为十六进制数字（禁止符号/空格）
//! - 单个 chunk 不得超过 max_chunk_size（防止内存爆涨）
//! - 拒绝负值/零长度陷阱（0 必须是 last-chunk）
//! - 处理 trailer 时重置状态

use crate::types::Http1Error;

/// trailer 缓冲区最大大小（与 MAX_HEADER_SIZE 典型限制一致，防止恶意超长 trailer 耗尽内存）
const MAX_TRAILER_BUF_SIZE: usize = 8192;

/// Chunked 解析状态机
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ChunkedState {
    /// 等待下一个 chunk-size 行
    SizeLine,
    /// 正在读取 chunk-data
    ChunkData,
    /// 等待 chunk-data 结束 CRLF
    ChunkDataEnd,
    /// 读取 trailer（可选）
    Trailer,
    /// 完成
    Done,
}

/// Chunked 解码器
#[derive(Debug, Clone)]
pub struct ChunkedDecoder {
    /// 当前状态
    pub state: ChunkedState,
    /// 剩余 chunk 字节数
    pub remaining: u64,
    /// 当前 chunk 大小
    pub current_chunk_size: u64,
    /// 已解码总字节数
    pub total_decoded: u64,
    /// 最大允许字节数
    pub max_total: u64,
    /// 是否已遇到 0 大小块
    pub saw_last: bool,
    /// 临时缓冲区（size 行累积）
    size_buf: [u8; 64],
    size_len: usize,
    /// 临时缓冲区（trailer 行跨 feed 累积）：H-1 修复，与 size_buf 同语义，
    /// 防止 partial trailer line 跨 feed 丢失
    trailer_buf: Vec<u8>,
}

impl ChunkedDecoder {
    /// 创建解码器
    #[inline]
    pub fn new(max_total: u64) -> Self {
        Self {
            state: ChunkedState::SizeLine,
            remaining: 0,
            current_chunk_size: 0,
            total_decoded: 0,
            max_total,
            saw_last: false,
            size_buf: [0u8; 64],
            size_len: 0,
            trailer_buf: Vec::new(),
        }
    }

    /// 重置状态
    #[inline]
    pub fn reset(&mut self) {
        self.state = ChunkedState::SizeLine;
        self.remaining = 0;
        self.current_chunk_size = 0;
        self.total_decoded = 0;
        self.saw_last = false;
        self.size_len = 0;
        self.trailer_buf.clear();
    }

    /// 是否已完成
    #[inline]
    pub fn is_done(&self) -> bool {
        matches!(self.state, ChunkedState::Done)
    }

    /// 推入数据并返回产出的解码字节
    ///
    /// # 参数
    /// - `input`: 输入字节流
    ///
    /// # 返回
    /// - Ok((output_vec, consumed)): 解码出的字节与消耗的输入字节数
    /// - Err(Http1Error): 解码错误
    pub fn feed(&mut self, input: &[u8]) -> Result<(Vec<u8>, usize), Http1Error> {
        let mut output = Vec::with_capacity(input.len());

        // 若 size_buf 中有未完成的数据，先拼接进本次输入的前部
        if self.size_len > 0 {
            let buffered = self.size_len;
            // checked 算术：输入长度由对端控制，溢出 fail-closed
            let combined_len = buffered.checked_add(input.len()).ok_or_else(|| {
                Http1Error::ChunkedError("buffered input length overflow".into())
            })?;
            let mut combined = Vec::with_capacity(combined_len);
            combined.extend_from_slice(&self.size_buf[..buffered]);
            combined.extend_from_slice(input);
            self.size_len = 0;
            let consumed = self.feed_combined(&combined, &mut output)?;
            // 只计本次输入消耗的字节数
            let own_consumed = consumed.saturating_sub(buffered);
            return Ok((output, own_consumed));
        }

        // H-1: 若 trailer_buf 中有未完成的 trailer 数据，先拼接进本次输入的前部
        if !self.trailer_buf.is_empty() {
            let buffered = self.trailer_buf.len();
            let combined_len = buffered.checked_add(input.len()).ok_or_else(|| {
                Http1Error::ChunkedError("buffered trailer input length overflow".into())
            })?;
            let mut combined = Vec::with_capacity(combined_len);
            combined.extend_from_slice(&self.trailer_buf);
            combined.extend_from_slice(input);
            self.trailer_buf.clear();
            let consumed = self.feed_combined(&combined, &mut output)?;
            // 只计本次输入消耗的字节数
            let own_consumed = consumed.saturating_sub(buffered);
            return Ok((output, own_consumed));
        }

        let consumed = self.feed_combined(input, &mut output)?;
        Ok((output, consumed))
    }

    /// 内部解析（已合并 size_buf）
    fn feed_combined(
        &mut self,
        input: &[u8],
        output: &mut Vec<u8>,
    ) -> Result<usize, Http1Error> {
        let mut consumed = 0;

        while consumed < input.len() && !matches!(self.state, ChunkedState::Done) {
            match self.state {
                ChunkedState::SizeLine => {
                    let rest = &input[consumed..];
                    if let Some(crlf) = find_crlf(rest) {
                        let line = &rest[..crlf];
                        self.parse_size_line(line)?;
                        consumed += crlf + 2;
                    } else {
                        let available = rest.len().min(self.size_buf.len() - self.size_len);
                        if available > 0 {
                            self.size_buf[self.size_len..self.size_len + available]
                                .copy_from_slice(&rest[..available]);
                            self.size_len += available;
                            consumed += available;
                        }
                        if self.size_len >= self.size_buf.len() {
                            return Err(Http1Error::ChunkedError(
                                "chunk-size line too long".into(),
                            ));
                        }
                        break;
                    }
                }
                ChunkedState::ChunkData => {
                    let take = (input.len() - consumed) as u64;
                    let take = take.min(self.remaining);
                    if take == 0 {
                        // 需要更多数据
                        break;
                    }
                    output.extend_from_slice(
                        &input[consumed..consumed + take as usize],
                    );
                    consumed += take as usize;
                    self.remaining -= take;
                    // checked 算术：累计解码字节数溢出按超限处理（fail-closed）
                    self.total_decoded = self.total_decoded.checked_add(take).ok_or(Http1Error::BodyTooLarge)?;
                    if self.total_decoded > self.max_total {
                        return Err(Http1Error::BodyTooLarge);
                    }
                    if self.remaining == 0 {
                        self.state = ChunkedState::ChunkDataEnd;
                    }
                }
                ChunkedState::ChunkDataEnd => {
                    let rest = &input[consumed..];
                    if rest.len() < 2 {
                        break;
                    }
                    if rest[0] == b'\r' && rest[1] == b'\n' {
                        consumed += 2;
                        self.state = ChunkedState::SizeLine;
                    } else {
                        return Err(Http1Error::ChunkedError(
                            "missing CRLF after chunk-data".into(),
                        ));
                    }
                }
                ChunkedState::Trailer => {
                    let rest = &input[consumed..];
                    // 数据不足暂不消费：空 rest 等待下次喂入补齐
                    if rest.is_empty() {
                        break;
                    }
                    // 跳过 trailer 行直到遇到空行 (CRLF at pos 0)
                    if let Some(crlf) = find_crlf(rest) {
                        consumed += crlf + 2;
                        if crlf == 0 {
                            self.state = ChunkedState::Done;
                        }
                    } else {
                        // H-1: 未找到完整 CRLF，缓冲 partial trailer 行
                        // （含单独 \r 的半个 CRLF），下次 feed() 拼接后继续扫描，
                        // 禁止丢失跨 feed 的 partial bytes
                        self.trailer_buf.extend_from_slice(rest);
                        // 限制 trailer 缓冲区大小，防止恶意超长 trailer 耗尽内存
                        if self.trailer_buf.len() > MAX_TRAILER_BUF_SIZE {
                            return Err(Http1Error::ChunkedError(
                                "trailer buffer exceeds maximum size".into(),
                            ));
                        }
                        consumed += rest.len();
                        break;
                    }
                }
                // 循环条件已排除 Done；防御性 fail-closed，禁止 unreachable!() panic
                ChunkedState::Done => {
                    return Err(Http1Error::ChunkedError(
                        "chunked decoder in terminal state".into(),
                    ));
                }
            }
        }

        // 注意：输入耗尽停在 Trailer 不视为完成——终止空行（CRLF）必须真实到达，
        // 否则截断的 chunked body 会被误判为完整（fail-closed）

        Ok(consumed)
    }

    /// 解析 chunk-size 行
    ///
    /// 委托 [`crate::parser::parse_chunk_size`]（全 crate 唯一实现：分号
    /// chunk-ext 剥离、trim、显式 hexdigit 校验、溢出 fail-closed），
    /// 与服务端 feed 状态机语义严格一致，消除双解析器分歧。
    fn parse_size_line(&mut self, line: &[u8]) -> Result<(), Http1Error> {
        let parsed = crate::parser::parse_chunk_size(line)?;
        // usize → u64 显式形式化（32-bit 平台理论截断面 fail-closed；64-bit 无损）
        let val = u64::try_from(parsed).map_err(|_| Http1Error::BodyTooLarge)?;

        self.current_chunk_size = val;
        if val == 0 {
            self.saw_last = true;
            // RFC 7230 §4.1：last-chunk 的 CRLF 之后直接进入 trailer-part
            // （空 trailer 时下一个 CRLF 即终止空行；非空 trailer 逐行忽略）
            self.state = ChunkedState::Trailer;
        } else {
            if val > self.max_total {
                return Err(Http1Error::BodyTooLarge);
            }
            // checked 算术：total_decoded + val 溢出按超限处理（fail-closed）
            if self
                .total_decoded
                .checked_add(val)
                .map(|sum| sum > self.max_total)
                .unwrap_or(true)
            {
                return Err(Http1Error::BodyTooLarge);
            }
            self.remaining = val;
            self.state = ChunkedState::ChunkData;
        }
        Ok(())
    }
}

/// 查找 CRLF 位置
#[inline]
fn find_crlf(buf: &[u8]) -> Option<usize> {
    let mut i = 0;
    while i + 1 < buf.len() {
        if buf[i] == b'\r' && buf[i + 1] == b'\n' {
            return Some(i);
        }
        i += 1;
    }
    None
}

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

    #[test]
    fn test_simple_chunked() {
        let mut dec = ChunkedDecoder::new(1024);
        let input = b"5\r\nhello\r\n0\r\n\r\n";
        let (out, consumed) = dec.feed(input).unwrap();
        assert_eq!(out, b"hello");
        assert_eq!(consumed, input.len());
        assert!(dec.is_done());
    }

    #[test]
    fn test_multiple_chunks() {
        let mut dec = ChunkedDecoder::new(1024);
        let input = b"5\r\nhello\r\n6\r\n world\r\n0\r\n\r\n";
        let (out, consumed) = dec.feed(input).unwrap();
        assert_eq!(out, b"hello world");
        assert_eq!(consumed, input.len());
        assert!(dec.is_done());
    }

    #[test]
    fn test_incremental() {
        let mut dec = ChunkedDecoder::new(1024);
        let input = b"5\r\nhello\r\n0\r\n\r\n";
        let (out1, c1) = dec.feed(&input[..6]).unwrap();
        assert_eq!(c1, 6, "first feed should consume 6 bytes");
        let (out2, _c2) = dec.feed(&input[c1..]).unwrap();
        let mut all = out1;
        all.extend_from_slice(&out2);
        assert_eq!(all, b"hello");
        assert!(dec.is_done());
    }

    #[test]
    fn test_chunk_extension() {
        let mut dec = ChunkedDecoder::new(1024);
        let input = b"5;foo=bar\r\nhello\r\n0\r\n\r\n";
        let (out, _) = dec.feed(input).unwrap();
        assert_eq!(out, b"hello");
        assert!(dec.is_done());
    }

    #[test]
    fn test_invalid_hex() {
        let mut dec = ChunkedDecoder::new(1024);
        let input = b"ZZ\r\nx\r\n0\r\n\r\n";
        let r = dec.feed(input);
        assert!(r.is_err());
    }

    #[test]
    fn test_too_large() {
        let mut dec = ChunkedDecoder::new(2);
        let input = b"5\r\nhello\r\n0\r\n\r\n";
        let r = dec.feed(input);
        assert!(r.is_err());
    }

    #[test]
    fn test_missing_crlf() {
        let mut dec = ChunkedDecoder::new(1024);
        let input = b"5\r\nhello\rBAD";
        let r = dec.feed(input);
        assert!(r.is_err());
    }

    #[test]
    fn test_empty_chunk_zero_size() {
        let mut dec = ChunkedDecoder::new(1024);
        let input = b"0\r\n\r\n";
        let (out, consumed) = dec.feed(input).unwrap();
        assert!(out.is_empty());
        assert_eq!(consumed, input.len());
        assert!(dec.is_done());
        assert!(dec.saw_last);
    }

    #[test]
    fn test_single_byte_chunk() {
        let mut dec = ChunkedDecoder::new(1024);
        let input = b"1\r\nA\r\n0\r\n\r\n";
        let (out, _) = dec.feed(input).unwrap();
        assert_eq!(out, b"A");
        assert!(dec.is_done());
    }

    #[test]
    fn test_large_hex_chunk_size() {
        let mut dec = ChunkedDecoder::new(1024 * 1024);
        let size = 0xFF;
        let mut input = Vec::new();
        input.extend_from_slice(format!("{:X}\r\n", size).as_bytes());
        input.extend_from_slice(&vec![b'x'; size]);
        input.extend_from_slice(b"\r\n0\r\n\r\n");
        let (out, _) = dec.feed(&input).unwrap();
        assert_eq!(out.len(), size);
        assert!(dec.is_done());
    }

    #[test]
    fn test_chunk_size_with_trailing_spaces() {
        let mut dec = ChunkedDecoder::new(1024);
        let input = b"5   \r\nhello\r\n0\r\n\r\n";
        let (out, _) = dec.feed(input).unwrap();
        assert_eq!(out, b"hello");
        assert!(dec.is_done());
    }

    #[test]
    fn test_chunk_size_with_tabs() {
        let mut dec = ChunkedDecoder::new(1024);
        let input = b"5\t\r\nhello\r\n0\r\n\r\n";
        let (out, _) = dec.feed(input).unwrap();
        assert_eq!(out, b"hello");
        assert!(dec.is_done());
    }

    #[test]
    fn test_empty_chunk_size_line() {
        let mut dec = ChunkedDecoder::new(1024);
        let input = b"\r\nhello\r\n0\r\n\r\n";
        let r = dec.feed(input);
        assert!(r.is_err());
    }

    #[test]
    fn test_chunk_size_line_too_long() {
        let mut dec = ChunkedDecoder::new(1024);
        let mut long_line = vec![b'A'; 100];
        long_line.extend_from_slice(b"\r\n");
        let r = dec.feed(&long_line);
        assert!(r.is_err());
    }

    #[test]
    fn test_trailer_state_received() {
        let mut dec = ChunkedDecoder::new(1024);
        let input = b"5\r\nhello\r\n0\r\n\r\n";
        let (out, _) = dec.feed(input).unwrap();
        assert_eq!(out, b"hello");
        assert!(dec.is_done());
    }

    #[test]
    fn test_trailer_empty_immediate_done() {
        let mut dec = ChunkedDecoder::new(1024);
        let input = b"5\r\nhello\r\n0\r\n";
        dec.feed(input).unwrap();
        let _ = dec.feed(b"\r\n");
        assert!(dec.is_done());
    }

    #[test]
    fn test_reset_decoder() {
        let mut dec = ChunkedDecoder::new(1024);
        let input = b"5\r\nhello\r\n0\r\n\r\n";
        dec.feed(input).unwrap();
        assert!(dec.is_done());
        dec.reset();
        assert_eq!(dec.state, ChunkedState::SizeLine);
        assert_eq!(dec.total_decoded, 0);
        assert!(!dec.saw_last);
        assert_eq!(dec.current_chunk_size, 0);
    }

    #[test]
    fn test_chunked_state_variants() {
        let states = [
            ChunkedState::SizeLine,
            ChunkedState::ChunkData,
            ChunkedState::ChunkDataEnd,
            ChunkedState::Trailer,
            ChunkedState::Done,
        ];
        for (i, s) in states.iter().enumerate() {
            assert_eq!(*s, states[i]);
        }
        assert_ne!(ChunkedState::SizeLine, ChunkedState::Done);
    }

    #[test]
    fn test_decoder_clone() {
        let dec = ChunkedDecoder::new(1024);
        let dec2 = dec.clone();
        assert_eq!(dec.state, dec2.state);
        assert_eq!(dec.max_total, dec2.max_total);
    }

    #[test]
    fn test_total_decoded_tracking() {
        let mut dec = ChunkedDecoder::new(1024);
        let input = b"5\r\nhello\r\n6\r\n world\r\n0\r\n\r\n";
        let (_, _) = dec.feed(input).unwrap();
        assert_eq!(dec.total_decoded, 11);
    }

    #[test]
    fn test_uppercase_hex() {
        let mut dec = ChunkedDecoder::new(1024);
        let input = b"A\r\n0123456789\r\n0\r\n\r\n";
        let (out, _) = dec.feed(input).unwrap();
        assert_eq!(out.len(), 10);
        assert!(dec.is_done());
    }

    #[test]
    fn test_lowercase_hex() {
        let mut dec = ChunkedDecoder::new(1024);
        let input = b"a\r\n0123456789\r\n0\r\n\r\n";
        let (out, _) = dec.feed(input).unwrap();
        assert_eq!(out.len(), 10);
        assert!(dec.is_done());
    }

    #[test]
    fn test_negative_chunk_size_rejected() {
        let mut dec = ChunkedDecoder::new(1024);
        let input = b"-5\r\nhello\r\n0\r\n\r\n";
        let r = dec.feed(input);
        assert!(r.is_err());
    }

    #[test]
    fn test_body_too_large_accumulated() {
        let mut dec = ChunkedDecoder::new(10);
        let input = b"6\r\nhello \r\n6\r\nworld!\r\n0\r\n\r\n";
        let r = dec.feed(input);
        assert!(r.is_err());
    }

    #[test]
    fn test_total_decoded_overflow_fail_closed() {
        // 恶意输入：第二个 chunk-size 接近 u64::MAX，
        // total_decoded + val 溢出必须 fail-closed 返回 BodyTooLarge 而非回绕
        let mut dec = ChunkedDecoder::new(u64::MAX);
        // 0xFFFFFFFFFFFFFFFC = u64::MAX - 3；total_decoded(5) + val 溢出 u64
        let input = b"5\r\nhello\r\nFFFFFFFFFFFFFFFC\r\nx\r\n0\r\n\r\n";
        let r = dec.feed(input);
        assert!(
            matches!(r, Err(Http1Error::BodyTooLarge)),
            "total_decoded 溢出必须返回 BodyTooLarge，实际 {r:?}"
        );
    }

    #[test]
    fn test_incremental_small_chunks() {
        let mut dec = ChunkedDecoder::new(1024);
        let input = b"5\r\nhello\r\n0\r\n\r\n";
        let mut output = Vec::new();
        let mut pos = 0;
        while pos < input.len() {
            let end = (pos + 2).min(input.len());
            let (out, consumed) = dec.feed(&input[pos..end]).unwrap();
            output.extend_from_slice(&out);
            pos += consumed;
            if consumed == 0 {
                pos += 1;
            }
        }
        assert_eq!(output, b"hello");
        assert!(dec.is_done());
    }
}