go_http/parse/
transfer.rs

1// SPDX-License-Identifier: Apache-2.0
2
3/// Body framing — port of Go net/http `readTransfer` / `writeTransfer`.
4///
5/// Resolves whether a message has a body and how it is framed (chunked vs
6/// content-length vs connection-close), then wraps the raw reader in the
7/// appropriate limiting/decoding reader.
8use std::io::{self, Read, Take};
9
10use super::{chunk::ChunkedReader, ParseError};
11use crate::header::Header;
12
13// ---------------------------------------------------------------------------
14// Body — the opaque body reader type
15// ---------------------------------------------------------------------------
16
17/// An HTTP message body.
18pub enum Body {
19    /// Body of exactly `n` bytes.
20    Limited(Take<Box<dyn Read + Send>>),
21    /// Chunked transfer-encoded body.
22    Chunked(ChunkedReader<Box<dyn Read + Send>>),
23    /// Body read until connection close (response only).
24    Unbounded(Box<dyn Read + Send>),
25    /// No body (HEAD, 204, 304, …).
26    Empty,
27    /// Any body wrapped with a hard byte cap; returns an error on overflow.
28    Capped { inner: Box<Body>, remaining: u64 },
29}
30
31impl Read for Body {
32    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
33        match self {
34            Self::Limited(r)   => r.read(buf),
35            Self::Chunked(r)   => r.read(buf),
36            Self::Unbounded(r) => r.read(buf),
37            Self::Empty        => Ok(0),
38            Self::Capped { inner, remaining } => {
39                if buf.is_empty() {
40                    return Ok(0);
41                }
42                if *remaining > 0 {
43                    let max = buf.len().min(*remaining as usize);
44                    let n = inner.read(&mut buf[..max])?;
45                    *remaining -= n as u64;
46                    return Ok(n);
47                }
48                // Cap is exhausted: check whether inner has naturally ended.
49                match inner.read(&mut buf[..1])? {
50                    0 => Ok(0),
51                    _ => Err(io::Error::other("request body too large")),
52                }
53            }
54        }
55    }
56}
57
58impl Body {
59    /// Read the entire body into a `Vec<u8>`.
60    ///
61    /// Used by the HTTP client to buffer the response body before releasing
62    /// the underlying `TcpStream` back to the connection pool.  Buffering
63    /// first ensures the stream is fully drained and that no `try_clone()`
64    /// alias of the stream is live when the next request starts, which would
65    /// cause concurrent reads on the same socket and SIGSEGV on Linux/epoll.
66    pub fn read_to_vec(&mut self) -> io::Result<Vec<u8>> {
67        let mut buf = Vec::new();
68        self.read_to_end(&mut buf)?;
69        Ok(buf)
70    }
71
72    /// Wrap this body with a hard byte cap.  Reads beyond `limit` bytes return
73    /// an `io::Error` rather than silently truncating.
74    pub fn capped(self, limit: u64) -> Self {
75        Body::Capped { inner: Box::new(self), remaining: limit }
76    }
77
78    /// Return trailer headers if this is a chunked body that has been fully
79    /// read.  Returns an empty header for all other body types.
80    pub fn trailers(&self) -> &Header {
81        static EMPTY: std::sync::OnceLock<Header> = std::sync::OnceLock::new();
82        match self {
83            Self::Chunked(r)              => &r.trailers,
84            Self::Capped { inner, .. }    => inner.trailers(),
85            _                             => EMPTY.get_or_init(Header::new),
86        }
87    }
88
89    /// Consume the body and return any trailer headers (chunked bodies only).
90    pub fn into_trailers(self) -> Header {
91        match self {
92            Self::Chunked(r)           => r.trailers,
93            Self::Capped { inner, .. } => inner.into_trailers(),
94            _                          => Header::new(),
95        }
96    }
97}
98
99// ---------------------------------------------------------------------------
100// Body presence / framing resolution
101// ---------------------------------------------------------------------------
102
103/// Who is asking: determines whether "no Content-Length, no TE" means a body
104/// (responses: read until close) or no body (requests).
105#[derive(Clone, Copy, PartialEq, Eq)]
106pub enum MessageKind {
107    Request,
108    Response { status: u16, method: Option<RequestMethod> },
109}
110
111/// The request method that prompted a response (used to decide body for HEAD).
112#[derive(Clone, Copy, PartialEq, Eq)]
113pub enum RequestMethod {
114    Head,
115    Connect,
116    Other,
117}
118
119/// Resolve the body reader for a message given its headers and kind.
120///
121/// Port of Go's `readTransfer`: chooses chunked, content-length, unbounded, or
122/// empty framing, then wraps `r` accordingly.
123///
124/// Returns `(body, trailers_key)` — `trailers_key` is `true` if the `Trailer`
125/// header was present, signaling that the caller should harvest trailer headers
126/// from the chunked reader after reading is complete.
127pub fn resolve_body(
128    r: Box<dyn Read + Send>,
129    headers: &Header,
130    kind: MessageKind,
131) -> Result<Body, ParseError> {
132    // RFC 7230 §3.3: certain responses never have a body.
133    if let MessageKind::Response { status, method } = kind {
134        let no_body = status == 204
135            || status == 304
136            || (100..200).contains(&status)
137            || method == Some(RequestMethod::Head)
138            || method == Some(RequestMethod::Connect) && (200..300).contains(&status);
139        if no_body {
140            return Ok(Body::Empty);
141        }
142    }
143
144    // Check Transfer-Encoding.
145    let te = headers.get("Transfer-Encoding").unwrap_or("").to_ascii_lowercase();
146    if te.contains("chunked") {
147        return Ok(Body::Chunked(ChunkedReader::new(r)));
148    }
149
150    // Check Content-Length.
151    if let Some(cl_str) = headers.get("Content-Length") {
152        let n: u64 = cl_str
153            .trim()
154            .parse()
155            .map_err(|_| ParseError::InvalidContentLength)?;
156        return Ok(Body::Limited(r.take(n)));
157    }
158
159    // For requests with no framing: no body.
160    // For responses with no framing: read until EOF.
161    match kind {
162        MessageKind::Request => Ok(Body::Empty),
163        MessageKind::Response { .. } => Ok(Body::Unbounded(r)),
164    }
165}
166
167/// Serialize body framing headers for an outgoing message.
168///
169/// Port of Go's `writeTransfer`: if the body length is known, writes
170/// `Content-Length`; otherwise selects chunked encoding and writes
171/// `Transfer-Encoding: chunked`.
172pub fn write_framing_headers(
173    headers: &mut Header,
174    content_length: Option<u64>,
175) {
176    if let Some(n) = content_length {
177        headers.set("Content-Length", n.to_string());
178        headers.del("Transfer-Encoding");
179    } else {
180        headers.set("Transfer-Encoding", "chunked");
181        headers.del("Content-Length");
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188    use std::io::{Cursor, Read};
189
190    fn boxed(b: &'static [u8]) -> Box<dyn Read + Send> {
191        Box::new(Cursor::new(b))
192    }
193
194    #[test]
195    fn content_length_body() {
196        let mut h = Header::new();
197        h.set("Content-Length", "5");
198        let mut body = resolve_body(boxed(b"Hello World"), &h, MessageKind::Request).unwrap();
199        let mut out = Vec::new();
200        body.read_to_end(&mut out).unwrap();
201        assert_eq!(out, b"Hello");
202    }
203
204    #[test]
205    fn chunked_body() {
206        let raw: &'static [u8] = b"5\r\nHello\r\n0\r\n\r\n";
207        let mut h = Header::new();
208        h.set("Transfer-Encoding", "chunked");
209        let mut body = resolve_body(boxed(raw), &h, MessageKind::Request).unwrap();
210        let mut out = Vec::new();
211        body.read_to_end(&mut out).unwrap();
212        assert_eq!(out, b"Hello");
213    }
214
215    #[test]
216    fn no_content_request() {
217        let h = Header::new();
218        let mut body = resolve_body(boxed(b"leftover"), &h, MessageKind::Request).unwrap();
219        let mut out = Vec::new();
220        body.read_to_end(&mut out).unwrap();
221        assert!(out.is_empty());
222    }
223
224    #[test]
225    fn no_body_204() {
226        let h = Header::new();
227        let mut body = resolve_body(
228            boxed(b"ignored"),
229            &h,
230            MessageKind::Response { status: 204, method: None },
231        )
232        .unwrap();
233        let mut out = Vec::new();
234        body.read_to_end(&mut out).unwrap();
235        assert!(out.is_empty());
236    }
237
238    #[test]
239    fn write_framing_content_length() {
240        let mut h = Header::new();
241        write_framing_headers(&mut h, Some(42));
242        assert_eq!(h.get("Content-Length"), Some("42"));
243        assert_eq!(h.get("Transfer-Encoding"), None);
244    }
245
246    #[test]
247    fn write_framing_chunked() {
248        let mut h = Header::new();
249        write_framing_headers(&mut h, None);
250        assert_eq!(h.get("Transfer-Encoding"), Some("chunked"));
251        assert_eq!(h.get("Content-Length"), None);
252    }
253}
go_http/parse/transfer.rs

go_http/parse/
transfer.rs
