xitca_http/h1/proto/
trasnder_coding.rs

1use core::{fmt, mem};
2
3use std::io;
4
5use tracing::{trace, warn};
6
7use crate::{
8    bytes::{Buf, BufMut, Bytes, BytesMut},
9    http::header::{HeaderMap, HeaderName, HeaderValue},
10};
11
12use super::{buf_write::H1BufWrite, error::ProtoError};
13
14/// Maximum number of bytes allowed for all trailer fields.
15const TRAILER_MAX_HEADER_SIZE: usize = 1024 * 16;
16
17/// Coder for different Transfer-Decoding/Transfer-Encoding.
18#[derive(Clone, Debug, Eq, PartialEq)]
19pub enum TransferCoding {
20    /// Default coder indicates the Request/Response have successfully code it's associated body.
21    Eof,
22    /// Corrupted coder that can not be used anymore.
23    Corrupted,
24    /// Coder used when a Content-Length header is passed with a positive integer.
25    Length(u64),
26    /// Decoder used when Transfer-Encoding is `chunked`.
27    DecodeChunked {
28        state: ChunkedState,
29        size: u64,
30        trailers: Trailers,
31    },
32    /// Encoder for when Transfer-Encoding includes `chunked`.
33    EncodeChunked,
34    /// Upgrade type coder that pass through body as is without transforming.
35    Upgrade,
36}
37
38/// State for accumulating chunked trailer headers during decode.
39/// The `buf` is lazily boxed only when trailers are actually present.
40#[derive(Clone, Debug, Eq, PartialEq)]
41pub struct Trailers {
42    buf: Option<Box<BytesMut>>,
43    len: usize,
44    limit: usize,
45    size_limit: usize,
46}
47
48impl Trailers {
49    // TODO: expose size_limit to HttpServiceConfig
50    pub fn new(header_limit: usize) -> Self {
51        Self {
52            buf: None,
53            len: 0,
54            limit: header_limit,
55            size_limit: TRAILER_MAX_HEADER_SIZE,
56        }
57    }
58
59    fn try_put(&mut self, byte: u8) -> io::Result<()> {
60        if self.buf.is_some() {
61            self.put(byte)?;
62        }
63        Ok(())
64    }
65
66    fn put(&mut self, byte: u8) -> io::Result<()> {
67        if let Some(buf) = self.buf.as_deref_mut() {
68            buf.put_u8(byte);
69            if buf.len() > self.size_limit {
70                return Err(io::Error::new(
71                    io::ErrorKind::InvalidData,
72                    "chunk trailers bytes over limit",
73                ));
74            }
75        } else {
76            let mut buf = BytesMut::with_capacity(64);
77            buf.put_u8(byte);
78            self.buf = Some(Box::new(buf));
79        };
80
81        Ok(())
82    }
83
84    fn incr_len(&mut self) -> io::Result<()> {
85        self.len += 1;
86        if self.len > self.limit {
87            return Err(io::Error::new(
88                io::ErrorKind::InvalidData,
89                "chunk trailers count overflow",
90            ));
91        }
92        Ok(())
93    }
94
95    fn take(&mut self) -> Option<Self> {
96        self.buf.is_some().then(|| mem::replace(self, Trailers::new(0)))
97    }
98
99    fn decode(self) -> io::Result<HeaderMap> {
100        let buf = self.buf.expect("trailer buf must be initialized");
101        let mut headers = vec![httparse::EMPTY_HEADER; self.len];
102        match httparse::parse_headers(&buf, &mut headers) {
103            Ok(httparse::Status::Complete((_, parsed))) => {
104                let mut map = HeaderMap::with_capacity(parsed.len());
105                for header in parsed {
106                    let name = HeaderName::from_bytes(header.name.as_bytes())
107                        .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "invalid trailer header name"))?;
108                    let value = HeaderValue::from_bytes(header.value)
109                        .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "invalid trailer header value"))?;
110                    map.append(name, value);
111                }
112                Ok(map)
113            }
114            Ok(httparse::Status::Partial) => {
115                Err(io::Error::new(io::ErrorKind::InvalidInput, "partial trailer headers"))
116            }
117            Err(e) => Err(io::Error::new(io::ErrorKind::InvalidInput, e)),
118        }
119    }
120}
121
122impl TransferCoding {
123    #[inline]
124    pub const fn eof() -> Self {
125        Self::Eof
126    }
127
128    #[inline]
129    pub const fn length(len: u64) -> Self {
130        Self::Length(len)
131    }
132
133    #[inline]
134    pub fn decode_chunked(header_limit: usize) -> Self {
135        Self::DecodeChunked {
136            state: ChunkedState::Size,
137            size: 0,
138            trailers: Trailers::new(header_limit),
139        }
140    }
141
142    #[inline]
143    pub const fn encode_chunked() -> Self {
144        Self::EncodeChunked
145    }
146
147    #[inline]
148    pub const fn upgrade() -> Self {
149        Self::Upgrade
150    }
151
152    /// Check if Self is in EOF state. An EOF state means TransferCoding is ended gracefully
153    /// and can not decode any value. See [TransferCoding::decode] for detail.
154    #[inline]
155    pub fn is_eof(&self) -> bool {
156        match self {
157            Self::Eof => true,
158            Self::EncodeChunked => unreachable!("TransferCoding can't decide eof state when encoding chunked data"),
159            _ => false,
160        }
161    }
162
163    #[inline]
164    pub fn is_upgrade(&self) -> bool {
165        matches!(self, Self::Upgrade)
166    }
167}
168
169#[derive(Clone, Debug, Eq, PartialEq)]
170pub enum ChunkedState {
171    Size,
172    SizeLws,
173    Extension,
174    SizeLf,
175    Body,
176    BodyCr,
177    BodyLf,
178    Trailer,
179    TrailerLf,
180    EndCr,
181    EndLf,
182    End,
183}
184
185macro_rules! byte (
186    ($rdr:ident) => ({
187        if $rdr.len() > 0 {
188            let b = $rdr[0];
189            $rdr.advance(1);
190            b
191        } else {
192            return Ok(None);
193        }
194    })
195);
196
197impl ChunkedState {
198    pub fn step(
199        &mut self,
200        body: &mut BytesMut,
201        size: &mut u64,
202        buf: &mut Option<Bytes>,
203        trailers: &mut Trailers,
204    ) -> io::Result<Option<Self>> {
205        match *self {
206            Self::Size => Self::read_size(body, size),
207            Self::SizeLws => Self::read_size_lws(body),
208            Self::Extension => Self::read_extension(body),
209            Self::SizeLf => Self::read_size_lf(body, size),
210            Self::Body => Self::read_body(body, size, buf),
211            Self::BodyCr => Self::read_body_cr(body),
212            Self::BodyLf => Self::read_body_lf(body),
213            Self::Trailer => Self::read_trailer(body, trailers),
214            Self::TrailerLf => Self::read_trailer_lf(body, trailers),
215            Self::EndCr => Self::read_end_cr(body, trailers),
216            Self::EndLf => Self::read_end_lf(body, trailers),
217            Self::End => Ok(Some(Self::End)),
218        }
219    }
220
221    fn read_size(rdr: &mut BytesMut, size: &mut u64) -> io::Result<Option<Self>> {
222        macro_rules! or_overflow {
223            ($e:expr) => (
224                match $e {
225                    Some(val) => val,
226                    None => return Err(io::Error::new(
227                        io::ErrorKind::InvalidData,
228                        "invalid chunk size: overflow",
229                    )),
230                }
231            )
232        }
233
234        let radix = 16;
235        match byte!(rdr) {
236            b @ b'0'..=b'9' => {
237                *size = or_overflow!(size.checked_mul(radix));
238                *size = or_overflow!(size.checked_add((b - b'0') as u64));
239            }
240            b @ b'a'..=b'f' => {
241                *size = or_overflow!(size.checked_mul(radix));
242                *size = or_overflow!(size.checked_add((b + 10 - b'a') as u64));
243            }
244            b @ b'A'..=b'F' => {
245                *size = or_overflow!(size.checked_mul(radix));
246                *size = or_overflow!(size.checked_add((b + 10 - b'A') as u64));
247            }
248            b'\t' | b' ' => return Ok(Some(ChunkedState::SizeLws)),
249            b';' => return Ok(Some(ChunkedState::Extension)),
250            b'\r' => return Ok(Some(ChunkedState::SizeLf)),
251            _ => {
252                return Err(io::Error::new(
253                    io::ErrorKind::InvalidInput,
254                    "Invalid chunk size line: Invalid Size",
255                ));
256            }
257        }
258
259        Ok(Some(ChunkedState::Size))
260    }
261
262    fn read_size_lws(rdr: &mut BytesMut) -> io::Result<Option<Self>> {
263        match byte!(rdr) {
264            // LWS can follow the chunk size, but no more digits can come
265            b'\t' | b' ' => Ok(Some(Self::SizeLws)),
266            b';' => Ok(Some(Self::Extension)),
267            b'\r' => Ok(Some(Self::SizeLf)),
268            _ => Err(io::Error::new(
269                io::ErrorKind::InvalidInput,
270                "Invalid chunk size linear white space",
271            )),
272        }
273    }
274
275    fn read_extension(rdr: &mut BytesMut) -> io::Result<Option<Self>> {
276        match byte!(rdr) {
277            b'\r' => Ok(Some(Self::SizeLf)),
278            b'\n' => Err(io::Error::new(
279                io::ErrorKind::InvalidData,
280                "invalid chunk extension contains newline",
281            )),
282            _ => Ok(Some(Self::Extension)), // no supported extensions
283        }
284    }
285
286    fn read_size_lf(rdr: &mut BytesMut, size: &u64) -> io::Result<Option<Self>> {
287        match byte!(rdr) {
288            b'\n' if *size > 0 => Ok(Some(Self::Body)),
289            b'\n' if *size == 0 => Ok(Some(Self::EndCr)),
290            _ => Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid chunk size LF")),
291        }
292    }
293
294    fn read_body(rdr: &mut BytesMut, rem: &mut u64, buf: &mut Option<Bytes>) -> io::Result<Option<Self>> {
295        if rdr.is_empty() {
296            Ok(None)
297        } else {
298            *buf = Some(bounded_split(rem, rdr));
299            if *rem > 0 {
300                Ok(Some(Self::Body))
301            } else {
302                Ok(Some(Self::BodyCr))
303            }
304        }
305    }
306
307    fn read_body_cr(rdr: &mut BytesMut) -> io::Result<Option<Self>> {
308        match byte!(rdr) {
309            b'\r' => Ok(Some(Self::BodyLf)),
310            _ => Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid chunk body CR")),
311        }
312    }
313
314    fn read_body_lf(rdr: &mut BytesMut) -> io::Result<Option<Self>> {
315        match byte!(rdr) {
316            b'\n' => Ok(Some(Self::Size)),
317            _ => Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid chunk body LF")),
318        }
319    }
320
321    fn read_trailer(rdr: &mut BytesMut, trailers: &mut Trailers) -> io::Result<Option<Self>> {
322        trace!(target: "h1_decode", "read_trailer");
323        let byte = byte!(rdr);
324        trailers.put(byte)?;
325        match byte {
326            b'\r' => Ok(Some(Self::TrailerLf)),
327            _ => Ok(Some(Self::Trailer)),
328        }
329    }
330
331    fn read_trailer_lf(rdr: &mut BytesMut, trailers: &mut Trailers) -> io::Result<Option<Self>> {
332        let byte = byte!(rdr);
333        match byte {
334            b'\n' => {
335                trailers.incr_len()?;
336                trailers.put(byte)?;
337                Ok(Some(Self::EndCr))
338            }
339            _ => Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid trailer end LF")),
340        }
341    }
342
343    fn read_end_cr(rdr: &mut BytesMut, trailers: &mut Trailers) -> io::Result<Option<Self>> {
344        let byte = byte!(rdr);
345        match byte {
346            b'\r' => {
347                trailers.try_put(byte)?;
348                Ok(Some(Self::EndLf))
349            }
350            _ => {
351                trailers.put(byte)?;
352                Ok(Some(Self::Trailer))
353            }
354        }
355    }
356
357    fn read_end_lf(rdr: &mut BytesMut, trailers: &mut Trailers) -> io::Result<Option<Self>> {
358        let byte = byte!(rdr);
359        match byte {
360            b'\n' => {
361                trailers.try_put(byte)?;
362                Ok(Some(Self::End))
363            }
364            _ => Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid chunk end LF")),
365        }
366    }
367}
368
369impl TransferCoding {
370    pub fn try_set(&mut self, other: Self) -> Result<(), ProtoError> {
371        match (&self, &other) {
372            // multiple set to plain chunked is allowed. This can happen from Connect method
373            // and/or Connection header.
374            // skip set when the request body is zero length.
375            (TransferCoding::Upgrade, TransferCoding::Upgrade) | (_, TransferCoding::Length(0)) => Ok(()),
376            // multiple set to decoded chunked/content-length are forbidden.
377            // mutation between decoded chunked/content-length/plain chunked is forbidden.
378            (TransferCoding::Upgrade, _)
379            | (TransferCoding::DecodeChunked { .. }, _)
380            | (TransferCoding::Length(..), _) => Err(ProtoError::HeaderName),
381            _ => {
382                *self = other;
383                Ok(())
384            }
385        }
386    }
387
388    #[inline]
389    pub fn set_eof(&mut self) {
390        *self = Self::Eof;
391    }
392
393    #[inline]
394    pub fn set_corrupted(&mut self) {
395        *self = Self::Corrupted;
396    }
397
398    /// Encode message. Return `EOF` state of encoder
399    pub fn encode<W>(&mut self, mut bytes: Bytes, buf: &mut W)
400    where
401        W: H1BufWrite,
402    {
403        // Skip encode empty bytes.
404        // This is to avoid unnecessary extending on h1::proto::buf::ListBuf when user
405        // provided empty bytes by accident.
406        if bytes.is_empty() {
407            return;
408        }
409
410        match *self {
411            Self::Upgrade => buf.write_buf_bytes(bytes),
412            Self::EncodeChunked => buf.write_buf_bytes_chunked(bytes),
413            Self::Length(ref mut rem) => {
414                let len = bytes.len() as u64;
415                if *rem >= len {
416                    buf.write_buf_bytes(bytes);
417                    *rem -= len;
418                } else {
419                    let rem = mem::replace(rem, 0u64);
420                    buf.write_buf_bytes(bytes.split_to(rem as usize));
421                }
422            }
423            Self::Eof => warn!(target: "h1_encode", "TransferCoding::Eof should not encode response body"),
424            _ => unreachable!(),
425        }
426    }
427
428    /// Encode eof with optional trailer headers.
429    ///
430    /// For chunked encoding, forbidden trailer fields (per RFC 9110) are silently filtered out.
431    pub fn encode_eof<W>(&mut self, trailers: Option<HeaderMap>, buf: &mut W)
432    where
433        W: H1BufWrite,
434    {
435        match *self {
436            Self::Eof | Self::Upgrade | Self::Length(0) => {}
437            Self::EncodeChunked => match trailers {
438                Some(trailers) => buf.write_buf_trailers(trailers),
439                None => buf.write_buf_static(b"0\r\n\r\n"),
440            },
441            Self::Length(n) => unreachable!("UnexpectedEof for Length Body with {} remaining", n),
442            _ => unreachable!(),
443        }
444    }
445
446    /// decode body. See [ChunkResult] for detailed outcome.
447    pub fn decode(&mut self, src: &mut BytesMut) -> ChunkResult {
448        match *self {
449            // when decoder reaching eof state it would return ChunkResult::Eof and followed by
450            // ChunkResult::AlreadyEof if decode is called again.
451            // This multi stage behaviour is depended on by the caller to know the exact timing of
452            // when eof happens. (Expensive one time operations can be happening at Eof)
453            Self::Length(0)
454            | Self::DecodeChunked {
455                state: ChunkedState::End,
456                ..
457            } => {
458                *self = Self::Eof;
459                ChunkResult::OnEof
460            }
461            Self::Eof => ChunkResult::AlreadyEof,
462            Self::Corrupted => ChunkResult::Corrupted,
463            ref _this if src.is_empty() => ChunkResult::InsufficientData,
464            Self::Length(ref mut rem) => ChunkResult::Ok(bounded_split(rem, src)),
465            Self::Upgrade => ChunkResult::Ok(src.split().freeze()),
466            Self::DecodeChunked {
467                ref mut state,
468                ref mut size,
469                ref mut trailers,
470            } => {
471                loop {
472                    let mut buf = None;
473                    // advances the chunked state
474                    *state = match state.step(src, size, &mut buf, trailers) {
475                        Ok(Some(state)) => state,
476                        Ok(None) => return ChunkResult::InsufficientData,
477                        Err(e) => return ChunkResult::Err(e),
478                    };
479
480                    if matches!(state, ChunkedState::End) {
481                        if let Some(trailers) = trailers.take() {
482                            match trailers.decode() {
483                                Ok(headers) => return ChunkResult::Trailers(headers),
484                                Err(e) => return ChunkResult::Err(e),
485                            }
486                        }
487                        return self.decode(src);
488                    }
489
490                    if let Some(buf) = buf {
491                        return ChunkResult::Ok(buf);
492                    }
493                }
494            }
495            _ => unreachable!(),
496        }
497    }
498}
499
500#[derive(Debug)]
501pub enum ChunkResult {
502    /// non empty chunk data produced by coder.
503    Ok(Bytes),
504    /// trailer headers decoded from chunked transfer encoding.
505    Trailers(HeaderMap),
506    /// io error type produced by coder that can be bubbled up to upstream caller.
507    Err(io::Error),
508    /// insufficient data. More input bytes required.
509    InsufficientData,
510    /// coder reached EOF state and no more chunk can be produced.
511    OnEof,
512    /// coder already reached EOF state and no more chunk can be produced.
513    /// used to hint calling stop filling input buffer with more data and/or calling method again.
514    AlreadyEof,
515    /// see [TransferCoding::Corrupted].
516    Corrupted,
517}
518
519impl fmt::Display for ChunkResult {
520    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
521        match *self {
522            Self::Ok(_) => f.write_str("chunked data."),
523            Self::Trailers(_) => f.write_str("trailer headers."),
524            Self::Err(ref e) => fmt::Display::fmt(e, f),
525            Self::InsufficientData => f.write_str("no sufficient data. More input bytes required."),
526            Self::OnEof => f.write_str("coder reached EOF state. no more chunk can be produced."),
527            Self::AlreadyEof => f.write_str("coder already reached EOF state. no more chunk can be produced."),
528            Self::Corrupted => f.write_str("coder corrupted. can not be used anymore."),
529        }
530    }
531}
532
533impl From<io::Error> for ChunkResult {
534    fn from(e: io::Error) -> Self {
535        Self::Err(e)
536    }
537}
538
539fn bounded_split(rem: &mut u64, buf: &mut BytesMut) -> Bytes {
540    let len = buf.len() as u64;
541    if *rem >= len {
542        *rem -= len;
543        buf.split().freeze()
544    } else {
545        let rem = mem::replace(rem, 0);
546        buf.split_to(rem as usize).freeze()
547    }
548}
549
550#[cfg(test)]
551mod test {
552    use super::*;
553
554    #[test]
555    fn test_read_chunk_size() {
556        use std::io::ErrorKind::{InvalidData, InvalidInput, UnexpectedEof};
557
558        fn read(s: &str) -> u64 {
559            let mut state = ChunkedState::Size;
560            let rdr = &mut BytesMut::from(s);
561            let mut size = 0;
562            loop {
563                let result = state.step(rdr, &mut size, &mut None, &mut Trailers::new(64));
564                state = result.unwrap_or_else(|_| panic!("read_size failed for {s:?}")).unwrap();
565                if state == ChunkedState::Body || state == ChunkedState::EndCr {
566                    break;
567                }
568            }
569            size
570        }
571
572        fn read_err(s: &str, expected_err: io::ErrorKind) {
573            let mut state = ChunkedState::Size;
574            let rdr = &mut BytesMut::from(s);
575            let mut size = 0;
576            loop {
577                let result = state.step(rdr, &mut size, &mut None, &mut Trailers::new(64));
578                state = match result {
579                    Ok(Some(s)) => s,
580                    Ok(None) => return assert_eq!(expected_err, UnexpectedEof),
581                    Err(e) => {
582                        assert_eq!(
583                            expected_err,
584                            e.kind(),
585                            "Reading {:?}, expected {:?}, but got {:?}",
586                            s,
587                            expected_err,
588                            e.kind()
589                        );
590                        return;
591                    }
592                };
593                if state == ChunkedState::Body || state == ChunkedState::End {
594                    panic!("Was Ok. Expected Err for {s:?}");
595                }
596            }
597        }
598
599        assert_eq!(1, read("1\r\n"));
600        assert_eq!(1, read("01\r\n"));
601        assert_eq!(0, read("0\r\n"));
602        assert_eq!(0, read("00\r\n"));
603        assert_eq!(10, read("A\r\n"));
604        assert_eq!(10, read("a\r\n"));
605        assert_eq!(255, read("Ff\r\n"));
606        assert_eq!(255, read("Ff   \r\n"));
607        // Missing LF or CRLF
608        read_err("F\rF", InvalidInput);
609        read_err("F", UnexpectedEof);
610        // Invalid hex digit
611        read_err("X\r\n", InvalidInput);
612        read_err("1X\r\n", InvalidInput);
613        read_err("-\r\n", InvalidInput);
614        read_err("-1\r\n", InvalidInput);
615        // Acceptable (if not fully valid) extensions do not influence the size
616        assert_eq!(1, read("1;extension\r\n"));
617        assert_eq!(10, read("a;ext name=value\r\n"));
618        assert_eq!(1, read("1;extension;extension2\r\n"));
619        assert_eq!(1, read("1;;;  ;\r\n"));
620        assert_eq!(2, read("2; extension...\r\n"));
621        assert_eq!(3, read("3   ; extension=123\r\n"));
622        assert_eq!(3, read("3   ;\r\n"));
623        assert_eq!(3, read("3   ;   \r\n"));
624        // Invalid extensions cause an error
625        read_err("1 invalid extension\r\n", InvalidInput);
626        read_err("1 A\r\n", InvalidInput);
627        read_err("1;no CRLF", UnexpectedEof);
628        read_err("1;reject\nnewlines\r\n", InvalidData);
629        // Overflow
630        read_err("f0000000000000003\r\n", InvalidData);
631    }
632
633    #[test]
634    fn test_read_chunked_single_read() {
635        let mock_buf = &mut BytesMut::from("10\r\n1234567890abcdef\r\n0\r\n");
636
637        match TransferCoding::decode_chunked(64).decode(mock_buf) {
638            ChunkResult::Ok(buf) => {
639                assert_eq!(16, buf.len());
640                let result = String::from_utf8(buf.as_ref().to_vec()).expect("decode String");
641                assert_eq!("1234567890abcdef", &result);
642            }
643            state => panic!("{}", state),
644        }
645    }
646
647    #[test]
648    fn test_read_chunked_trailer_with_missing_lf() {
649        let mock_buf = &mut BytesMut::from("10\r\n1234567890abcdef\r\n0\r\nbad\r\r\n");
650
651        let mut decoder = TransferCoding::decode_chunked(64);
652
653        match decoder.decode(mock_buf) {
654            ChunkResult::Ok(_) => {}
655            state => panic!("{}", state),
656        }
657
658        match decoder.decode(mock_buf) {
659            ChunkResult::Err(e) => assert_eq!(e.kind(), io::ErrorKind::InvalidInput),
660            state => panic!("{}", state),
661        }
662    }
663
664    #[test]
665    fn test_read_chunked_after_eof() {
666        let mock_buf = &mut BytesMut::from("10\r\n1234567890abcdef\r\n0\r\n\r\n");
667        let mut decoder = TransferCoding::decode_chunked(64);
668
669        // normal read
670        match decoder.decode(mock_buf) {
671            ChunkResult::Ok(buf) => {
672                assert_eq!(16, buf.len());
673                let result = String::from_utf8(buf.as_ref().to_vec()).unwrap();
674                assert_eq!("1234567890abcdef", &result);
675            }
676            state => panic!("{}", state),
677        }
678
679        // eof read
680        match decoder.decode(mock_buf) {
681            ChunkResult::OnEof => {}
682            state => panic!("{}", state),
683        }
684
685        // already meet eof
686        match decoder.decode(mock_buf) {
687            ChunkResult::AlreadyEof => {}
688            state => panic!("{}", state),
689        }
690    }
691
692    #[test]
693    fn test_read_chunked_with_trailers() {
694        // chunked body: "Hello" (5 bytes) followed by two trailer headers
695        let mock_buf = &mut BytesMut::from(
696            "5\r\nHello\r\n0\r\nExpires: Wed, 21 Oct 2015 07:28:00 GMT\r\nX-Checksum: abc123\r\n\r\n",
697        );
698        let mut decoder = TransferCoding::decode_chunked(64);
699
700        // first decode yields the body data
701        match decoder.decode(mock_buf) {
702            ChunkResult::Ok(buf) => {
703                assert_eq!(buf.as_ref(), b"Hello");
704            }
705            state => panic!("expected data chunk, got: {}", state),
706        }
707
708        // second decode yields the trailer headers
709        match decoder.decode(mock_buf) {
710            ChunkResult::Trailers(headers) => {
711                assert_eq!(headers.len(), 2);
712                assert_eq!(headers.get("Expires").unwrap(), "Wed, 21 Oct 2015 07:28:00 GMT");
713                assert_eq!(headers.get("X-Checksum").unwrap(), "abc123");
714            }
715            state => panic!("expected trailers, got: {}", state),
716        }
717
718        // third decode yields eof
719        match decoder.decode(mock_buf) {
720            ChunkResult::OnEof => {}
721            state => panic!("expected OnEof, got: {}", state),
722        }
723    }
724
725    #[test]
726    fn encode_chunked() {
727        let mut encoder = TransferCoding::encode_chunked();
728        let dst = &mut BytesMut::default();
729
730        let msg1 = Bytes::from("foo bar");
731        encoder.encode(msg1, dst);
732
733        assert_eq!(dst.as_ref(), b"7\r\nfoo bar\r\n");
734
735        let msg2 = Bytes::from("baz quux herp");
736        encoder.encode(msg2, dst);
737
738        assert_eq!(dst.as_ref(), b"7\r\nfoo bar\r\nD\r\nbaz quux herp\r\n");
739
740        encoder.encode_eof(None, dst);
741
742        assert_eq!(dst.as_ref(), b"7\r\nfoo bar\r\nD\r\nbaz quux herp\r\n0\r\n\r\n");
743    }
744
745    #[test]
746    fn encode_length() {
747        let max_len = 8;
748        let mut encoder = TransferCoding::length(max_len as u64);
749
750        let dst = &mut BytesMut::default();
751
752        let msg1 = Bytes::from("foo bar");
753        encoder.encode(msg1, dst);
754
755        assert_eq!(dst.as_ref(), b"foo bar");
756
757        for _ in 0..8 {
758            let msg2 = Bytes::from("baz");
759            encoder.encode(msg2, dst);
760
761            assert_eq!(dst.as_ref().len(), max_len);
762            assert_eq!(dst.as_ref(), b"foo barb");
763        }
764
765        encoder.encode_eof(None, dst);
766        assert_eq!(dst.as_ref().len(), max_len);
767        assert_eq!(dst.as_ref(), b"foo barb");
768    }
769
770    #[test]
771    fn encode_chunked_with_trailers() {
772        let mut encoder = TransferCoding::encode_chunked();
773        let dst = &mut BytesMut::default();
774
775        let msg = Bytes::from("hello");
776        encoder.encode(msg, dst);
777
778        let mut trailers = HeaderMap::new();
779        trailers.insert("x-checksum", HeaderValue::from_static("abc123"));
780        trailers.insert("x-status", HeaderValue::from_static("ok"));
781        trailers.append("x-status", HeaderValue::from_static("done"));
782
783        encoder.encode_eof(Some(trailers), dst);
784
785        assert_eq!(
786            dst.as_ref(),
787            b"5\r\nhello\r\n0\r\nx-checksum: abc123\r\nx-status: ok\r\nx-status: done\r\n\r\n"
788        );
789    }
790
791    #[test]
792    fn encode_chunked_trailers_filters_forbidden() {
793        use crate::http::header;
794
795        let mut encoder = TransferCoding::encode_chunked();
796        let dst = &mut BytesMut::default();
797
798        let msg = Bytes::from("data");
799        encoder.encode(msg, dst);
800
801        let mut trailers = HeaderMap::new();
802        trailers.insert("x-checksum", HeaderValue::from_static("abc123"));
803        // forbidden fields per RFC 9110
804        trailers.insert(header::CONTENT_LENGTH, HeaderValue::from_static("99"));
805        trailers.insert(header::CONTENT_TYPE, HeaderValue::from_static("text/plain"));
806        trailers.insert(header::HOST, HeaderValue::from_static("example.com"));
807        trailers.insert(header::TRANSFER_ENCODING, HeaderValue::from_static("chunked"));
808
809        encoder.encode_eof(Some(trailers), dst);
810
811        // only x-checksum should remain
812        assert_eq!(dst.as_ref(), b"4\r\ndata\r\n0\r\nx-checksum: abc123\r\n\r\n");
813    }
814
815    #[test]
816    fn encode_chunked_trailers_empty_after_filter() {
817        use crate::http::header;
818
819        let mut encoder = TransferCoding::encode_chunked();
820        let dst = &mut BytesMut::default();
821
822        encoder.encode(Bytes::from("x"), dst);
823
824        let mut trailers = HeaderMap::new();
825        trailers.insert(header::CONTENT_LENGTH, HeaderValue::from_static("1"));
826
827        encoder.encode_eof(Some(trailers), dst);
828
829        // all trailers filtered, should still produce valid chunked terminator
830        assert_eq!(dst.as_ref(), b"1\r\nx\r\n0\r\n\r\n");
831    }
832}
xitca_http/h1/proto/trasnder_coding.rs

xitca_http/h1/proto/
trasnder_coding.rs
