rama_http/layer/
map_response_body.rs

1//! Apply a transformation to the response body.
2//!
3//! # Example
4//!
5//! ```
6//! use rama_core::bytes::Bytes;
7//! use rama_http::{Body, Request, Response, StreamingBody, body::{Frame, SizeHint}};
8//! use std::convert::Infallible;
9//! use std::{pin::Pin, task::{Context, Poll}};
10//! use rama_core::{Layer, Service};
11//! use rama_core::service::service_fn;
12//! use rama_http::layer::map_response_body::MapResponseBodyLayer;
13//! use rama_core::error::BoxError;
14//! use rama_core::futures::ready;
15//!
16//! // A wrapper for a `http_body::Body` that prints the size of data chunks
17//! pin_project_lite::pin_project! {
18//!     struct PrintChunkSizesBody<B> {
19//!         #[pin]
20//!         inner: B,
21//!     }
22//! }
23//!
24//! impl<B> PrintChunkSizesBody<B> {
25//!     fn new(inner: B) -> Self {
26//!         Self { inner }
27//!     }
28//! }
29//!
30//! impl<B> StreamingBody for PrintChunkSizesBody<B>
31//!     where B: StreamingBody<Data = Bytes, Error = BoxError>,
32//! {
33//!     type Data = Bytes;
34//!     type Error = BoxError;
35//!
36//!     fn poll_frame(
37//!         mut self: Pin<&mut Self>,
38//!         cx: &mut Context<'_>,
39//!     ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
40//!         let inner_body = self.as_mut().project().inner;
41//!         if let Some(frame) = ready!(inner_body.poll_frame(cx)?) {
42//!             if let Some(chunk) = frame.data_ref() {
43//!                 println!("chunk size = {}", chunk.len());
44//!             } else {
45//!                 eprintln!("no data chunk found");
46//!             }
47//!             Poll::Ready(Some(Ok(frame)))
48//!         } else {
49//!             Poll::Ready(None)
50//!         }
51//!     }
52//!
53//!     fn is_end_stream(&self) -> bool {
54//!         self.inner.is_end_stream()
55//!     }
56//!
57//!     fn size_hint(&self) -> SizeHint {
58//!         self.inner.size_hint()
59//!     }
60//! }
61//!
62//! async fn handle(_: Request) -> Result<Response, Infallible> {
63//!     // ...
64//!     # Ok(Response::new(Body::default()))
65//! }
66//!
67//! # #[tokio::main]
68//! # async fn main() -> Result<(), Box<dyn std::error::Error>> {
69//! let mut svc = (
70//!     // Wrap response bodies in `PrintChunkSizesBody`
71//!     MapResponseBodyLayer::new(PrintChunkSizesBody::new),
72//! ).into_layer(service_fn(handle));
73//!
74//! // Call the service
75//! let request = Request::new(Body::from("foobar"));
76//!
77//! svc.serve(request).await?;
78//! # Ok(())
79//! # }
80//! ```
81
82use crate::{Request, Response};
83use rama_core::{Layer, Service, bytes::Bytes, error::BoxError};
84use rama_http_types::StreamingBody;
85use rama_utils::macros::define_inner_service_accessors;
86use std::fmt;
87
88/// Apply a transformation to the response body.
89///
90/// See the [module docs](crate::layer::map_response_body) for an example.
91#[derive(Clone)]
92pub struct MapResponseBodyLayer<F> {
93    f: F,
94}
95
96impl<F> MapResponseBodyLayer<F> {
97    /// Create a new [`MapResponseBodyLayer`].
98    ///
99    /// `F` is expected to be a function that takes a body and returns another body.
100    pub const fn new(f: F) -> Self {
101        Self { f }
102    }
103}
104
105impl<Body> MapResponseBodyLayer<fn(Body) -> crate::Body>
106where
107    Body: StreamingBody<Data = Bytes, Error: Into<BoxError>> + Send + Sync + 'static,
108{
109    /// Turn the [`StreamingBody`] into a [`crate::Body`] using [`crate::Body::new`].
110    ///
111    /// This is shorthand for `MapResponseBodyLayer::new(Body::new)`.
112    pub const fn new_boxed_streaming_body() -> Self {
113        Self::new(crate::Body::new)
114    }
115}
116
117impl<Body> MapResponseBodyLayer<fn(Body) -> crate::Body>
118where
119    crate::Body: From<Body>,
120{
121    /// Turn the body into a [`crate::Body`] using [`crate::Body::from`].
122    ///
123    /// This is shorthand for `MapResponseBodyLayer::new(Body::from)`.
124    pub const fn into_boxed_streaming_body() -> Self {
125        Self::new(crate::Body::from)
126    }
127}
128
129impl<S, F> Layer<S> for MapResponseBodyLayer<F>
130where
131    F: Clone,
132{
133    type Service = MapResponseBody<S, F>;
134
135    fn layer(&self, inner: S) -> Self::Service {
136        MapResponseBody::new(inner, self.f.clone())
137    }
138
139    fn into_layer(self, inner: S) -> Self::Service {
140        MapResponseBody::new(inner, self.f)
141    }
142}
143
144impl<F> fmt::Debug for MapResponseBodyLayer<F> {
145    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
146        f.debug_struct("MapResponseBodyLayer")
147            .field("f", &std::any::type_name::<F>())
148            .finish()
149    }
150}
151
152/// Apply a transformation to the response body.
153///
154/// See the [module docs](crate::layer::map_response_body) for an example.
155#[derive(Clone)]
156pub struct MapResponseBody<S, F> {
157    inner: S,
158    f: F,
159}
160
161impl<S, F> MapResponseBody<S, F> {
162    /// Create a new [`MapResponseBody`].
163    ///
164    /// `F` is expected to be a function that takes a body and returns another body.
165    pub const fn new(service: S, f: F) -> Self {
166        Self { inner: service, f }
167    }
168
169    define_inner_service_accessors!();
170}
171
172impl<S, Body> MapResponseBody<S, fn(Body) -> crate::Body>
173where
174    Body: StreamingBody<Data = Bytes, Error: Into<BoxError>> + Send + Sync + 'static,
175{
176    /// Turn the [`StreamingBody`] into a [`crate::Body`] using [`crate::Body::new`].
177    ///
178    /// This is shorthand for `MapResponseBody::new(..., Body::new)`.
179    pub const fn new_boxed_streaming_body(inner: S) -> Self {
180        Self::new(inner, crate::Body::new)
181    }
182}
183
184impl<S, Body> MapResponseBody<S, fn(Body) -> crate::Body>
185where
186    crate::Body: From<Body>,
187{
188    /// Turn the body into a [`crate::Body`] using [`crate::Body::from`].
189    ///
190    /// This is shorthand for `MapResponseBody::new(..., Body::from)`.
191    pub const fn into_boxed_streaming_body(inner: S) -> Self {
192        Self::new(inner, crate::Body::from)
193    }
194}
195
196impl<F, S, ReqBody, ResBody, NewResBody> Service<Request<ReqBody>> for MapResponseBody<S, F>
197where
198    S: Service<Request<ReqBody>, Output = Response<ResBody>>,
199    ReqBody: Send + 'static,
200    ResBody: Send + Sync + 'static,
201    NewResBody: Send + Sync + 'static,
202    F: Fn(ResBody) -> NewResBody + Clone + Send + Sync + 'static,
203{
204    type Output = Response<NewResBody>;
205    type Error = S::Error;
206
207    async fn serve(&self, req: Request<ReqBody>) -> Result<Self::Output, Self::Error> {
208        let res = self.inner.serve(req).await?;
209        Ok(res.map(self.f.clone()))
210    }
211}
212
213impl<S, F> fmt::Debug for MapResponseBody<S, F>
214where
215    S: fmt::Debug,
216{
217    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
218        f.debug_struct("MapResponseBody")
219            .field("inner", &self.inner)
220            .field("f", &std::any::type_name::<F>())
221            .finish()
222    }
223}
rama_http/layer/map_response_body.rs

rama_http/layer/
map_response_body.rs
