use crate::{
body::{Decoder, Step, TrailerMode},
error::HttpError,
head::{BodyFraming, HeadScanner},
limits::Limits,
};
#[non_exhaustive]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum ChunkStep {
Data {
skip: usize,
len: usize,
},
NeedMore {
consumed: usize,
},
Trailers {
skip: usize,
len: usize,
},
Done {
consumed: usize,
},
}
pub struct ChunkedDecoder {
decoder: Decoder,
trailers: HeadScanner,
limits: Limits,
total: u64,
trailers_handed: bool,
}
impl core::fmt::Debug for ChunkedDecoder {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("ChunkedDecoder")
.field("total", &self.total)
.field("done", &self.is_done())
.field("limits", &self.limits)
.finish()
}
}
impl ChunkedDecoder {
#[must_use]
pub fn new(limits: &Limits) -> Self {
let mut decoder = Decoder::new();
decoder.start(BodyFraming::Chunked);
Self {
decoder,
trailers: HeadScanner::new(),
limits: *limits,
total: 0,
trailers_handed: false,
}
}
pub fn decode(&mut self, input: &[u8], eof: bool) -> Result<ChunkStep, HttpError> {
if self.trailers_handed {
self.trailers_handed = false;
self.decoder.take_trailers();
}
let mut consumed = 0;
let step = self.decoder.step_with(
input,
eof,
&self.limits,
&mut self.trailers,
TrailerMode::Locate,
&mut consumed,
)?;
Ok(match step {
Step::Data(len) => {
self.total += len as u64;
ChunkStep::Data {
skip: consumed,
len,
}
}
Step::Trailers => {
self.trailers_handed = true;
ChunkStep::Trailers {
skip: consumed,
len: self.decoder.trailer_len(),
}
}
Step::Done => ChunkStep::Done { consumed },
Step::NeedMore => ChunkStep::NeedMore { consumed },
})
}
#[must_use]
pub fn is_done(&self) -> bool {
!self.trailers_handed && self.decoder.is_done()
}
#[must_use]
pub fn total(&self) -> u64 {
self.total
}
}
#[cfg(test)]
mod tests {
use alloc::vec::Vec;
use super::*;
type Decoded = (Vec<u8>, Vec<u8>, Vec<u8>);
fn run(wire: &[u8], piece: usize, limits: &Limits) -> Result<Decoded, HttpError> {
let mut d = ChunkedDecoder::new(limits);
let mut buf: Vec<u8> = Vec::new();
let mut fed = 0;
let (mut body, mut trailers) = (Vec::new(), Vec::new());
loop {
let eof = fed == wire.len();
match d.decode(&buf, eof)? {
ChunkStep::Data { skip, len } => {
body.extend_from_slice(&buf[skip..skip + len]);
let _ = buf.drain(..skip + len);
}
ChunkStep::Trailers { skip, len } => {
trailers.extend_from_slice(&buf[skip..skip + len]);
let _ = buf.drain(..skip + len);
}
ChunkStep::Done { consumed } => {
let _ = buf.drain(..consumed);
buf.extend_from_slice(&wire[fed..]);
return Ok((body, trailers, buf));
}
ChunkStep::NeedMore { consumed } => {
let _ = buf.drain(..consumed);
let end = (fed + piece).min(wire.len());
buf.extend_from_slice(&wire[fed..end]);
fed = end;
}
}
}
}
#[test]
fn test_decode_any_split_gives_the_same_body_and_trailers() {
let wire = b"4;a=b\r\nWiki\r\n5\r\npedia\r\n0\r\nX-T: 1\r\nY: 2\r\n\r\nGET";
let whole = run(wire, wire.len(), &Limits::new()).unwrap();
assert_eq!(whole.0, b"Wikipedia");
assert_eq!(whole.1, b"X-T: 1\r\nY: 2\r\n\r\n");
assert_eq!(whole.2, b"GET");
for piece in 1..wire.len() {
assert_eq!(run(wire, piece, &Limits::new()).unwrap(), whole, "{piece}");
}
}
#[test]
fn test_decode_without_trailers_reports_done() {
let (body, trailers, rest) = run(b"1\r\nx\r\n0\r\n\r\n", 3, &Limits::new()).unwrap();
assert_eq!(body, b"x");
assert!(trailers.is_empty());
assert!(rest.is_empty());
}
#[test]
fn test_decode_errors_match_the_connection_grammar() {
let cases: &[(&[u8], HttpError)] = &[
(b"x\r\n", HttpError::InvalidChunkSize),
(b"00000000000000001\r\n", HttpError::InvalidChunkSize),
(b"5;\r\n", HttpError::InvalidChunkExtension),
(b"1\r\naX", HttpError::InvalidChunkTerminator),
(b"5;a=b\nX", HttpError::BareLf),
(b"0\r\nBad Name: x\r\n\r\n", HttpError::InvalidHeaderName),
(b"0\r\n folded\r\n\r\n", HttpError::ObsFold),
(b"3\r\nab", HttpError::UnexpectedEof),
];
for (wire, err) in cases {
assert_eq!(
run(wire, wire.len(), &Limits::new()).err(),
Some(*err),
"{wire:?}"
);
}
let limits = Limits::new().with_max_body(4);
assert_eq!(
run(b"5\r\nhello\r\n0\r\n\r\n", 64, &limits).err(),
Some(HttpError::BodyTooLarge)
);
let limits = Limits::new().with_max_trailer_fields(1);
assert_eq!(
run(b"0\r\nA: 1\r\nB: 2\r\n\r\n", 64, &limits).err(),
Some(HttpError::TooManyHeaders)
);
}
#[test]
fn test_total_and_done_track_the_body() {
let mut d = ChunkedDecoder::new(&Limits::new());
assert_eq!(
d.decode(b"2\r\nab\r\n", false),
Ok(ChunkStep::Data { skip: 3, len: 2 })
);
assert_eq!(d.total(), 2);
assert!(!d.is_done());
assert_eq!(
d.decode(b"\r\n0\r\nT: v\r\n\r\n", false),
Ok(ChunkStep::Trailers { skip: 5, len: 8 })
);
assert!(!d.is_done());
assert_eq!(d.decode(b"", false), Ok(ChunkStep::Done { consumed: 0 }));
assert!(d.is_done());
assert_eq!(
d.decode(b"more", false),
Ok(ChunkStep::Done { consumed: 0 })
);
}
}