Skip to main content

aws_smithy_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//! This module provides body handling utilities for HTTP 1.x using the
9//! `http-body` and `http-body-util` crates.
10
11use crate::error::{BoxError, Error};
12use bytes::Bytes;
13use std::fmt;
14
15// Preserve the public bound used by previously generated servers.
16pub use http_body::Body as HttpBody;
17
18/// Internal helpers used by generated server code.
19///
20/// Public so generated code can access them. No public API stability or semver guarantees.
21#[doc(hidden)]
22pub mod __private;
23
24// ============================================================================
25// BoxBody - Type-Erased Body
26// ============================================================================
27
28/// The primary body type returned by the generated `smithy-rs` service.
29///
30/// This is a type-erased body that wraps `UnsyncBoxBody` from `http-body-util`.
31/// It is `Send` but not `Sync`, making it suitable for most HTTP handlers.
32pub type BoxBody = http_body_util::combinators::UnsyncBoxBody<Bytes, Error>;
33
34/// A thread-safe body type for operations that require `Sync`.
35///
36/// This is used specifically for event streaming operations and lambda handlers
37/// that need thread safety guarantees.
38pub type BoxBodySync = http_body_util::combinators::BoxBody<Bytes, Error>;
39
40/// A request body on the schema-routing pipeline.
41///
42/// A request entering through hyper keeps its [`Incoming`](hyper::body::Incoming) body
43/// inside the `Hyper` state, without boxing. The routing service replaces it with the
44/// `Buffered` state after collecting for a body-first protocol. Other compatible HTTP bodies
45/// enter through [`RequestBody::new`], which erases them into a boxed state.
46///
47/// [`RequestBody::prefixed`] replays bytes already read ahead of the remaining body. This allows
48/// middleware to inspect a prefix, then forward the full request without collecting the rest.
49///
50/// The states are private so buffering and transport-specific optimizations can evolve.
51///
52/// `RequestBody` is `Send` and `Sync`, so it can be handed to [`SdkBody::from_body_1_x`]
53/// when an operation input has a streaming member; `SdkBody` only accepts `Sync` bodies.
54///
55/// [`SdkBody::from_body_1_x`]: aws_smithy_types::body::SdkBody::from_body_1_x
56#[derive(Debug)]
57pub struct RequestBody(BodyInner);
58
59/// The default request body: the schema pipeline over hyper's transport body.
60pub type Body = RequestBody;
61
62#[derive(Debug)]
63enum BodyInner {
64    /// A hyper incoming body. May have been partially polled.
65    Hyper(hyper::body::Incoming),
66    /// An erased body from a non-transport source. May have been partially polled.
67    Boxed(BoxBodySync),
68    /// Buffered bytes, followed by an optional unread tail and any already-captured trailers.
69    Buffered {
70        bytes: Option<Bytes>,
71        trailers: Option<http::HeaderMap>,
72        // None means complete collection. A tail remains until polling it returns Ready(None).
73        tail: Option<Box<RequestBody>>,
74    },
75}
76
77impl RequestBody {
78    /// Wraps any compatible HTTP body without polling it.
79    ///
80    /// A hyper incoming body and an already-wrapped `RequestBody` stay unerased; any other
81    /// body type is erased into a boxed state.
82    pub fn new<T>(body: T) -> Self
83    where
84        T: http_body::Body<Data = Bytes> + Send + Sync + 'static,
85        T::Error: Into<BoxError>,
86    {
87        try_downcast::<Self, T>(body)
88            .or_else(|body| try_downcast::<hyper::body::Incoming, T>(body).map(|body| Self(BodyInner::Hyper(body))))
89            .unwrap_or_else(|body| Self(BodyInner::Boxed(boxed_sync(body))))
90    }
91
92    /// Builds an already-buffered body, preserving any trailers read with the content.
93    pub(crate) fn buffered(bytes: Bytes, trailers: Option<http::HeaderMap>) -> Self {
94        Self(BodyInner::Buffered {
95            bytes: Some(bytes),
96            trailers,
97            tail: None,
98        })
99    }
100
101    /// Replays bytes already read from a request ahead of its remaining body.
102    ///
103    /// Include all consumed data in `prefix` so downstream readers receive the full request.
104    /// Construction does not poll or erase `tail`. An empty prefix with no captured trailers
105    /// returns `tail` unchanged.
106    ///
107    /// Pass `None` for `trailers` when only data has been read: unread trailers arrive from
108    /// `tail` normally. If trailer frames were already read, pass them here to replay them
109    /// after `tail` reaches end-of-stream.
110    pub fn prefixed(prefix: Bytes, trailers: Option<http::HeaderMap>, tail: Self) -> Self {
111        if prefix.is_empty() && trailers.is_none() {
112            return tail;
113        }
114        Self(BodyInner::Buffered {
115            bytes: Some(prefix).filter(|bytes| !bytes.is_empty()),
116            trailers,
117            tail: Some(Box::new(tail)),
118        })
119    }
120
121    // Only unpolled, completely buffered bodies qualify for complete-content reuse.
122    pub(crate) fn buffered_content(&self) -> Option<&Bytes> {
123        match &self.0 {
124            BodyInner::Buffered { bytes, tail: None, .. } => bytes.as_ref(),
125            _ => None,
126        }
127    }
128
129    /// Creates an empty request body.
130    pub fn empty() -> Self {
131        Self(BodyInner::Buffered {
132            bytes: None,
133            trailers: None,
134            tail: None,
135        })
136    }
137
138    /// Creates a request body containing buffered bytes.
139    pub fn from_bytes(bytes: Bytes) -> Self {
140        Self::buffered(bytes, None)
141    }
142}
143
144impl Default for RequestBody {
145    fn default() -> Self {
146        Self::empty()
147    }
148}
149
150/// The transport body enters unerased, in the `Hyper` state.
151///
152/// Other compatible HTTP bodies can be wrapped with [`RequestBody::new`].
153impl From<hyper::body::Incoming> for Body {
154    fn from(body: hyper::body::Incoming) -> Self {
155        Self(BodyInner::Hyper(body))
156    }
157}
158
159impl From<Bytes> for RequestBody {
160    fn from(bytes: Bytes) -> Self {
161        Self::from_bytes(bytes)
162    }
163}
164
165impl From<Vec<u8>> for RequestBody {
166    fn from(bytes: Vec<u8>) -> Self {
167        Self::from_bytes(bytes.into())
168    }
169}
170
171impl From<&'static [u8]> for RequestBody {
172    fn from(bytes: &'static [u8]) -> Self {
173        Self::from_bytes(Bytes::from_static(bytes))
174    }
175}
176
177impl From<String> for RequestBody {
178    fn from(string: String) -> Self {
179        Self::from_bytes(string.into())
180    }
181}
182
183impl From<&'static str> for RequestBody {
184    fn from(string: &'static str) -> Self {
185        Self::from_bytes(Bytes::from_static(string.as_bytes()))
186    }
187}
188
189impl http_body::Body for RequestBody {
190    type Data = Bytes;
191    type Error = Error;
192
193    fn poll_frame(
194        mut self: std::pin::Pin<&mut Self>,
195        cx: &mut std::task::Context<'_>,
196    ) -> std::task::Poll<Option<Result<http_body::Frame<Bytes>, Error>>> {
197        match &mut self.0 {
198            BodyInner::Hyper(body) => std::pin::Pin::new(body)
199                .poll_frame(cx)
200                .map(|frame| frame.map(|result| result.map_err(Error::new))),
201            BodyInner::Boxed(body) => std::pin::Pin::new(body).poll_frame(cx),
202            BodyInner::Buffered { bytes, trailers, tail } => {
203                if let Some(bytes) = bytes.take().filter(|bytes| !bytes.is_empty()) {
204                    return std::task::Poll::Ready(Some(Ok(http_body::Frame::data(bytes))));
205                }
206                if let Some(remaining) = tail {
207                    match std::pin::Pin::new(remaining.as_mut()).poll_frame(cx) {
208                        std::task::Poll::Ready(None) => *tail = None,
209                        other => return other,
210                    }
211                }
212                std::task::Poll::Ready(trailers.take().map(http_body::Frame::trailers).map(Ok))
213            }
214        }
215    }
216
217    fn is_end_stream(&self) -> bool {
218        match &self.0 {
219            BodyInner::Hyper(body) => body.is_end_stream(),
220            BodyInner::Boxed(body) => body.is_end_stream(),
221            BodyInner::Buffered { bytes, trailers, tail } => {
222                bytes.as_ref().is_none_or(|bytes| bytes.is_empty())
223                    && tail.as_ref().is_none_or(|tail| tail.is_end_stream())
224                    && trailers.is_none()
225            }
226        }
227    }
228    fn size_hint(&self) -> http_body::SizeHint {
229        match &self.0 {
230            BodyInner::Hyper(body) => body.size_hint(),
231            BodyInner::Boxed(body) => body.size_hint(),
232            BodyInner::Buffered { bytes, tail, .. } => {
233                let prefix = bytes.as_ref().map_or(0, |bytes| bytes.len() as u64);
234                let remaining = tail
235                    .as_ref()
236                    .map_or_else(|| http_body::SizeHint::with_exact(0), |tail| tail.size_hint());
237                let mut hint = http_body::SizeHint::new();
238                hint.set_lower(prefix.saturating_add(remaining.lower()));
239                if let Some(upper) = remaining.upper().and_then(|upper| prefix.checked_add(upper)) {
240                    hint.set_upper(upper);
241                }
242                hint
243            }
244        }
245    }
246}
247
248// ============================================================================
249// Body Construction Functions
250// ============================================================================
251
252// `boxed` is used in the codegen of the implementation of the operation `Handler` trait.
253/// Convert an HTTP body implementing [`http_body::Body`] into a [`BoxBody`].
254pub fn boxed<B>(body: B) -> BoxBody
255where
256    B: http_body::Body<Data = Bytes> + Send + 'static,
257    B::Error: Into<BoxError>,
258{
259    use http_body_util::BodyExt;
260
261    try_downcast(body).unwrap_or_else(|body| body.map_err(Error::new).boxed_unsync())
262}
263
264/// Convert an HTTP body implementing [`http_body::Body`] into a [`BoxBodySync`].
265pub fn boxed_sync<B>(body: B) -> BoxBodySync
266where
267    B: http_body::Body<Data = Bytes> + Send + Sync + 'static,
268    B::Error: Into<BoxError>,
269{
270    use http_body_util::BodyExt;
271    body.map_err(Error::new).boxed()
272}
273
274/// Convert bytes or similar types into a [`BoxBody`].
275pub fn to_boxed<B>(body: B) -> BoxBody
276where
277    B: Into<Bytes>,
278{
279    boxed(http_body_util::Full::new(body.into()))
280}
281
282/// Convert bytes or similar types into a [`BoxBodySync`].
283pub fn to_boxed_sync<B>(body: B) -> BoxBodySync
284where
285    B: Into<Bytes>,
286{
287    boxed_sync(http_body_util::Full::new(body.into()))
288}
289
290fn try_downcast<T, K>(k: K) -> Result<T, K>
291where
292    T: 'static,
293    K: Send + 'static,
294{
295    let mut k = Some(k);
296    if let Some(k) = <dyn std::any::Any>::downcast_mut::<Option<T>>(&mut k) {
297        Ok(k.take().unwrap())
298    } else {
299        Err(k.unwrap())
300    }
301}
302
303/// Create an empty body.
304pub fn empty() -> BoxBody {
305    boxed(http_body_util::Empty::<Bytes>::new())
306}
307
308/// Create an empty sync body.
309pub fn empty_sync() -> BoxBodySync {
310    boxed_sync(http_body_util::Empty::<Bytes>::new())
311}
312
313/// Create a body from bytes.
314pub fn from_bytes(bytes: Bytes) -> BoxBody {
315    boxed(http_body_util::Full::new(bytes))
316}
317
318// ============================================================================
319// Size-limited Body Collection
320// ============================================================================
321
322/// An error produced by [`collect_body_limited`] when the body exceeds the configured limit.
323///
324/// This is intentionally a simple struct so it can be downcast via
325/// `Box<dyn std::error::Error>` from code that processes the error chain.
326#[doc(hidden)]
327#[derive(Debug, Clone, Copy)]
328pub struct BodyLimitExceeded {
329    /// The configured maximum, in bytes.
330    pub limit: usize,
331}
332
333impl fmt::Display for BodyLimitExceeded {
334    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
335        write!(
336            f,
337            "request body exceeded the configured maximum of {} bytes",
338            self.limit
339        )
340    }
341}
342
343impl std::error::Error for BodyLimitExceeded {}
344
345/// The error returned by [`collect_body_limited`].
346///
347/// Either the underlying body produced an error, or the configured size limit was
348/// exceeded. The generic over `E` lets the helper work with body error types that
349/// are `Send` but not `Sync` (the generated-server code carries only a `Send` bound),
350/// while still letting callers `?` the result into a `RequestRejection`.
351#[doc(hidden)]
352#[derive(Debug)]
353pub enum CollectBodyError<E> {
354    /// The underlying body produced an error while being read.
355    Body(E),
356    /// The body exceeded the configured maximum size.
357    TooLarge(BodyLimitExceeded),
358}
359
360impl<E: fmt::Display> fmt::Display for CollectBodyError<E> {
361    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
362        match self {
363            Self::Body(e) => write!(f, "error reading request body: {e}"),
364            Self::TooLarge(e) => e.fmt(f),
365        }
366    }
367}
368
369impl<E: std::error::Error + 'static> std::error::Error for CollectBodyError<E> {
370    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
371        match self {
372            Self::Body(e) => Some(e),
373            Self::TooLarge(e) => Some(e),
374        }
375    }
376}
377
378/// Collect an HTTP body into a [`Bytes`] buffer, enforcing a maximum size.
379///
380/// If the body produces more than `limit` bytes (including via
381/// `Transfer-Encoding: chunked` or a mismatched `Content-Length`), the function
382/// returns an error *before* buffering additional data.
383///
384/// Passing `limit == 0` disables the check and collects the entire body (the
385/// historical behavior, *not* recommended — see the security notes on the
386/// `requestBodyMaxBytes` codegen setting).
387#[doc(hidden)]
388pub async fn collect_body_limited<B>(body: B, limit: usize) -> Result<Bytes, CollectBodyError<B::Error>>
389where
390    B: HttpBody,
391{
392    use http_body_util::BodyExt;
393
394    if limit == 0 {
395        return body
396            .collect()
397            .await
398            .map(|c| c.to_bytes())
399            .map_err(CollectBodyError::Body);
400    }
401
402    // Walk frames ourselves (rather than using `http_body_util::Limited`) so we
403    // don't require `B::Error: Send + Sync + 'static`. The generated server
404    // deserializer only carries a `Send` bound on body errors.
405    let lower = body.size_hint().lower() as usize;
406    if lower > limit {
407        return Err(CollectBodyError::TooLarge(BodyLimitExceeded { limit }));
408    }
409
410    let mut body = std::pin::pin!(body);
411    let mut collected = bytes::BytesMut::with_capacity(lower);
412    while let Some(frame) = std::future::poll_fn(|cx| body.as_mut().poll_frame(cx))
413        .await
414        .transpose()
415        .map_err(CollectBodyError::Body)?
416    {
417        if let Ok(data) = frame.into_data() {
418            use bytes::{Buf, BufMut};
419            let data_len = data.remaining();
420            if collected.len().saturating_add(data_len) > limit {
421                return Err(CollectBodyError::TooLarge(BodyLimitExceeded { limit }));
422            }
423            collected.put(data);
424        }
425        // Trailer / non-data frames are discarded; we don't surface trailers.
426    }
427
428    Ok(collected.freeze())
429}
430
431// ============================================================================
432// Stream Wrapping for Event Streaming
433// ============================================================================
434
435/// Wrap a stream of byte chunks into a BoxBody.
436///
437/// This is used for event streaming support. The stream should produce `Result<O, E>`
438/// where `O` can be converted into `Bytes` and `E` can be converted into an error.
439///
440/// In hyper 0.x, `Body::wrap_stream` was available directly on the body type.
441/// In hyper 1.x, the `stream` feature was removed, and the official approach is to use
442/// `http_body_util::StreamBody` to convert streams into bodies, which is what this
443/// function provides as a convenient wrapper.
444///
445/// For scenarios requiring `Sync` (e.g., lambda handlers), use [`wrap_stream_sync`] instead.
446pub fn wrap_stream<S, O, E>(stream: S) -> BoxBody
447where
448    S: futures_util::Stream<Item = Result<O, E>> + Send + 'static,
449    O: Into<Bytes> + 'static,
450    E: Into<BoxError> + 'static,
451{
452    use futures_util::TryStreamExt;
453    use http_body_util::StreamBody;
454
455    // Convert the stream of Result<O, E> into a stream of Result<Frame<Bytes>, Error>
456    let frame_stream = stream
457        .map_ok(|chunk| http_body::Frame::data(chunk.into()))
458        .map_err(|e| Error::new(e.into()));
459
460    boxed(StreamBody::new(frame_stream))
461}
462
463/// Wrap a stream of byte chunks into a BoxBodySync.
464///
465/// This is the thread-safe variant of [`wrap_stream`], used for event streaming operations
466/// that require `Sync` bounds, such as lambda handlers.
467///
468/// The stream should produce `Result<O, E>` where `O` can be converted into `Bytes` and
469/// `E` can be converted into an error.
470pub fn wrap_stream_sync<S, O, E>(stream: S) -> BoxBodySync
471where
472    S: futures_util::Stream<Item = Result<O, E>> + Send + Sync + 'static,
473    O: Into<Bytes> + 'static,
474    E: Into<BoxError> + 'static,
475{
476    use futures_util::TryStreamExt;
477    use http_body_util::StreamBody;
478
479    // Convert the stream of Result<O, E> into a stream of Result<Frame<Bytes>, Error>
480    let frame_stream = stream
481        .map_ok(|chunk| http_body::Frame::data(chunk.into()))
482        .map_err(|e| Error::new(e.into()));
483
484    boxed_sync(StreamBody::new(frame_stream))
485}
486
487#[cfg(test)]
488mod tests;