tower_http/
map_response_body.rs

Help
1//! Apply a transformation to the response body.
2//!
3//! # Example
4//!
5//! ```
6//! use bytes::Bytes;
7//! use http::{Request, Response};
8//! use http_body_util::Full;
9//! use std::convert::Infallible;
10//! use std::{pin::Pin, task::{ready, Context, Poll}};
11//! use tower::{ServiceBuilder, service_fn, ServiceExt, Service};
12//! use tower_http::map_response_body::MapResponseBodyLayer;
13//!
14//! // A wrapper for a `Full<Bytes>`
15//! struct BodyWrapper {
16//!     inner: Full<Bytes>,
17//! }
18//!
19//! impl BodyWrapper {
20//!     fn new(inner: Full<Bytes>) -> Self {
21//!         Self { inner }
22//!     }
23//! }
24//!
25//! impl http_body::Body for BodyWrapper {
26//!     // ...
27//!     # type Data = Bytes;
28//!     # type Error = tower::BoxError;
29//!     # fn poll_frame(
30//!     #     self: Pin<&mut Self>,
31//!     #     cx: &mut Context<'_>
32//!     # ) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> { unimplemented!() }
33//!     # fn is_end_stream(&self) -> bool { unimplemented!() }
34//!     # fn size_hint(&self) -> http_body::SizeHint { unimplemented!() }
35//! }
36//!
37//! async fn handle<B>(_: Request<B>) -> Result<Response<Full<Bytes>>, Infallible> {
38//!     // ...
39//!     # Ok(Response::new(Full::default()))
40//! }
41//!
42//! # #[tokio::main]
43//! # async fn main() -> Result<(), Box<dyn std::error::Error>> {
44//! let mut svc = ServiceBuilder::new()
45//!     // Wrap response bodies in `BodyWrapper`
46//!     .layer(MapResponseBodyLayer::new(BodyWrapper::new))
47//!     .service_fn(handle);
48//!
49//! // Call the service
50//! let request = Request::new(Full::<Bytes>::from("foobar"));
51//!
52//! svc.ready().await?.call(request).await?;
53//! # Ok(())
54//! # }
55//! ```
56
57use http::{Request, Response};
58use pin_project_lite::pin_project;
59use std::future::Future;
60use std::{
61    fmt,
62    pin::Pin,
63    task::{ready, Context, Poll},
64};
65use tower_layer::Layer;
66use tower_service::Service;
67
68/// Apply a transformation to the response body.
69///
70/// See the [module docs](crate::map_response_body) for an example.
71#[derive(Clone)]
72pub struct MapResponseBodyLayer<F> {
73    f: F,
74}
75
76impl<F> MapResponseBodyLayer<F> {
77    /// Create a new [`MapResponseBodyLayer`].
78    ///
79    /// `F` is expected to be a function that takes a body and returns another body.
80    pub fn new(f: F) -> Self {
81        Self { f }
82    }
83}
84
85impl<S, F> Layer<S> for MapResponseBodyLayer<F>
86where
87    F: Clone,
88{
89    type Service = MapResponseBody<S, F>;
90
91    fn layer(&self, inner: S) -> Self::Service {
92        MapResponseBody::new(inner, self.f.clone())
93    }
94}
95
96impl<F> fmt::Debug for MapResponseBodyLayer<F> {
97    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
98        f.debug_struct("MapResponseBodyLayer")
99            .field("f", &std::any::type_name::<F>())
100            .finish()
101    }
102}
103
104/// Apply a transformation to the response body.
105///
106/// See the [module docs](crate::map_response_body) for an example.
107#[derive(Clone)]
108pub struct MapResponseBody<S, F> {
109    inner: S,
110    f: F,
111}
112
113impl<S, F> MapResponseBody<S, F> {
114    /// Create a new [`MapResponseBody`].
115    ///
116    /// `F` is expected to be a function that takes a body and returns another body.
117    pub fn new(service: S, f: F) -> Self {
118        Self { inner: service, f }
119    }
120
121    /// Returns a new [`Layer`] that wraps services with a `MapResponseBodyLayer` middleware.
122    ///
123    /// [`Layer`]: tower_layer::Layer
124    pub fn layer(f: F) -> MapResponseBodyLayer<F> {
125        MapResponseBodyLayer::new(f)
126    }
127
128    define_inner_service_accessors!();
129}
130
131impl<F, S, ReqBody, ResBody, NewResBody> Service<Request<ReqBody>> for MapResponseBody<S, F>
132where
133    S: Service<Request<ReqBody>, Response = Response<ResBody>>,
134    F: FnMut(ResBody) -> NewResBody + Clone,
135{
136    type Response = Response<NewResBody>;
137    type Error = S::Error;
138    type Future = ResponseFuture<S::Future, F>;
139
140    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
141        self.inner.poll_ready(cx)
142    }
143
144    fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
145        ResponseFuture {
146            inner: self.inner.call(req),
147            f: self.f.clone(),
148        }
149    }
150}
151
152impl<S, F> fmt::Debug for MapResponseBody<S, F>
153where
154    S: fmt::Debug,
155{
156    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
157        f.debug_struct("MapResponseBody")
158            .field("inner", &self.inner)
159            .field("f", &std::any::type_name::<F>())
160            .finish()
161    }
162}
163
164pin_project! {
165    /// Response future for [`MapResponseBody`].
166    pub struct ResponseFuture<Fut, F> {
167        #[pin]
168        inner: Fut,
169        f: F,
170    }
171}
172
173impl<Fut, F, ResBody, E, NewResBody> Future for ResponseFuture<Fut, F>
174where
175    Fut: Future<Output = Result<Response<ResBody>, E>>,
176    F: FnMut(ResBody) -> NewResBody,
177{
178    type Output = Result<Response<NewResBody>, E>;
179
180    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
181        let this = self.project();
182        let res = ready!(this.inner.poll(cx)?);
183        Poll::Ready(Ok(res.map(this.f)))
184    }
185}
tower_http/map_response_body.rs

tower_http/
map_response_body.rs
