rama_http_types/body/
mod.rs

1//! Asynchronous HTTP request or response body.
2//!
3//! See [`StreamingBody`] for more details.
4
5use crate::body::util::BodyExt;
6use pin_project_lite::pin_project;
7use rama_core::bytes::Bytes;
8use rama_core::error::BoxError;
9use rama_core::futures::TryStream;
10use rama_core::futures::stream::Stream;
11use rama_core::stream::json;
12use rama_utils::str::arcstr::ArcStr;
13use serde::de::DeserializeOwned;
14use sse::{EventDataRead, EventStream};
15use std::pin::Pin;
16use std::task::{Context, Poll};
17use sync_wrapper::SyncWrapper;
18
19// Forked `http-body` / `http-body-util` (see fork README), relocated under
20// `body` so the module path matches the file layout.
21pub mod http_body;
22mod http_body_util;
23
24pub use http_body::Body as StreamingBody;
25pub use http_body::Frame;
26pub use http_body::SizeHint;
27
28// `http_body_util` is private; its API is surfaced through this curated alias.
29pub mod util {
30    pub use super::http_body_util::*;
31}
32
33mod limit;
34pub use limit::BodyLimit;
35
36mod ext;
37pub use ext::BodyExtractExt;
38
39#[doc(inline)]
40pub use util::{CollectError, CollectErrorKind, CollectOptions};
41
42pub mod sse;
43
44mod infinite;
45pub use infinite::InfiniteReader;
46
47mod on_drop;
48pub use on_drop::OnDropBody;
49
50mod guarded;
51pub use guarded::GuardedBody;
52
53mod optional;
54pub use optional::OptionalBody;
55
56mod capture;
57pub use capture::{
58    BodyCaptureEvent, BodyCaptureSink, BufferedBodyCapture, CaptureBody, CaptureCanceled,
59    CaptureHandle, CaptureLimit, CaptureOutcome, CapturedBody,
60};
61
62// Implementations copied over from http-body but addapted to work with our Requests/Response types
63
64impl<B: StreamingBody> StreamingBody for crate::Request<B> {
65    type Data = B::Data;
66    type Error = B::Error;
67
68    fn poll_frame(
69        self: Pin<&mut Self>,
70        cx: &mut Context<'_>,
71    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
72        // SAFETY:
73        // A pin projection.
74        unsafe { self.map_unchecked_mut(Self::body_mut).poll_frame(cx) }
75    }
76
77    fn is_end_stream(&self) -> bool {
78        self.body().is_end_stream()
79    }
80
81    fn size_hint(&self) -> SizeHint {
82        self.body().size_hint()
83    }
84}
85
86impl<B: StreamingBody> StreamingBody for crate::Response<B> {
87    type Data = B::Data;
88    type Error = B::Error;
89
90    fn poll_frame(
91        self: Pin<&mut Self>,
92        cx: &mut Context<'_>,
93    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
94        // SAFETY:
95        // A pin projection.
96        unsafe { self.map_unchecked_mut(Self::body_mut).poll_frame(cx) }
97    }
98
99    fn is_end_stream(&self) -> bool {
100        self.body().is_end_stream()
101    }
102
103    fn size_hint(&self) -> SizeHint {
104        self.body().size_hint()
105    }
106}
107
108// Rama specific stuff
109
110type BoxBody = util::combinators::BoxBody<Bytes, BoxError>;
111
112fn boxed<B>(body: B) -> BoxBody
113where
114    B: StreamingBody<Data = Bytes, Error: Into<BoxError>> + Send + Sync + 'static,
115{
116    try_downcast(body).unwrap_or_else(|body| body.map_err(Into::into).boxed())
117}
118
119#[expect(clippy::unwrap_used)]
120pub(crate) fn try_downcast<T, K>(k: K) -> Result<T, K>
121where
122    T: 'static,
123    K: Send + 'static,
124{
125    let mut k = Some(k);
126    if let Some(k) = <dyn std::any::Any>::downcast_mut::<Option<T>>(&mut k) {
127        Ok(k.take().unwrap())
128    } else {
129        Err(k.unwrap())
130    }
131}
132
133/// The body type used in rama requests and responses.
134#[must_use]
135#[derive(Debug)]
136pub struct Body(BoxBody);
137
138impl Body {
139    /// Create a new `Body` that wraps another [`Body`].
140    pub fn new<B>(body: B) -> Self
141    where
142        B: StreamingBody<Data = Bytes, Error: Into<BoxError>> + Send + Sync + 'static,
143    {
144        try_downcast(body).unwrap_or_else(|body| Self(boxed(body)))
145    }
146
147    /// Create a new `Body` with a maximum size limit.
148    pub fn with_limit<B>(body: B, limit: usize) -> Self
149    where
150        B: StreamingBody<Data = Bytes, Error: Into<BoxError>> + Send + Sync + 'static,
151    {
152        Self::new(util::Limited::new(body, limit))
153    }
154
155    /// Create an empty body.
156    pub fn empty() -> Self {
157        Self::new(util::Empty::new())
158    }
159
160    /// Create a new `Body` from a [`Stream`].
161    ///
162    /// [`Stream`]: https://docs.rs/futures/latest/futures/stream/trait.Stream.html
163    pub fn from_stream<S>(stream: S) -> Self
164    where
165        S: TryStream<Ok: Into<Bytes>, Error: Into<BoxError>> + Send + 'static,
166    {
167        Self::new(StreamBody {
168            stream: SyncWrapper::new(stream),
169        })
170    }
171
172    /// Create a new [`Body`] from a [`Stream`] with a maximum size limit.
173    pub fn limited(self, limit: usize) -> Self {
174        Self::new(util::Limited::new(self.0, limit))
175    }
176
177    /// Wrap this body so that `on_drop` is called if the body is dropped before
178    /// being fully consumed (e.g. because the client disconnected mid-response).
179    ///
180    /// The closure fires exactly once and only on early drops — it is disarmed
181    /// when [`StreamingBody::poll_frame`] returns `Poll::Ready(None)`.
182    pub fn on_drop<F>(self, on_drop: F) -> Self
183    where
184        F: FnOnce() + Send + Sync + 'static,
185    {
186        Self::new(OnDropBody::new(self.0, on_drop))
187    }
188
189    /// Forward this body while asynchronously sending owned frame copies to a sink.
190    pub fn capture<S>(self, sink: S) -> Self
191    where
192        S: BodyCaptureSink,
193    {
194        Self::new(CaptureBody::new(self.0, sink))
195    }
196
197    /// Forward this body while retaining a bounded or unlimited copy in memory.
198    ///
199    /// The returned handle resolves exactly once after the forwarding body
200    /// completes, errors, or is dropped:
201    ///
202    /// ```
203    /// use rama_http_types::{Body, CaptureLimit, CaptureOutcome, body::util::BodyExt as _};
204    ///
205    /// # #[tokio::main]
206    /// # async fn main() {
207    /// let (body, capture) = Body::from("hello").capture_buffered(CaptureLimit::max_bytes(1024));
208    /// body.collect().await.unwrap();
209    ///
210    /// let captured = capture.wait().await.unwrap();
211    /// assert_eq!(captured.bytes(), "hello");
212    /// assert_eq!(captured.outcome(), CaptureOutcome::Complete);
213    /// assert!(!captured.is_truncated());
214    /// # }
215    /// ```
216    pub fn capture_buffered(self, limit: CaptureLimit) -> (Self, CaptureHandle) {
217        let (body, handle) = CaptureBody::buffered(self.0, limit);
218        (Self::new(body), handle)
219    }
220
221    /// Attach a headermap as trailer headers to this body.
222    ///
223    /// In case you want this is also available as `with_trailer`,
224    /// provided by a body extension which allows you to use
225    /// a [`Future`] with as output an option of a result of a headermap.
226    /// This method is a shortcut of that function.
227    pub fn with_trailer_headers(self, headers: crate::HeaderMap) -> Self {
228        Self::new(self.0.with_trailers(std::future::ready(Some(Ok(headers)))))
229    }
230
231    /// Convert the body into a [`Stream`] of data frames.
232    ///
233    /// Non-data frames (such as trailers) will be discarded. Use [`http_body_util::BodyStream`] if
234    /// you need a [`Stream`] of all frame types.
235    ///
236    /// [`http_body_util::BodyStream`]: https://docs.rs/http-body-util/latest/http_body_util/struct.BodyStream.html
237    pub fn into_data_stream(self) -> BodyDataStream {
238        BodyDataStream { inner: self }
239    }
240
241    /// Convert the body into a [`Stream`] of [`sse::Event`]s.
242    ///
243    /// <https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events>.
244    #[must_use]
245    pub fn into_event_stream<T: EventDataRead>(self) -> EventStream<BodyDataStream, T> {
246        EventStream::new(self.into_data_stream())
247    }
248
249    /// Convert the body into a [`JsonReadStream`][json::JsonReadStream].
250    ///
251    /// Stream of json objects, each object separated by a newline (`\n`).
252    #[must_use]
253    pub fn into_json_stream<T: DeserializeOwned>(self) -> json::JsonReadStream<T, BodyDataStream> {
254        let stream = self.into_data_stream();
255        json::JsonReadStream::new(stream)
256    }
257
258    /// Convert the body into a [`JsonReadStream`][json::JsonReadStream].
259    ///
260    /// Stream of json objects, each object separated by a newline (`\n`).
261    #[must_use]
262    pub fn into_json_stream_with_config<T: DeserializeOwned>(
263        self,
264        cfg: json::ParseConfig,
265    ) -> json::JsonReadStream<T, BodyDataStream> {
266        let stream = self.into_data_stream();
267        json::JsonReadStream::new_with_config(stream, cfg)
268    }
269
270    /// Convert the body into a [`Stream`] of [`sse::Event`]s with optional string data.
271    ///
272    /// <https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events>.
273    #[must_use]
274    pub fn into_string_data_event_stream(self) -> EventStream<BodyDataStream> {
275        EventStream::new(self.into_data_stream())
276    }
277
278    /// Stream a chunk of the response body.
279    ///
280    /// When the response body has been exhausted, this will return `None`.
281    pub async fn chunk(&mut self) -> Result<Option<Bytes>, BoxError> {
282        // loop to ignore unrecognized frames
283        loop {
284            if let Some(res) = self.frame().await {
285                let frame = res?;
286                if let Ok(buf) = frame.into_data() {
287                    return Ok(Some(buf));
288                }
289                // else continue
290            } else {
291                return Ok(None);
292            }
293        }
294    }
295}
296
297impl Default for Body {
298    fn default() -> Self {
299        Self::empty()
300    }
301}
302
303impl From<()> for Body {
304    fn from(_: ()) -> Self {
305        Self::empty()
306    }
307}
308
309macro_rules! body_from_impl {
310    ($ty:ty) => {
311        impl From<$ty> for Body {
312            fn from(buf: $ty) -> Self {
313                Self::new(crate::body::util::Full::from(buf))
314            }
315        }
316    };
317}
318
319body_from_impl!(&'static [u8]);
320body_from_impl!(std::borrow::Cow<'static, [u8]>);
321body_from_impl!(Vec<u8>);
322
323body_from_impl!(&'static str);
324body_from_impl!(std::borrow::Cow<'static, str>);
325body_from_impl!(String);
326
327impl From<&ArcStr> for Body {
328    fn from(buf: &ArcStr) -> Self {
329        Self::new(if let Some(s) = ArcStr::as_static(buf) {
330            crate::body::util::Full::from(s)
331        } else {
332            crate::body::util::Full::from(buf.to_string())
333        })
334    }
335}
336
337impl From<ArcStr> for Body {
338    #[inline(always)]
339    fn from(buf: ArcStr) -> Self {
340        Self::from(&buf)
341    }
342}
343
344body_from_impl!(Bytes);
345
346impl StreamingBody for Body {
347    type Data = Bytes;
348    type Error = BoxError;
349
350    #[inline]
351    fn poll_frame(
352        mut self: Pin<&mut Self>,
353        cx: &mut Context<'_>,
354    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
355        Pin::new(&mut self.0).poll_frame(cx)
356    }
357
358    #[inline]
359    fn size_hint(&self) -> SizeHint {
360        self.0.size_hint()
361    }
362
363    #[inline]
364    fn is_end_stream(&self) -> bool {
365        self.0.is_end_stream()
366    }
367}
368
369/// A stream of data frames.
370///
371/// Created with [`Body::into_data_stream`].
372#[derive(Debug)]
373#[must_use]
374pub struct BodyDataStream {
375    inner: Body,
376}
377
378impl Stream for BodyDataStream {
379    type Item = Result<Bytes, BoxError>;
380
381    #[inline]
382    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
383        loop {
384            match rama_core::futures::ready!(Pin::new(&mut self.inner).poll_frame(cx)?) {
385                Some(frame) => match frame.into_data() {
386                    Ok(data) => return Poll::Ready(Some(Ok(data))),
387                    Err(_frame) => {}
388                },
389                None => return Poll::Ready(None),
390            }
391        }
392    }
393}
394
395impl StreamingBody for BodyDataStream {
396    type Data = Bytes;
397    type Error = BoxError;
398
399    #[inline]
400    fn poll_frame(
401        mut self: Pin<&mut Self>,
402        cx: &mut Context<'_>,
403    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
404        Pin::new(&mut self.inner).poll_frame(cx)
405    }
406
407    #[inline]
408    fn is_end_stream(&self) -> bool {
409        self.inner.is_end_stream()
410    }
411
412    #[inline]
413    fn size_hint(&self) -> SizeHint {
414        self.inner.size_hint()
415    }
416}
417
418pin_project! {
419    struct StreamBody<S> {
420        #[pin]
421        stream: SyncWrapper<S>,
422    }
423}
424
425impl<S> StreamingBody for StreamBody<S>
426where
427    S: TryStream<Ok: Into<Bytes>, Error: Into<BoxError>>,
428{
429    type Data = Bytes;
430    type Error = BoxError;
431
432    fn poll_frame(
433        self: Pin<&mut Self>,
434        cx: &mut Context<'_>,
435    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
436        let stream = self.project().stream.get_pin_mut();
437        match rama_core::futures::ready!(stream.try_poll_next(cx)) {
438            Some(Ok(chunk)) => Poll::Ready(Some(Ok(Frame::data(chunk.into())))),
439            Some(Err(err)) => Poll::Ready(Some(Err(err.into()))),
440            None => Poll::Ready(None),
441        }
442    }
443}
444
445#[test]
446fn test_try_downcast() {
447    assert_eq!(try_downcast::<i32, _>(5_u32), Err(5_u32));
448    assert_eq!(try_downcast::<i32, _>(5_i32), Ok(5_i32));
449}
450
451#[tokio::test]
452async fn capture_convenience_forwards_and_observes_body() {
453    let (sink, mut events) = tokio::sync::mpsc::unbounded_channel();
454    let forwarded = Body::from("captured").capture(sink);
455
456    assert_eq!(
457        forwarded.collect().await.unwrap().to_bytes(),
458        Bytes::from_static(b"captured")
459    );
460    let BodyCaptureEvent::Frame(frame) = events.recv().await.unwrap() else {
461        panic!("the first capture event must contain the data frame");
462    };
463    assert_eq!(frame.into_data().unwrap(), Bytes::from_static(b"captured"));
464    assert!(matches!(
465        events.recv().await,
466        Some(BodyCaptureEvent::End(CaptureOutcome::Complete))
467    ));
468}
rama_http_types/body/mod.rs

rama_http_types/body/
mod.rs
