aws_smithy_legacy_http_server/
body.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6//! HTTP body utilities.
7
8// Used in the codegen in trait bounds.
9#[doc(hidden)]
10pub use http_body::Body as HttpBody;
11
12pub use hyper::body::Body;
13
14use bytes::Bytes;
15
16use crate::error::{BoxError, Error};
17
18/// The primary [`Body`] returned by the generated `smithy-rs` service.
19pub type BoxBody = http_body::combinators::UnsyncBoxBody<Bytes, Error>;
20
21// `boxed` is used in the codegen of the implementation of the operation `Handler` trait.
22/// Convert a [`http_body::Body`] into a [`BoxBody`].
23pub fn boxed<B>(body: B) -> BoxBody
24where
25    B: http_body::Body<Data = Bytes> + Send + 'static,
26    B::Error: Into<BoxError>,
27{
28    try_downcast(body).unwrap_or_else(|body| body.map_err(Error::new).boxed_unsync())
29}
30
31#[doc(hidden)]
32pub(crate) fn try_downcast<T, K>(k: K) -> Result<T, K>
33where
34    T: 'static,
35    K: Send + 'static,
36{
37    let mut k = Some(k);
38    if let Some(k) = <dyn std::any::Any>::downcast_mut::<Option<T>>(&mut k) {
39        Ok(k.take().unwrap())
40    } else {
41        Err(k.unwrap())
42    }
43}
44
45pub(crate) fn empty() -> BoxBody {
46    boxed(http_body::Empty::new())
47}
48
49/// Convert anything that can be converted into a [`hyper::body::Body`] into a [`BoxBody`].
50/// This simplifies codegen a little bit.
51#[doc(hidden)]
52pub fn to_boxed<B>(body: B) -> BoxBody
53where
54    Body: From<B>,
55{
56    boxed(Body::from(body))
57}
58
59// ============================================================================
60// Size-limited Body Collection
61// ============================================================================
62
63use std::fmt;
64use std::future::poll_fn;
65use std::pin::Pin;
66
67/// An error produced by [`collect_body_limited`] when the body exceeds the
68/// configured limit.
69#[doc(hidden)]
70#[derive(Debug, Clone, Copy)]
71pub struct BodyLimitExceeded {
72    /// The configured maximum, in bytes.
73    pub limit: usize,
74}
75
76impl fmt::Display for BodyLimitExceeded {
77    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
78        write!(
79            f,
80            "request body exceeded the configured maximum of {} bytes",
81            self.limit
82        )
83    }
84}
85
86impl std::error::Error for BodyLimitExceeded {}
87
88/// The error returned by [`collect_body_limited`].
89///
90/// Either the underlying body produced an error, or the configured size limit was
91/// exceeded. The generic over `E` lets the helper work with body error types that
92/// are `Send` but not `Sync` (the generated-server code carries only a `Send` bound).
93#[doc(hidden)]
94#[derive(Debug)]
95pub enum CollectBodyError<E> {
96    /// The underlying body produced an error while being read.
97    Body(E),
98    /// The body exceeded the configured maximum size.
99    TooLarge(BodyLimitExceeded),
100}
101
102impl<E: fmt::Display> fmt::Display for CollectBodyError<E> {
103    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
104        match self {
105            Self::Body(e) => write!(f, "error reading request body: {e}"),
106            Self::TooLarge(e) => e.fmt(f),
107        }
108    }
109}
110
111impl<E: std::error::Error + 'static> std::error::Error for CollectBodyError<E> {
112    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
113        match self {
114            Self::Body(e) => Some(e),
115            Self::TooLarge(e) => Some(e),
116        }
117    }
118}
119
120/// Collect an HTTP body into a [`Bytes`] buffer, enforcing a maximum size.
121///
122/// If the body produces more than `limit` bytes the function returns an error
123/// *before* buffering additional data. This bounds server memory regardless of
124/// what the client sends (for example, via `Transfer-Encoding: chunked`).
125///
126/// Passing `limit == 0` disables the check and collects the entire body (the
127/// historical behavior, *not* recommended — see the security notes on the
128/// `requestBodyMaxBytes` codegen setting).
129#[doc(hidden)]
130pub async fn collect_body_limited<B>(body: B, limit: usize) -> Result<Bytes, CollectBodyError<B::Error>>
131where
132    B: HttpBody,
133{
134    // `http-body` 0.4's core method `poll_data` requires a `Pin<&mut Self>`. Heap-pin
135    // so this works for any `B: HttpBody` without requiring `Unpin`.
136    let lower = body.size_hint().lower() as usize;
137    if lower > limit && limit > 0 {
138        return Err(CollectBodyError::TooLarge(BodyLimitExceeded { limit }));
139    }
140
141    let mut body: Pin<Box<B>> = Box::pin(body);
142    let mut buf: Vec<u8> = Vec::with_capacity(lower);
143
144    loop {
145        let chunk_opt = poll_fn(|cx| body.as_mut().poll_data(cx)).await;
146        match chunk_opt {
147            None => break,
148            Some(Err(e)) => return Err(CollectBodyError::Body(e)),
149            Some(Ok(mut chunk)) => {
150                use bytes::Buf;
151                let chunk_len = chunk.remaining();
152                if limit > 0 && buf.len().saturating_add(chunk_len) > limit {
153                    return Err(CollectBodyError::TooLarge(BodyLimitExceeded { limit }));
154                }
155                buf.reserve(chunk_len);
156                while chunk.has_remaining() {
157                    let slice = chunk.chunk();
158                    let len = slice.len();
159                    buf.extend_from_slice(slice);
160                    chunk.advance(len);
161                }
162            }
163        }
164    }
165
166    Ok(Bytes::from(buf))
167}
168
169#[cfg(test)]
170mod tests {
171    use super::*;
172
173    #[tokio::test]
174    async fn test_collect_body_limited_under_limit() {
175        let body = Body::from("hello");
176        let result = collect_body_limited(body, 1024).await;
177        assert_eq!(result.unwrap(), Bytes::from("hello"));
178    }
179
180    #[tokio::test]
181    async fn test_collect_body_limited_over_limit() {
182        let body = Body::from("this is way too long");
183        let result = collect_body_limited(body, 5).await;
184        assert!(matches!(
185            result,
186            Err(CollectBodyError::TooLarge(BodyLimitExceeded { limit: 5 }))
187        ));
188    }
189
190    #[tokio::test]
191    async fn test_collect_body_limited_chunked_exceeds_mid_stream() {
192        use futures_util::stream;
193
194        // Simulate a chunked body where the limit is crossed on the second chunk.
195        let chunks: Vec<Result<&[u8], std::io::Error>> = vec![
196            Ok(b"aaaa"), // 4 bytes
197            Ok(b"bbbb"), // +4 = 8, over limit of 6
198        ];
199        let body = Body::wrap_stream(stream::iter(chunks));
200        let result = collect_body_limited(body, 6).await;
201        assert!(matches!(result, Err(CollectBodyError::TooLarge(_))));
202    }
203
204    #[tokio::test]
205    async fn test_collect_body_limited_empty_body() {
206        let body = Body::empty();
207        let result = collect_body_limited(body, 100).await;
208        assert_eq!(result.unwrap(), Bytes::from(""));
209    }
210}
aws_smithy_legacy_http_server/body.rs

aws_smithy_legacy_http_server/
body.rs
