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

//! Middleware to use [`Cookies`].

use self::future::ResponseFuture;
use crate::Cookies;
use http::{header, Request, Response};
use std::task::{Context, Poll};
use tower_layer::Layer;
use tower_service::Service;

pub mod future;

/// Middleware to use [`Cookies`].
#[derive(Clone, Debug)]
pub struct CookieManager<S> {
    inner: S,
}

impl<S> CookieManager<S> {
    /// Create a new cookie manager.
    pub fn new(inner: S) -> Self {
        Self { inner }
    }
}

impl<ReqBody, ResBody, S> Service<Request<ReqBody>> for CookieManager<S>
where
    S: Service<Request<ReqBody>, Response = Response<ResBody>>,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = ResponseFuture<S::Future>;

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

    fn call(&mut self, mut req: Request<ReqBody>) -> Self::Future {
        let value = req
            .headers()
            .get_all(header::COOKIE)
            .iter()
            .cloned()
            .collect();
        let cookies = Cookies::new(value);
        req.extensions_mut().insert(cookies.clone());

        ResponseFuture {
            future: self.inner.call(req),
            cookies,
        }
    }
}

/// Layer to apply [`CookieManager`] middleware.
#[derive(Clone, Debug, Default)]
pub struct CookieManagerLayer {
    _priv: (),
}

impl CookieManagerLayer {
    /// Create a new cookie manager layer.
    pub fn new() -> Self {
        Self { _priv: () }
    }
}

impl<S> Layer<S> for CookieManagerLayer {
    type Service = CookieManager<S>;

    fn layer(&self, inner: S) -> Self::Service {
        CookieManager { inner }
    }
}