satex_layer/remove_header/
remove_response_header.rs

1use crate::make::MakeRouteLayer;
2use crate::remove_header::Removable;
3use http::HeaderName;
4use pin_project_lite::pin_project;
5use satex_core::component::Args;
6use satex_core::{component::Configurable, Error};
7use satex_macro::make;
8use std::pin::Pin;
9use std::task::{ready, Context, Poll};
10use std::{str::FromStr, sync::Arc};
11use tower::{Layer, Service};
12
13#[make(kind = RemoveResponseHeader)]
14pub struct MakeRemoveResponseHeaderRouteLayer {
15    name: String,
16}
17
18impl MakeRouteLayer for MakeRemoveResponseHeaderRouteLayer {
19    type Layer = RemoveResponseHeaderLayer;
20
21    fn make(&self, args: Args) -> Result<Self::Layer, Error> {
22        Config::with_args(args).and_then(|config| {
23            HeaderName::from_str(&config.name)
24                .map_err(Error::new)
25                .map(RemoveResponseHeaderLayer::new)
26        })
27    }
28}
29
30#[derive(Debug, Clone)]
31pub struct RemoveResponseHeaderLayer {
32    name: Arc<HeaderName>,
33}
34
35impl RemoveResponseHeaderLayer {
36    pub fn new(name: HeaderName) -> Self {
37        Self {
38            name: Arc::new(name),
39        }
40    }
41}
42
43impl<S> Layer<S> for RemoveResponseHeaderLayer {
44    type Service = RemoveResponseHeader<S>;
45
46    fn layer(&self, inner: S) -> Self::Service {
47        RemoveResponseHeader {
48            name: self.name.clone(),
49            inner,
50        }
51    }
52}
53
54#[derive(Debug, Clone)]
55pub struct RemoveResponseHeader<S> {
56    inner: S,
57    name: Arc<HeaderName>,
58}
59
60impl<S, Req, Res> Service<Req> for RemoveResponseHeader<S>
61where
62    S: Service<Req, Response=Res>,
63    Res: Removable,
64{
65    type Response = S::Response;
66
67    type Error = S::Error;
68
69    type Future = ResponseFuture<S::Future>;
70
71    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
72        self.inner.poll_ready(cx)
73    }
74
75    fn call(&mut self, request: Req) -> Self::Future {
76        let name = self.name.clone();
77        ResponseFuture::new(name, self.inner.call(request))
78    }
79}
80
81pin_project! {
82    #[project = ResponseFutureProj]
83    #[project_replace = ResponseFutureProjReplace]
84    pub enum ResponseFuture<F> {
85        Incomplete {
86            name: Arc<HeaderName>,
87            #[pin]
88            future: F,
89        },
90        Complete,
91    }
92}
93
94impl<F> ResponseFuture<F> {
95    pub fn new(name: Arc<HeaderName>, future: F) -> Self {
96        Self::Incomplete { name, future }
97    }
98}
99
100impl<F, Res, E> Future for ResponseFuture<F>
101where
102    F: Future<Output=Result<Res, E>>,
103    Res: Removable,
104{
105    type Output = F::Output;
106
107    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
108        match self.as_mut().project() {
109            ResponseFutureProj::Incomplete { future, .. } => match ready!(future.poll(cx)) {
110                Ok(mut res) => match self.project_replace(ResponseFuture::Complete) {
111                    ResponseFutureProjReplace::Incomplete { name, .. } => {
112                        res.remove(name.as_ref());
113                        Poll::Ready(Ok(res))
114                    }
115                    ResponseFutureProjReplace::Complete => unreachable!(),
116                },
117                Err(e) => Poll::Ready(Err(e)),
118            },
119            ResponseFutureProj::Complete { .. } => {
120                panic!("polled after completion")
121            }
122        }
123    }
124}
satex_layer/remove_header/remove_response_header.rs

satex_layer/remove_header/
remove_response_header.rs
