1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

use std::{
    pin::Pin,
    task::{Context, Poll},
};

use futures_util::{Future, FutureExt};
use pin_project_lite::pin_project;
use tower::{Layer, Service};

use crate::{request::BoxRequest, response::BoxResponse};

/// Layer for creating [`Modify`] instances.
#[derive(Clone)]
pub struct ModifyLayer<ModifyReq, ModifyResp> {
    req_fn: ModifyReq,
    resp_fn: ModifyResp,
}

impl<ModifyReq, ModifyResp> ModifyLayer<ModifyReq, ModifyResp>
where
    ModifyReq: Fn(&mut BoxRequest),
    ModifyResp: Fn(&mut BoxResponse),
{
    /// Create a new layer.
    pub fn new(req_fn: ModifyReq, resp_fn: ModifyResp) -> Self {
        Self { req_fn, resp_fn }
    }
}

impl<ModifyReq> ModifyLayer<ModifyReq, fn(&mut BoxResponse)>
where
    ModifyReq: Fn(&mut BoxRequest),
{
    /// Create a new layer that only modifies requests.
    pub fn new_request(f: ModifyReq) -> Self {
        Self::new(f, |_| ())
    }
}

impl<ModifyResp> ModifyLayer<fn(&mut BoxRequest), ModifyResp>
where
    ModifyResp: Fn(&mut BoxResponse),
{
    /// Create a new layer that only modifies responses.
    pub fn new_response(f: ModifyResp) -> Self {
        Self::new(|_| (), f)
    }
}

impl<ModifyReq, ModifyResp, S> Layer<S> for ModifyLayer<ModifyReq, ModifyResp>
where
    ModifyReq: Clone,
    ModifyResp: Clone,
{
    type Service = Modify<ModifyReq, ModifyResp, S>;

    fn layer(&self, inner: S) -> Self::Service {
        Modify::new(inner, self.req_fn.clone(), self.resp_fn.clone())
    }
}

/// Service that lets you modify / inspect requests and responses.
#[derive(Clone)]
pub struct Modify<ModifyReq, ModifyResp, S> {
    inner: S,
    req_fn: ModifyReq,
    resp_fn: ModifyResp,
}

impl<ModifyReq, ModifyResp, S> Modify<ModifyReq, ModifyResp, S> {
    /// Create a new service by wrapping a given service.
    pub fn new(inner: S, req_fn: ModifyReq, resp_fn: ModifyResp) -> Self {
        Self {
            inner,
            req_fn,
            resp_fn,
        }
    }
}

impl<ModifyReq, ModifyResp, S> Service<BoxRequest> for Modify<ModifyReq, ModifyResp, S>
where
    S: Service<BoxRequest, Response = BoxResponse>,
    ModifyReq: Fn(&mut BoxRequest),
    ModifyResp: Fn(&mut BoxResponse) + Clone,
{
    type Response = BoxResponse;

    type Error = S::Error;

    type Future = ModifyFuture<ModifyResp, S::Future>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Service::poll_ready(&mut self.inner, cx)
    }

    fn call(&mut self, mut req: BoxRequest) -> Self::Future {
        (self.req_fn)(&mut req);

        ModifyFuture {
            fut: Service::call(&mut self.inner, req),
            resp_fn: self.resp_fn.clone(),
        }
    }
}

pin_project! {
    /// Future for [`Modify`].
    pub struct ModifyFuture<ModifyResp, Fut> {
        #[pin]
        fut: Fut,
        resp_fn: ModifyResp,
    }
}

impl<ModifyResp, Fut, Err> Future for ModifyFuture<ModifyResp, Fut>
where
    Fut: Future<Output = Result<BoxResponse, Err>>,
    ModifyResp: Fn(&mut BoxResponse),
{
    type Output = Result<BoxResponse, Err>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut this = self.project();

        this.fut.poll_unpin(cx).map_ok(|mut resp| {
            (this.resp_fn)(&mut resp);
            resp
        })
    }
}