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
56// Implementations copied over from http-body but addapted to work with our Requests/Response types
57
58impl<B: StreamingBody> StreamingBody for crate::Request<B> {
59    type Data = B::Data;
60    type Error = B::Error;
61
62    fn poll_frame(
63        self: Pin<&mut Self>,
64        cx: &mut Context<'_>,
65    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
66        // SAFETY:
67        // A pin projection.
68        unsafe { self.map_unchecked_mut(Self::body_mut).poll_frame(cx) }
69    }
70
71    fn is_end_stream(&self) -> bool {
72        self.body().is_end_stream()
73    }
74
75    fn size_hint(&self) -> SizeHint {
76        self.body().size_hint()
77    }
78}
79
80impl<B: StreamingBody> StreamingBody for crate::Response<B> {
81    type Data = B::Data;
82    type Error = B::Error;
83
84    fn poll_frame(
85        self: Pin<&mut Self>,
86        cx: &mut Context<'_>,
87    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
88        // SAFETY:
89        // A pin projection.
90        unsafe { self.map_unchecked_mut(Self::body_mut).poll_frame(cx) }
91    }
92
93    fn is_end_stream(&self) -> bool {
94        self.body().is_end_stream()
95    }
96
97    fn size_hint(&self) -> SizeHint {
98        self.body().size_hint()
99    }
100}
101
102// Rama specific stuff
103
104type BoxBody = util::combinators::BoxBody<Bytes, BoxError>;
105
106fn boxed<B>(body: B) -> BoxBody
107where
108    B: StreamingBody<Data = Bytes, Error: Into<BoxError>> + Send + Sync + 'static,
109{
110    try_downcast(body).unwrap_or_else(|body| body.map_err(Into::into).boxed())
111}
112
113#[expect(clippy::unwrap_used)]
114pub(crate) fn try_downcast<T, K>(k: K) -> Result<T, K>
115where
116    T: 'static,
117    K: Send + 'static,
118{
119    let mut k = Some(k);
120    if let Some(k) = <dyn std::any::Any>::downcast_mut::<Option<T>>(&mut k) {
121        Ok(k.take().unwrap())
122    } else {
123        Err(k.unwrap())
124    }
125}
126
127/// The body type used in rama requests and responses.
128#[must_use]
129#[derive(Debug)]
130pub struct Body(BoxBody);
131
132impl Body {
133    /// Create a new `Body` that wraps another [`Body`].
134    pub fn new<B>(body: B) -> Self
135    where
136        B: StreamingBody<Data = Bytes, Error: Into<BoxError>> + Send + Sync + 'static,
137    {
138        try_downcast(body).unwrap_or_else(|body| Self(boxed(body)))
139    }
140
141    /// Create a new `Body` with a maximum size limit.
142    pub fn with_limit<B>(body: B, limit: usize) -> Self
143    where
144        B: StreamingBody<Data = Bytes, Error: Into<BoxError>> + Send + Sync + 'static,
145    {
146        Self::new(util::Limited::new(body, limit))
147    }
148
149    /// Create an empty body.
150    pub fn empty() -> Self {
151        Self::new(util::Empty::new())
152    }
153
154    /// Create a new `Body` from a [`Stream`].
155    ///
156    /// [`Stream`]: https://docs.rs/futures/latest/futures/stream/trait.Stream.html
157    pub fn from_stream<S>(stream: S) -> Self
158    where
159        S: TryStream<Ok: Into<Bytes>, Error: Into<BoxError>> + Send + 'static,
160    {
161        Self::new(StreamBody {
162            stream: SyncWrapper::new(stream),
163        })
164    }
165
166    /// Create a new [`Body`] from a [`Stream`] with a maximum size limit.
167    pub fn limited(self, limit: usize) -> Self {
168        Self::new(util::Limited::new(self.0, limit))
169    }
170
171    /// Wrap this body so that `on_drop` is called if the body is dropped before
172    /// being fully consumed (e.g. because the client disconnected mid-response).
173    ///
174    /// The closure fires exactly once and only on early drops — it is disarmed
175    /// when [`StreamingBody::poll_frame`] returns `Poll::Ready(None)`.
176    pub fn on_drop<F>(self, on_drop: F) -> Self
177    where
178        F: FnOnce() + Send + Sync + 'static,
179    {
180        Self::new(OnDropBody::new(self.0, on_drop))
181    }
182
183    /// Attach a headermap as trailer headers to this body.
184    ///
185    /// In case you want this is also available as `with_trailer`,
186    /// provided by a body extension which allows you to use
187    /// a [`Future`] with as output an option of a result of a headermap.
188    /// This method is a shortcut of that function.
189    pub fn with_trailer_headers(self, headers: crate::HeaderMap) -> Self {
190        Self::new(self.0.with_trailers(std::future::ready(Some(Ok(headers)))))
191    }
192
193    /// Convert the body into a [`Stream`] of data frames.
194    ///
195    /// Non-data frames (such as trailers) will be discarded. Use [`http_body_util::BodyStream`] if
196    /// you need a [`Stream`] of all frame types.
197    ///
198    /// [`http_body_util::BodyStream`]: https://docs.rs/http-body-util/latest/http_body_util/struct.BodyStream.html
199    pub fn into_data_stream(self) -> BodyDataStream {
200        BodyDataStream { inner: self }
201    }
202
203    /// Convert the body into a [`Stream`] of [`sse::Event`]s.
204    ///
205    /// <https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events>.
206    #[must_use]
207    pub fn into_event_stream<T: EventDataRead>(self) -> EventStream<BodyDataStream, T> {
208        EventStream::new(self.into_data_stream())
209    }
210
211    /// Convert the body into a [`JsonReadStream`][json::JsonReadStream].
212    ///
213    /// Stream of json objects, each object separated by a newline (`\n`).
214    #[must_use]
215    pub fn into_json_stream<T: DeserializeOwned>(self) -> json::JsonReadStream<T, BodyDataStream> {
216        let stream = self.into_data_stream();
217        json::JsonReadStream::new(stream)
218    }
219
220    /// Convert the body into a [`JsonReadStream`][json::JsonReadStream].
221    ///
222    /// Stream of json objects, each object separated by a newline (`\n`).
223    #[must_use]
224    pub fn into_json_stream_with_config<T: DeserializeOwned>(
225        self,
226        cfg: json::ParseConfig,
227    ) -> json::JsonReadStream<T, BodyDataStream> {
228        let stream = self.into_data_stream();
229        json::JsonReadStream::new_with_config(stream, cfg)
230    }
231
232    /// Convert the body into a [`Stream`] of [`sse::Event`]s with optional string data.
233    ///
234    /// <https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events>.
235    #[must_use]
236    pub fn into_string_data_event_stream(self) -> EventStream<BodyDataStream> {
237        EventStream::new(self.into_data_stream())
238    }
239
240    /// Stream a chunk of the response body.
241    ///
242    /// When the response body has been exhausted, this will return `None`.
243    pub async fn chunk(&mut self) -> Result<Option<Bytes>, BoxError> {
244        // loop to ignore unrecognized frames
245        loop {
246            if let Some(res) = self.frame().await {
247                let frame = res?;
248                if let Ok(buf) = frame.into_data() {
249                    return Ok(Some(buf));
250                }
251                // else continue
252            } else {
253                return Ok(None);
254            }
255        }
256    }
257}
258
259impl Default for Body {
260    fn default() -> Self {
261        Self::empty()
262    }
263}
264
265impl From<()> for Body {
266    fn from(_: ()) -> Self {
267        Self::empty()
268    }
269}
270
271macro_rules! body_from_impl {
272    ($ty:ty) => {
273        impl From<$ty> for Body {
274            fn from(buf: $ty) -> Self {
275                Self::new(crate::body::util::Full::from(buf))
276            }
277        }
278    };
279}
280
281body_from_impl!(&'static [u8]);
282body_from_impl!(std::borrow::Cow<'static, [u8]>);
283body_from_impl!(Vec<u8>);
284
285body_from_impl!(&'static str);
286body_from_impl!(std::borrow::Cow<'static, str>);
287body_from_impl!(String);
288
289impl From<&ArcStr> for Body {
290    fn from(buf: &ArcStr) -> Self {
291        Self::new(if let Some(s) = ArcStr::as_static(buf) {
292            crate::body::util::Full::from(s)
293        } else {
294            crate::body::util::Full::from(buf.to_string())
295        })
296    }
297}
298
299impl From<ArcStr> for Body {
300    #[inline(always)]
301    fn from(buf: ArcStr) -> Self {
302        Self::from(&buf)
303    }
304}
305
306body_from_impl!(Bytes);
307
308impl StreamingBody for Body {
309    type Data = Bytes;
310    type Error = BoxError;
311
312    #[inline]
313    fn poll_frame(
314        mut self: Pin<&mut Self>,
315        cx: &mut Context<'_>,
316    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
317        Pin::new(&mut self.0).poll_frame(cx)
318    }
319
320    #[inline]
321    fn size_hint(&self) -> SizeHint {
322        self.0.size_hint()
323    }
324
325    #[inline]
326    fn is_end_stream(&self) -> bool {
327        self.0.is_end_stream()
328    }
329}
330
331/// A stream of data frames.
332///
333/// Created with [`Body::into_data_stream`].
334#[derive(Debug)]
335#[must_use]
336pub struct BodyDataStream {
337    inner: Body,
338}
339
340impl Stream for BodyDataStream {
341    type Item = Result<Bytes, BoxError>;
342
343    #[inline]
344    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
345        loop {
346            match rama_core::futures::ready!(Pin::new(&mut self.inner).poll_frame(cx)?) {
347                Some(frame) => match frame.into_data() {
348                    Ok(data) => return Poll::Ready(Some(Ok(data))),
349                    Err(_frame) => {}
350                },
351                None => return Poll::Ready(None),
352            }
353        }
354    }
355}
356
357impl StreamingBody for BodyDataStream {
358    type Data = Bytes;
359    type Error = BoxError;
360
361    #[inline]
362    fn poll_frame(
363        mut self: Pin<&mut Self>,
364        cx: &mut Context<'_>,
365    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
366        Pin::new(&mut self.inner).poll_frame(cx)
367    }
368
369    #[inline]
370    fn is_end_stream(&self) -> bool {
371        self.inner.is_end_stream()
372    }
373
374    #[inline]
375    fn size_hint(&self) -> SizeHint {
376        self.inner.size_hint()
377    }
378}
379
380pin_project! {
381    struct StreamBody<S> {
382        #[pin]
383        stream: SyncWrapper<S>,
384    }
385}
386
387impl<S> StreamingBody for StreamBody<S>
388where
389    S: TryStream<Ok: Into<Bytes>, Error: Into<BoxError>>,
390{
391    type Data = Bytes;
392    type Error = BoxError;
393
394    fn poll_frame(
395        self: Pin<&mut Self>,
396        cx: &mut Context<'_>,
397    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
398        let stream = self.project().stream.get_pin_mut();
399        match rama_core::futures::ready!(stream.try_poll_next(cx)) {
400            Some(Ok(chunk)) => Poll::Ready(Some(Ok(Frame::data(chunk.into())))),
401            Some(Err(err)) => Poll::Ready(Some(Err(err.into()))),
402            None => Poll::Ready(None),
403        }
404    }
405}
406
407#[test]
408fn test_try_downcast() {
409    assert_eq!(try_downcast::<i32, _>(5_u32), Err(5_u32));
410    assert_eq!(try_downcast::<i32, _>(5_i32), Ok(5_i32));
411}
rama_http_types/body/mod.rs

rama_http_types/body/
mod.rs
