1  use futures_util::{future::MapOk, TryFutureExt};
2  use std::fmt;
3  use std::task::{Context, Poll};
4  use tower_layer::Layer;
5  use tower_service::Service;
6  
7  /// Service returned by the [`map_response`] combinator.
8  ///
9  /// [`map_response`]: crate::util::ServiceExt::map_response
10  #[derive(Clone)]
11  pub struct MapResponse<S, F> {
12      inner: S,
13      f: F,
14  }
15  
16  impl<S, F> fmt::Debug for MapResponse<S, F>
17  where
18      S: fmt::Debug,
19  {
fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result20      fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
21          f.debug_struct("MapResponse")
22              .field("inner", &self.inner)
23              .field("f", &format_args!("{}", std::any::type_name::<F>()))
24              .finish()
25      }
26  }
27  
28  /// A [`Layer`] that produces a [`MapResponse`] service.
29  ///
30  /// [`Layer`]: tower_layer::Layer
31  #[derive(Debug, Clone)]
32  pub struct MapResponseLayer<F> {
33      f: F,
34  }
35  
36  opaque_future! {
37      /// Response future from [`MapResponse`] services.
38      ///
39      /// [`MapResponse`]: crate::util::MapResponse
40      pub type MapResponseFuture<F, N> = MapOk<F, N>;
41  }
42  
43  impl<S, F> MapResponse<S, F> {
44      /// Creates a new `MapResponse` service.
new(inner: S, f: F) -> Self45      pub fn new(inner: S, f: F) -> Self {
46          MapResponse { f, inner }
47      }
48  
49      /// Returns a new [`Layer`] that produces [`MapResponse`] services.
50      ///
51      /// This is a convenience function that simply calls [`MapResponseLayer::new`].
52      ///
53      /// [`Layer`]: tower_layer::Layer
layer(f: F) -> MapResponseLayer<F>54      pub fn layer(f: F) -> MapResponseLayer<F> {
55          MapResponseLayer { f }
56      }
57  }
58  
59  impl<S, F, Request, Response> Service<Request> for MapResponse<S, F>
60  where
61      S: Service<Request>,
62      F: FnOnce(S::Response) -> Response + Clone,
63  {
64      type Response = Response;
65      type Error = S::Error;
66      type Future = MapResponseFuture<S::Future, F>;
67  
68      #[inline]
poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>>69      fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
70          self.inner.poll_ready(cx)
71      }
72  
73      #[inline]
call(&mut self, request: Request) -> Self::Future74      fn call(&mut self, request: Request) -> Self::Future {
75          MapResponseFuture::new(self.inner.call(request).map_ok(self.f.clone()))
76      }
77  }
78  
79  impl<F> MapResponseLayer<F> {
80      /// Creates a new [`MapResponseLayer`] layer.
new(f: F) -> Self81      pub fn new(f: F) -> Self {
82          MapResponseLayer { f }
83      }
84  }
85  
86  impl<S, F> Layer<S> for MapResponseLayer<F>
87  where
88      F: Clone,
89  {
90      type Service = MapResponse<S, F>;
91  
layer(&self, inner: S) -> Self::Service92      fn layer(&self, inner: S) -> Self::Service {
93          MapResponse {
94              f: self.f.clone(),
95              inner,
96          }
97      }
98  }
99