trillium_http/

body.rs

use crate::{Headers, h2::H2Body, h3::H3Body};
use BodyType::{Empty, Static, Streaming};
use futures_lite::{AsyncRead, AsyncReadExt, io::Cursor, ready};
use pin_project_lite::pin_project;
use std::{
    borrow::Cow,
    fmt::{self, Debug, Formatter},
    io::{Error, Result},
    pin::Pin,
    task::{Context, Poll},
};
use sync_wrapper::SyncWrapper;

/// Trait for streaming body sources that can optionally produce trailers.
///
/// Implement this on types that compute trailer headers dynamically as the body
/// is read — for example, a hashing wrapper that produces a `Digest` trailer
/// after all bytes have been streamed.
///
/// For plain [`AsyncRead`] sources with no trailers, use [`Body::new_streaming`].
/// `BodySource` is only needed when trailers must be produced.
pub trait BodySource: AsyncRead + Send + 'static {
    /// Returns the trailers for this body, called after the body has been fully read.
    ///
    /// Implementations may clear internal state on this call; the result is
    /// only meaningful after [`AsyncRead::poll_read`] has returned `Ok(0)`.
    fn trailers(self: Pin<&mut Self>) -> Option<Headers>;
}

pin_project! {
    struct PlainBody<T> {
        #[pin]
        async_read: T,
    }
}

impl<T: AsyncRead> AsyncRead for PlainBody<T> {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<Result<usize>> {
        self.project().async_read.poll_read(cx, buf)
    }
}

impl<T: AsyncRead + Send + 'static> BodySource for PlainBody<T> {
    fn trailers(self: Pin<&mut Self>) -> Option<Headers> {
        None
    }
}

/// The trillium representation of a http body. This can contain
/// either `&'static [u8]` content, `Vec<u8>` content, or a boxed
/// [`AsyncRead`]/[`BodySource`] type.
#[derive(Debug, Default)]
pub struct Body(pub(crate) BodyType);

impl Body {
    /// Construct a new body from a streaming [`AsyncRead`] source. If
    /// you have the body content in memory already, prefer
    /// [`Body::new_static`] or one of the From conversions.
    pub fn new_streaming(async_read: impl AsyncRead + Send + 'static, len: Option<u64>) -> Self {
        Self::new_with_trailers(PlainBody { async_read }, len)
    }

    /// Construct a new body from a [`BodySource`] that can produce trailers after
    /// the body has been fully read.
    ///
    /// Use this when trailers must be computed dynamically from the body bytes,
    /// for example to append a content hash.
    pub fn new_with_trailers(body: impl BodySource, len: Option<u64>) -> Self {
        Self(Streaming {
            async_read: SyncWrapper::new(Box::pin(body)),
            len,
            done: false,
            progress: 0,
            chunked_framing: true,
        })
    }

    /// Disable RFC 9112 chunked-encoding framing emitted by [`AsyncRead`] for streaming
    /// bodies of unknown length.
    ///
    /// By default, when a streaming body has no known length, this type's [`AsyncRead`]
    /// implementation emits the wire-format chunked framing (chunk-size prefix, terminating
    /// `0\r\n` marker) so the h1 codec can write its bytes directly. That framing is wrong
    /// for any consumer that wants raw body bytes — e.g., installing a Body as the override
    /// source on a client `Conn` for cache replay or middleware tee.
    ///
    /// This method is `#[doc(hidden)]` and `unstable`-feature-gated; it exists for internal
    /// use by trillium-client. External code has no reason to set this flag.
    #[doc(hidden)]
    #[cfg(feature = "unstable")]
    #[must_use]
    pub fn without_chunked_framing(mut self) -> Self {
        if let Streaming {
            ref mut chunked_framing,
            ..
        } = self.0
        {
            *chunked_framing = false;
        }
        self
    }

    pub(crate) fn ensure_chunked_framing(&mut self) -> &mut Self {
        if let Streaming {
            ref mut chunked_framing,
            ..
        } = self.0
        {
            *chunked_framing = true;
        }

        self
    }

    /// Returns trailers from the body source, if any.
    ///
    /// Only meaningful after the body has been fully read (i.e., [`AsyncRead::poll_read`]
    /// has returned `Ok(0)`). Returns `None` for bodies constructed with
    /// [`Body::new_streaming`] or [`Body::new_static`].
    #[doc(hidden)] // this isn't really a user-facing interface
    pub fn trailers(&mut self) -> Option<Headers> {
        match &mut self.0 {
            Streaming {
                async_read, done, ..
            } if *done => async_read.get_mut().as_mut().trailers(),
            _ => None,
        }
    }

    /// Construct a fixed-length Body from a `Vec<u8>` or `&'static
    /// [u8]`.
    pub fn new_static(content: impl Into<Cow<'static, [u8]>>) -> Self {
        Self(Static {
            content: content.into(),
            cursor: 0,
        })
    }

    /// Retrieve a borrow of the static content in this body. If this
    /// body is a streaming body or an empty body, this will return
    /// None.
    pub fn static_bytes(&self) -> Option<&[u8]> {
        match &self.0 {
            Static { content, .. } => Some(content.as_ref()),
            _ => None,
        }
    }

    /// Transform this Body into a dyn [`AsyncRead`]. This will wrap
    /// static content in a [`Cursor`]. Note that this is different
    /// from reading directly from the Body, which includes chunked
    /// encoding.
    pub fn into_reader(self) -> Pin<Box<dyn AsyncRead + Send + Sync + 'static>> {
        match self.0 {
            Streaming { async_read, .. } => Box::pin(SyncAsyncReader(async_read)),
            Static { content, .. } => Box::pin(Cursor::new(content)),
            Empty => Box::pin(Cursor::new("")),
        }
    }

    /// Consume this body and return the full content. If the body was
    /// constructed with [`Body::new_streaming`], this will read the
    /// entire streaming body into memory, awaiting the streaming
    /// source's completion. This function will return an error if a
    /// streaming body has already been partially or fully read.
    ///
    /// # Errors
    ///
    /// This returns an error variant if either of the following conditions are met:
    ///
    /// there is an io error when reading from the underlying transport such as a disconnect
    /// the body has already been read to completion
    pub async fn into_bytes(self) -> Result<Cow<'static, [u8]>> {
        match self.0 {
            Static { content, .. } => Ok(content),

            Streaming {
                async_read,
                len,
                progress: 0,
                done: false,
                ..
            } => {
                let mut async_read = async_read.into_inner();
                let mut buf = len
                    .and_then(|c| c.try_into().ok())
                    .map(Vec::with_capacity)
                    .unwrap_or_default();

                async_read.read_to_end(&mut buf).await?;

                Ok(Cow::Owned(buf))
            }

            Empty => Ok(Cow::Borrowed(b"")),

            Streaming { .. } => Err(Error::other("body already read to completion")),
        }
    }

    /// Retrieve the number of bytes that have been read from this
    /// body
    pub fn bytes_read(&self) -> u64 {
        self.0.bytes_read()
    }

    /// returns the content length of this body, if known and
    /// available.
    pub fn len(&self) -> Option<u64> {
        self.0.len()
    }

    /// determine if the this body represents no data
    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    /// determine if the this body represents static content
    pub fn is_static(&self) -> bool {
        matches!(self.0, Static { .. })
    }

    /// determine if the this body represents streaming content
    pub fn is_streaming(&self) -> bool {
        matches!(self.0, Streaming { .. })
    }

    /// Attempt to clone this body. Returns `None` for streaming bodies, which are one-shot.
    ///
    /// Static bodies (constructed via [`Body::new_static`] or any `From` conversion for
    /// `Vec<u8>`, `&'static [u8]`, `String`, `&'static str`, etc.) clone cheaply — just a
    /// `Cow` clone, which is a pointer copy for borrowed `&'static` content and a `Vec` clone
    /// for owned content. The clone resets read progress, so it can be sent again from the
    /// beginning.
    ///
    /// Empty bodies always clone successfully.
    ///
    /// This is useful for client middleware that needs to retransmit a body — e.g., redirect
    /// handlers, retry handlers, or auth-refresh handlers.
    #[doc(hidden)]
    #[cfg(feature = "unstable")]
    pub fn try_clone(&self) -> Option<Self> {
        match &self.0 {
            Empty => Some(Self::default()),
            Static { content, .. } => Some(Self(Static {
                content: content.clone(),
                cursor: 0,
            })),
            Streaming { .. } => None,
        }
    }

    /// Convert this body into an `H3Body` for reading
    #[cfg(feature = "unstable")]
    pub fn into_h3(self) -> H3Body {
        H3Body::new(self)
    }

    /// Convert this body into an `H3Body` for reading
    #[cfg(not(feature = "unstable"))]
    pub(crate) fn into_h3(self) -> H3Body {
        H3Body::new(self)
    }

    /// Convert this body into an [`H2Body`] for reading by the h2 send pump.
    ///
    /// h2 frames DATA at the connection layer, so the body bytes that reach the send pump
    /// must be plain payload — not chunk-encoded. [`H2Body`] strips the chunked-transfer
    /// wrapping that [`Body::poll_read`] applies for the h1 path on streaming bodies of
    /// unknown length, and forwards trailers so the send pump can emit trailing HEADERS.
    pub(crate) fn into_h2(self) -> H2Body {
        H2Body::new(self)
    }
}

#[allow(
    clippy::cast_sign_loss,
    clippy::cast_possible_truncation,
    clippy::cast_precision_loss
)]
fn max_bytes_to_read(buf_len: usize) -> usize {
    assert!(
        buf_len >= 6,
        "buffers of length {buf_len} are too small for this implementation.
            if this is a problem for you, please open an issue"
    );

    // #[allow(clippy::cast_precision_loss)] applied to the function
    // is for this line. We do not expect our buffers to be on the
    // order of petabytes, so we will not fall outside of the range of
    // integers that can be represented by f64
    let bytes_remaining_after_two_cr_lns = (buf_len - 4) as f64;

    // #[allow(clippy::cast_sign_loss)] applied to the function is for
    // this line. This is ok because we know buf_len is already a
    // usize and we are just converting it to an f64 in order to do
    // float log2(x)/4
    //
    // the maximum number of bytes that the hex representation of remaining bytes might take
    let max_bytes_of_hex_framing = (bytes_remaining_after_two_cr_lns).log2() / 4f64;

    // #[allow(clippy::cast_sign_loss)] applied to the function is for
    // this line.  This is ok because max_bytes_of_hex_framing will
    // always be smaller than bytes_remaining_after_two_cr_lns, and so
    // there is no risk of sign loss
    (bytes_remaining_after_two_cr_lns - max_bytes_of_hex_framing.ceil()) as usize
}

impl AsyncRead for Body {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<Result<usize>> {
        match &mut self.0 {
            Empty => Poll::Ready(Ok(0)),
            Static { content, cursor } => {
                let length = content.len();
                if length == *cursor {
                    return Poll::Ready(Ok(0));
                }
                let bytes = (length - *cursor).min(buf.len());
                buf[0..bytes].copy_from_slice(&content[*cursor..*cursor + bytes]);
                *cursor += bytes;
                Poll::Ready(Ok(bytes))
            }

            Streaming {
                async_read,
                len: Some(len),
                done,
                progress,
                ..
            } => {
                if *done {
                    return Poll::Ready(Ok(0));
                }

                let max_bytes_to_read = (*len - *progress)
                    .try_into()
                    .unwrap_or(buf.len())
                    .min(buf.len());

                let bytes = ready!(
                    async_read
                        .get_mut()
                        .as_mut()
                        .poll_read(cx, &mut buf[..max_bytes_to_read])
                )?;

                if bytes == 0 {
                    *done = true;
                } else {
                    *progress += bytes as u64;
                }

                Poll::Ready(Ok(bytes))
            }

            Streaming {
                async_read,
                len: None,
                done,
                progress,
                chunked_framing,
            } => {
                if *done {
                    return Poll::Ready(Ok(0));
                }

                if !*chunked_framing {
                    let bytes = ready!(async_read.get_mut().as_mut().poll_read(cx, buf))?;
                    if bytes == 0 {
                        *done = true;
                    } else {
                        *progress += bytes as u64;
                    }
                    return Poll::Ready(Ok(bytes));
                }

                let max_bytes_to_read = max_bytes_to_read(buf.len());

                let bytes = ready!(
                    async_read
                        .get_mut()
                        .as_mut()
                        .poll_read(cx, &mut buf[..max_bytes_to_read])
                )?;

                if bytes == 0 {
                    *done = true;
                    // Write only the last-chunk marker (`0\r\n`). The caller must then
                    // emit the trailer-section (possibly empty) followed by the
                    // terminating `\r\n` to complete RFC 9112 §7.1.2 chunked framing.
                    //
                    // This split is structural, not a missed opportunity to encapsulate:
                    //   * Trailers come from `BodySource::trailers() -> Option<Headers>` after EOF,
                    //     not from this `AsyncRead` path. They are structured `Headers` data, not
                    //     bytes.
                    //   * Formatting them needs `HttpContext` config (e.g.
                    //     `panic_on_invalid_response_headers`) that `Body` does not carry, and
                    //     reuses the same `write_headers_or_trailers` helper used for the
                    //     response-header section.
                    //   * Trailers can be arbitrarily large; emitting them from inside `poll_read`
                    //     would force a multi-poll state machine to span buffers. The caller writes
                    //     them in one shot via `BufWriter::buffer_mut()`, which has no such
                    //     constraint.
                    //
                    // Caller stitch lives in `conn/h1.rs::Conn::send` after the
                    // `bufwriter.copy_from(&mut body, ...)` drain.
                    buf[..3].copy_from_slice(b"0\r\n");
                    return Poll::Ready(Ok(3));
                }

                *progress += bytes as u64;

                let start = format!("{bytes:X}\r\n");
                let start_length = start.len();
                let total = bytes + start_length + 2;
                buf.copy_within(..bytes, start_length);
                buf[..start_length].copy_from_slice(start.as_bytes());
                buf[total - 2..total].copy_from_slice(b"\r\n");
                Poll::Ready(Ok(total))
            }
        }
    }
}

struct SyncAsyncReader(SyncWrapper<Pin<Box<dyn BodySource>>>);
impl Debug for SyncAsyncReader {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        f.debug_struct("SyncAsyncReader").finish()
    }
}
impl AsyncRead for SyncAsyncReader {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<Result<usize>> {
        self.get_mut().0.get_mut().as_mut().poll_read(cx, buf)
    }
}

#[derive(Default)]
pub(crate) enum BodyType {
    #[default]
    Empty,

    Static {
        content: Cow<'static, [u8]>,
        cursor: usize,
    },

    Streaming {
        async_read: SyncWrapper<Pin<Box<dyn BodySource>>>,
        progress: u64,
        len: Option<u64>,
        done: bool,
        /// When true (the default), [`Body`]'s [`AsyncRead`] impl emits RFC 9112
        /// chunked-encoding framing for the `len: None` case so the h1 codec can
        /// write the bytes directly. When false (set via
        /// [`Body::without_chunked_framing`]), the same path passes through raw
        /// bytes from the inner [`BodySource`].
        chunked_framing: bool,
    },
}

impl Debug for BodyType {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            Empty => f.debug_tuple("BodyType::Empty").finish(),
            Static { content, cursor } => f
                .debug_struct("BodyType::Static")
                .field("content", &String::from_utf8_lossy(content))
                .field("cursor", cursor)
                .finish(),
            Streaming {
                len,
                done,
                progress,
                ..
            } => f
                .debug_struct("BodyType::Streaming")
                .field("async_read", &format_args!(".."))
                .field("len", &len)
                .field("done", &done)
                .field("progress", &progress)
                .finish(),
        }
    }
}

impl BodyType {
    fn is_empty(&self) -> bool {
        match *self {
            Empty => true,
            Static { ref content, .. } => content.is_empty(),
            Streaming { len, .. } => len == Some(0),
        }
    }

    fn len(&self) -> Option<u64> {
        match *self {
            Empty => Some(0),
            Static { ref content, .. } => Some(content.len() as u64),
            Streaming { len, .. } => len,
        }
    }

    fn bytes_read(&self) -> u64 {
        match *self {
            Empty => 0,
            Static { cursor, .. } => cursor as u64,
            Streaming { progress, .. } => progress,
        }
    }
}

impl From<String> for Body {
    fn from(s: String) -> Self {
        s.into_bytes().into()
    }
}

impl From<&'static str> for Body {
    fn from(s: &'static str) -> Self {
        s.as_bytes().into()
    }
}

impl From<&'static [u8]> for Body {
    fn from(content: &'static [u8]) -> Self {
        Self::new_static(content)
    }
}

impl From<Vec<u8>> for Body {
    fn from(content: Vec<u8>) -> Self {
        Self::new_static(content)
    }
}

impl From<Cow<'static, [u8]>> for Body {
    fn from(value: Cow<'static, [u8]>) -> Self {
        Self::new_static(value)
    }
}

impl From<Cow<'static, str>> for Body {
    fn from(value: Cow<'static, str>) -> Self {
        match value {
            Cow::Borrowed(b) => b.into(),
            Cow::Owned(o) => o.into(),
        }
    }
}

#[cfg(test)]
mod test_bytes_to_read {
    #[test]
    fn simple_check_of_known_values() {
        // the marked rows are the most important part of this test,
        // and a nonobvious but intentional consequence of the
        // implementation. in order to avoid overflowing, we must use
        // one fewer than the available buffer bytes because
        // increasing the read size increase the number of framed
        // bytes by two. This occurs when the hex representation of
        // the content bytes is near an increase in order of magnitude
        // (F->10, FF->100, FFF-> 1000, etc)
        let values = vec![
            (6, 1),       // 1
            (7, 2),       // 2
            (20, 15),     // F
            (21, 15),     // F <-
            (22, 16),     // 10
            (23, 17),     // 11
            (260, 254),   // FE
            (261, 254),   // FE <-
            (262, 255),   // FF <-
            (263, 256),   // 100
            (4100, 4093), // FFD
            (4101, 4093), // FFD <-
            (4102, 4094), // FFE <-
            (4103, 4095), // FFF <-
            (4104, 4096), // 1000
        ];

        for (input, expected) in values {
            let actual = super::max_bytes_to_read(input);
            assert_eq!(
                actual, expected,
                "\n\nexpected max_bytes_to_read({input}) to be {expected}, but it was {actual}"
            );

            // testing the test:
            let used_bytes = expected + 4 + format!("{expected:X}").len();
            assert!(
                used_bytes == input || used_bytes == input - 1,
                "\n\nfor an input of {}, expected used bytes to be {} or {}, but was {}",
                input,
                input,
                input - 1,
                used_bytes
            );
        }
    }
}