aws_smithy_http_server/

body.rs

/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0
 */

//! HTTP body utilities.
//!
//! This module provides body handling utilities for HTTP 1.x using the
//! `http-body` and `http-body-util` crates.

use crate::error::{BoxError, Error};
use bytes::Bytes;
use std::fmt;

// Used in the codegen in trait bounds.
#[doc(hidden)]
pub use http_body::Body as HttpBody;

// ============================================================================
// BoxBody - Type-Erased Body
// ============================================================================

/// The primary body type returned by the generated `smithy-rs` service.
///
/// This is a type-erased body that wraps `UnsyncBoxBody` from `http-body-util`.
/// It is `Send` but not `Sync`, making it suitable for most HTTP handlers.
pub type BoxBody = http_body_util::combinators::UnsyncBoxBody<Bytes, Error>;

/// A thread-safe body type for operations that require `Sync`.
///
/// This is used specifically for event streaming operations and lambda handlers
/// that need thread safety guarantees.
pub type BoxBodySync = http_body_util::combinators::BoxBody<Bytes, Error>;

// ============================================================================
// Body Construction Functions
// ============================================================================

// `boxed` is used in the codegen of the implementation of the operation `Handler` trait.
/// Convert an HTTP body implementing [`http_body::Body`] into a [`BoxBody`].
pub fn boxed<B>(body: B) -> BoxBody
where
    B: http_body::Body<Data = Bytes> + Send + 'static,
    B::Error: Into<BoxError>,
{
    use http_body_util::BodyExt;

    try_downcast(body).unwrap_or_else(|body| body.map_err(Error::new).boxed_unsync())
}

/// Convert an HTTP body implementing [`http_body::Body`] into a [`BoxBodySync`].
pub fn boxed_sync<B>(body: B) -> BoxBodySync
where
    B: http_body::Body<Data = Bytes> + Send + Sync + 'static,
    B::Error: Into<BoxError>,
{
    use http_body_util::BodyExt;
    body.map_err(Error::new).boxed()
}

#[doc(hidden)]
pub(crate) fn try_downcast<T, K>(k: K) -> Result<T, K>
where
    T: 'static,
    K: Send + 'static,
{
    let mut k = Some(k);
    if let Some(k) = <dyn std::any::Any>::downcast_mut::<Option<T>>(&mut k) {
        Ok(k.take().unwrap())
    } else {
        Err(k.unwrap())
    }
}

/// Create an empty body.
pub fn empty() -> BoxBody {
    boxed(http_body_util::Empty::<Bytes>::new())
}

/// Create an empty sync body.
pub fn empty_sync() -> BoxBodySync {
    boxed_sync(http_body_util::Empty::<Bytes>::new())
}

/// Convert bytes or similar types into a [`BoxBody`].
///
/// This simplifies codegen a little bit.
#[doc(hidden)]
pub fn to_boxed<B>(body: B) -> BoxBody
where
    B: Into<Bytes>,
{
    boxed(http_body_util::Full::new(body.into()))
}

/// Convert bytes or similar types into a [`BoxBodySync`].
///
/// This simplifies codegen a little bit.
#[doc(hidden)]
pub fn to_boxed_sync<B>(body: B) -> BoxBodySync
where
    B: Into<Bytes>,
{
    boxed_sync(http_body_util::Full::new(body.into()))
}

/// Create a body from bytes.
pub fn from_bytes(bytes: Bytes) -> BoxBody {
    boxed(http_body_util::Full::new(bytes))
}

// ============================================================================
// Size-limited Body Collection
// ============================================================================

/// An error produced by [`collect_body_limited`] when the body exceeds the configured limit.
///
/// This is intentionally a simple struct so it can be downcast via
/// `Box<dyn std::error::Error>` from code that processes the error chain.
#[doc(hidden)]
#[derive(Debug, Clone, Copy)]
pub struct BodyLimitExceeded {
    /// The configured maximum, in bytes.
    pub limit: usize,
}

impl fmt::Display for BodyLimitExceeded {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "request body exceeded the configured maximum of {} bytes",
            self.limit
        )
    }
}

impl std::error::Error for BodyLimitExceeded {}

/// The error returned by [`collect_body_limited`].
///
/// Either the underlying body produced an error, or the configured size limit was
/// exceeded. The generic over `E` lets the helper work with body error types that
/// are `Send` but not `Sync` (the generated-server code carries only a `Send` bound),
/// while still letting callers `?` the result into a `RequestRejection`.
#[doc(hidden)]
#[derive(Debug)]
pub enum CollectBodyError<E> {
    /// The underlying body produced an error while being read.
    Body(E),
    /// The body exceeded the configured maximum size.
    TooLarge(BodyLimitExceeded),
}

impl<E: fmt::Display> fmt::Display for CollectBodyError<E> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Body(e) => write!(f, "error reading request body: {e}"),
            Self::TooLarge(e) => e.fmt(f),
        }
    }
}

impl<E: std::error::Error + 'static> std::error::Error for CollectBodyError<E> {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::Body(e) => Some(e),
            Self::TooLarge(e) => Some(e),
        }
    }
}

/// Collect an HTTP body into a [`Bytes`] buffer, enforcing a maximum size.
///
/// If the body produces more than `limit` bytes (including via
/// `Transfer-Encoding: chunked` or a mismatched `Content-Length`), the function
/// returns an error *before* buffering additional data.
///
/// Passing `limit == 0` disables the check and collects the entire body (the
/// historical behavior, *not* recommended — see the security notes on the
/// `requestBodyMaxBytes` codegen setting).
#[doc(hidden)]
pub async fn collect_body_limited<B>(body: B, limit: usize) -> Result<Bytes, CollectBodyError<B::Error>>
where
    B: HttpBody,
{
    use http_body_util::BodyExt;

    if limit == 0 {
        return body
            .collect()
            .await
            .map(|c| c.to_bytes())
            .map_err(CollectBodyError::Body);
    }

    // Walk frames ourselves (rather than using `http_body_util::Limited`) so we
    // don't require `B::Error: Send + Sync + 'static`. The generated server
    // deserializer only carries a `Send` bound on body errors.
    let lower = body.size_hint().lower() as usize;
    if lower > limit {
        return Err(CollectBodyError::TooLarge(BodyLimitExceeded { limit }));
    }

    let mut body = std::pin::pin!(body);
    let mut collected = bytes::BytesMut::with_capacity(lower);
    while let Some(frame) = std::future::poll_fn(|cx| body.as_mut().poll_frame(cx))
        .await
        .transpose()
        .map_err(CollectBodyError::Body)?
    {
        if let Ok(data) = frame.into_data() {
            use bytes::{Buf, BufMut};
            let data_len = data.remaining();
            if collected.len().saturating_add(data_len) > limit {
                return Err(CollectBodyError::TooLarge(BodyLimitExceeded { limit }));
            }
            collected.put(data);
        }
        // Trailer / non-data frames are discarded; we don't surface trailers.
    }

    Ok(collected.freeze())
}

// ============================================================================
// Stream Wrapping for Event Streaming
// ============================================================================

/// Wrap a stream of byte chunks into a BoxBody.
///
/// This is used for event streaming support. The stream should produce `Result<O, E>`
/// where `O` can be converted into `Bytes` and `E` can be converted into an error.
///
/// In hyper 0.x, `Body::wrap_stream` was available directly on the body type.
/// In hyper 1.x, the `stream` feature was removed, and the official approach is to use
/// `http_body_util::StreamBody` to convert streams into bodies, which is what this
/// function provides as a convenient wrapper.
///
/// For scenarios requiring `Sync` (e.g., lambda handlers), use [`wrap_stream_sync`] instead.
pub fn wrap_stream<S, O, E>(stream: S) -> BoxBody
where
    S: futures_util::Stream<Item = Result<O, E>> + Send + 'static,
    O: Into<Bytes> + 'static,
    E: Into<BoxError> + 'static,
{
    use futures_util::TryStreamExt;
    use http_body_util::StreamBody;

    // Convert the stream of Result<O, E> into a stream of Result<Frame<Bytes>, Error>
    let frame_stream = stream
        .map_ok(|chunk| http_body::Frame::data(chunk.into()))
        .map_err(|e| Error::new(e.into()));

    boxed(StreamBody::new(frame_stream))
}

/// Wrap a stream of byte chunks into a BoxBodySync.
///
/// This is the thread-safe variant of [`wrap_stream`], used for event streaming operations
/// that require `Sync` bounds, such as lambda handlers.
///
/// The stream should produce `Result<O, E>` where `O` can be converted into `Bytes` and
/// `E` can be converted into an error.
pub fn wrap_stream_sync<S, O, E>(stream: S) -> BoxBodySync
where
    S: futures_util::Stream<Item = Result<O, E>> + Send + Sync + 'static,
    O: Into<Bytes> + 'static,
    E: Into<BoxError> + 'static,
{
    use futures_util::TryStreamExt;
    use http_body_util::StreamBody;

    // Convert the stream of Result<O, E> into a stream of Result<Frame<Bytes>, Error>
    let frame_stream = stream
        .map_ok(|chunk| http_body::Frame::data(chunk.into()))
        .map_err(|e| Error::new(e.into()));

    boxed_sync(StreamBody::new(frame_stream))
}

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

    /// Collect all bytes from a body (test utility).
    ///
    /// This uses `http_body_util::BodyExt::collect()` to read all body chunks
    /// into a single `Bytes` buffer.
    async fn collect_bytes<B>(body: B) -> Result<Bytes, Error>
    where
        B: HttpBody,
        B::Error: Into<BoxError>,
    {
        use http_body_util::BodyExt;

        let collected = body.collect().await.map_err(Error::new)?;
        Ok(collected.to_bytes())
    }

    #[tokio::test]
    async fn test_empty_body() {
        let body = empty();
        let bytes = collect_bytes(body).await.unwrap();
        assert_eq!(bytes.len(), 0);
    }

    #[tokio::test]
    async fn test_from_bytes() {
        let data = Bytes::from("hello world");
        let body = from_bytes(data.clone());
        let collected = collect_bytes(body).await.unwrap();
        assert_eq!(collected, data);
    }

    #[tokio::test]
    async fn test_to_boxed_string() {
        let s = "hello world";
        let body = to_boxed(s);
        let collected = collect_bytes(body).await.unwrap();
        assert_eq!(collected, Bytes::from(s));
    }

    #[tokio::test]
    async fn test_to_boxed_vec() {
        let vec = vec![1u8, 2, 3, 4, 5];
        let body = to_boxed(vec.clone());
        let collected = collect_bytes(body).await.unwrap();
        assert_eq!(collected.as_ref(), vec.as_slice());
    }

    #[tokio::test]
    async fn test_boxed() {
        use http_body_util::Full;
        let full_body = Full::new(Bytes::from("test data"));
        let boxed_body: BoxBody = boxed(full_body);
        let collected = collect_bytes(boxed_body).await.unwrap();
        assert_eq!(collected, Bytes::from("test data"));
    }

    #[tokio::test]
    async fn test_boxed_sync() {
        use http_body_util::Full;
        let full_body = Full::new(Bytes::from("sync test"));
        let boxed_body: BoxBodySync = boxed_sync(full_body);
        let collected = collect_bytes(boxed_body).await.unwrap();
        assert_eq!(collected, Bytes::from("sync test"));
    }

    #[tokio::test]
    async fn test_wrap_stream_single_chunk() {
        use futures_util::stream;

        let data = Bytes::from("single chunk");
        let stream = stream::iter(vec![Ok::<_, std::io::Error>(data.clone())]);

        let body = wrap_stream(stream);
        let collected = collect_bytes(body).await.unwrap();
        assert_eq!(collected, data);
    }

    #[tokio::test]
    async fn test_wrap_stream_multiple_chunks() {
        use futures_util::stream;

        let chunks = vec![
            Ok::<_, std::io::Error>(Bytes::from("chunk1")),
            Ok(Bytes::from("chunk2")),
            Ok(Bytes::from("chunk3")),
        ];
        let expected = Bytes::from("chunk1chunk2chunk3");

        let stream = stream::iter(chunks);
        let body = wrap_stream(stream);
        let collected = collect_bytes(body).await.unwrap();
        assert_eq!(collected, expected);
    }

    #[tokio::test]
    async fn test_wrap_stream_empty() {
        use futures_util::stream;

        let stream = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::new())]);

        let body = wrap_stream(stream);
        let collected = collect_bytes(body).await.unwrap();
        assert_eq!(collected.len(), 0);
    }

    #[tokio::test]
    async fn test_wrap_stream_error() {
        use futures_util::stream;

        let chunks = vec![
            Ok::<_, std::io::Error>(Bytes::from("chunk1")),
            Err(std::io::Error::other("test error")),
        ];

        let stream = stream::iter(chunks);
        let body = wrap_stream(stream);
        let result = collect_bytes(body).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_wrap_stream_various_types() {
        use futures_util::stream;

        // Test that Into<Bytes> works for various types

        // Test with &str
        let chunks = vec![Ok::<_, std::io::Error>("string slice"), Ok("another string")];
        let stream = stream::iter(chunks);
        let body = wrap_stream(stream);
        let collected = collect_bytes(body).await.unwrap();
        assert_eq!(collected, Bytes::from("string sliceanother string"));

        // Test with String
        let chunks = vec![
            Ok::<_, std::io::Error>(String::from("owned ")),
            Ok(String::from("strings")),
        ];
        let stream = stream::iter(chunks);
        let body = wrap_stream(stream);
        let collected = collect_bytes(body).await.unwrap();
        assert_eq!(collected, Bytes::from("owned strings"));

        // Test with Vec<u8>
        let chunks = vec![
            Ok::<_, std::io::Error>(vec![72u8, 101, 108, 108, 111]), // "Hello"
            Ok(vec![32u8, 87, 111, 114, 108, 100]),                  // " World"
        ];
        let stream = stream::iter(chunks);
        let body = wrap_stream(stream);
        let collected = collect_bytes(body).await.unwrap();
        assert_eq!(collected, Bytes::from("Hello World"));

        // Test with &[u8]
        let chunks = vec![
            Ok::<_, std::io::Error>(&[98u8, 121, 116, 101] as &[u8]), // "byte"
            Ok(&[115u8, 33] as &[u8]),                                // "s!"
        ];
        let stream = stream::iter(chunks);
        let body = wrap_stream(stream);
        let collected = collect_bytes(body).await.unwrap();
        assert_eq!(collected, Bytes::from("bytes!"));

        // Test with custom struct implementing Into<Bytes>
        struct CustomChunk {
            data: String,
        }

        impl From<CustomChunk> for Bytes {
            fn from(chunk: CustomChunk) -> Bytes {
                Bytes::from(chunk.data)
            }
        }

        let chunks = vec![
            Ok::<_, std::io::Error>(CustomChunk { data: "custom ".into() }),
            Ok(CustomChunk { data: "struct".into() }),
        ];
        let stream = stream::iter(chunks);
        let body = wrap_stream(stream);
        let collected = collect_bytes(body).await.unwrap();
        assert_eq!(collected, Bytes::from("custom struct"));
    }

    #[tokio::test]
    async fn test_wrap_stream_custom_stream_type() {
        use bytes::Bytes;
        use std::pin::Pin;
        use std::task::{Context, Poll};

        // Custom stream type that implements futures_util::Stream
        struct CustomStream {
            chunks: Vec<Result<Bytes, std::io::Error>>,
        }

        impl CustomStream {
            fn new(chunks: Vec<Result<Bytes, std::io::Error>>) -> Self {
                Self { chunks }
            }
        }

        impl futures_util::Stream for CustomStream {
            type Item = Result<Bytes, std::io::Error>;

            fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
                if self.chunks.is_empty() {
                    Poll::Ready(None)
                } else {
                    Poll::Ready(Some(self.chunks.remove(0)))
                }
            }
        }

        let stream = CustomStream::new(vec![Ok(Bytes::from("custom ")), Ok(Bytes::from("stream"))]);

        let body = wrap_stream(stream);
        let collected = collect_bytes(body).await.unwrap();
        assert_eq!(collected, Bytes::from("custom stream"));
    }

    #[tokio::test]
    async fn test_wrap_stream_custom_error_type() {
        use bytes::Bytes;
        use futures_util::stream;

        // Custom error type that implements Into<BoxError>
        #[derive(Debug, Clone)]
        struct CustomError {
            message: String,
        }

        impl std::fmt::Display for CustomError {
            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                write!(f, "CustomError: {}", self.message)
            }
        }

        impl std::error::Error for CustomError {}

        // Test successful case with custom error type
        let chunks = vec![
            Ok::<_, CustomError>(Bytes::from("custom ")),
            Ok(Bytes::from("error type")),
        ];
        let stream = stream::iter(chunks);
        let body = wrap_stream(stream);
        let collected = collect_bytes(body).await.unwrap();
        assert_eq!(collected, Bytes::from("custom error type"));

        // Test error case with custom error type
        let chunks = vec![
            Ok::<_, CustomError>(Bytes::from("data")),
            Err(CustomError {
                message: "custom error".into(),
            }),
        ];
        let stream = stream::iter(chunks);
        let body = wrap_stream(stream);
        let result = collect_bytes(body).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_wrap_stream_incremental_consumption() {
        use bytes::Bytes;
        use http_body_util::BodyExt;
        use std::pin::Pin;
        use std::task::{Context, Poll};

        struct IncrementalStream {
            chunks: Vec<Result<Bytes, std::io::Error>>,
        }

        impl IncrementalStream {
            fn new(chunks: Vec<Result<Bytes, std::io::Error>>) -> Self {
                Self { chunks }
            }
        }

        impl futures_util::Stream for IncrementalStream {
            type Item = Result<Bytes, std::io::Error>;

            fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
                if self.chunks.is_empty() {
                    Poll::Ready(None)
                } else {
                    Poll::Ready(Some(self.chunks.remove(0)))
                }
            }
        }

        let stream = IncrementalStream::new(vec![
            Ok(Bytes::from("chunk1")),
            Ok(Bytes::from("chunk2")),
            Ok(Bytes::from("chunk3")),
        ]);

        let mut body = wrap_stream(stream);

        let frame1 = body.frame().await.unwrap().unwrap();
        assert!(frame1.is_data());
        assert_eq!(frame1.into_data().unwrap(), Bytes::from("chunk1"));

        let frame2 = body.frame().await.unwrap().unwrap();
        assert!(frame2.is_data());
        assert_eq!(frame2.into_data().unwrap(), Bytes::from("chunk2"));

        let frame3 = body.frame().await.unwrap().unwrap();
        assert!(frame3.is_data());
        assert_eq!(frame3.into_data().unwrap(), Bytes::from("chunk3"));

        let frame4 = body.frame().await;
        assert!(frame4.is_none());
    }

    #[tokio::test]
    async fn test_wrap_stream_sync_single_chunk() {
        use futures_util::stream;

        let data = Bytes::from("sync single chunk");
        let stream = stream::iter(vec![Ok::<_, std::io::Error>(data.clone())]);

        let body = wrap_stream_sync(stream);
        let collected = collect_bytes(body).await.unwrap();
        assert_eq!(collected, data);
    }

    #[tokio::test]
    async fn test_wrap_stream_sync_multiple_chunks() {
        use futures_util::stream;

        let chunks = vec![
            Ok::<_, std::io::Error>(Bytes::from("sync1")),
            Ok(Bytes::from("sync2")),
            Ok(Bytes::from("sync3")),
        ];
        let expected = Bytes::from("sync1sync2sync3");

        let stream = stream::iter(chunks);
        let body = wrap_stream_sync(stream);
        let collected = collect_bytes(body).await.unwrap();
        assert_eq!(collected, expected);
    }

    #[tokio::test]
    async fn test_empty_sync_body() {
        let body = empty_sync();
        let bytes = collect_bytes(body).await.unwrap();
        assert_eq!(bytes.len(), 0);
    }

    #[tokio::test]
    async fn test_to_boxed_sync() {
        let data = Bytes::from("sync boxed data");
        let body = to_boxed_sync(data.clone());
        let collected = collect_bytes(body).await.unwrap();
        assert_eq!(collected, data);
    }

    // Compile-time tests to ensure Send/Sync bounds are correct
    // Following the pattern used by hyper and axum
    fn _assert_send<T: Send>() {}
    fn _assert_sync<T: Sync>() {}

    fn _assert_send_sync_bounds() {
        // BoxBodySync must be both Send and Sync
        _assert_send::<BoxBodySync>();
        _assert_sync::<BoxBodySync>();

        // BoxBody must be Send (but is intentionally NOT Sync - it's UnsyncBoxBody)
        _assert_send::<BoxBody>();
    }

    #[tokio::test]
    async fn test_wrap_stream_sync_produces_sync_body() {
        use futures_util::stream;

        let data = Bytes::from("test sync");
        let stream = stream::iter(vec![Ok::<_, std::io::Error>(data.clone())]);

        let body = wrap_stream_sync(stream);

        // Compile-time check: ensure the body is Sync
        fn check_sync<T: Sync>(_: &T) {}
        check_sync(&body);

        let collected = collect_bytes(body).await.unwrap();
        assert_eq!(collected, data);
    }

    #[test]
    fn test_empty_sync_is_sync() {
        let body = empty_sync();
        fn check_sync<T: Sync>(_: &T) {}
        check_sync(&body);
    }

    #[test]
    fn test_boxed_sync_is_sync() {
        use http_body_util::Full;
        let body = boxed_sync(Full::new(Bytes::from("test")));
        fn check_sync<T: Sync>(_: &T) {}
        check_sync(&body);
    }

    #[tokio::test]
    async fn test_collect_body_limited_under_limit() {
        let body = from_bytes(Bytes::from("hello"));
        let result = collect_body_limited(body, 1024).await;
        assert_eq!(result.unwrap(), Bytes::from("hello"));
    }

    #[tokio::test]
    async fn test_collect_body_limited_over_limit() {
        let body = from_bytes(Bytes::from("this is way too long"));
        let result = collect_body_limited(body, 5).await;
        assert!(matches!(
            result,
            Err(CollectBodyError::TooLarge(BodyLimitExceeded { limit: 5 }))
        ));
    }

    #[tokio::test]
    async fn test_collect_body_limited_chunked_exceeds_mid_stream() {
        use futures_util::stream;

        // Simulate a chunked body where the limit is crossed on the second chunk.
        let chunks = vec![
            Ok::<_, std::io::Error>(http_body::Frame::data(Bytes::from("aaaa"))), // 4 bytes
            Ok(http_body::Frame::data(Bytes::from("bbbb"))),                      // +4 = 8, over limit of 6
        ];
        let body = http_body_util::StreamBody::new(stream::iter(chunks));
        let result = collect_body_limited(body, 6).await;
        assert!(matches!(result, Err(CollectBodyError::TooLarge(_))));
    }

    #[tokio::test]
    async fn test_collect_body_limited_empty_body() {
        let body = empty();
        let result = collect_body_limited(body, 100).await;
        assert_eq!(result.unwrap(), Bytes::from(""));
    }
}