Skip to main content

axum_security/cookie/
service.rs

1use std::{
2    convert::Infallible,
3    pin::Pin,
4    task::{Context, Poll},
5};
6
7use axum::{extract::Request, response::IntoResponse};
8use tower::{Layer, Service};
9
10use crate::cookie::CookieContext;
11
12type BoxFuture<T> = Pin<Box<dyn Future<Output = T> + Send>>;
13
14pub struct CookieService<S, SERV> {
15    inner: CookieContext<S>,
16    rest: SERV,
17}
18
19impl<S, SERV> Service<Request> for CookieService<S, SERV>
20where
21    SERV: Service<Request, Error = Infallible> + Clone + Send + 'static,
22    <SERV as Service<Request>>::Response: IntoResponse,
23    <SERV as Service<Request>>::Future: Send,
24    S: Clone + Send + Sync + 'static,
25{
26    type Response = axum::response::Response;
27    type Error = Infallible;
28    type Future = BoxFuture<Result<Self::Response, Self::Error>>;
29
30    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
31        self.rest.poll_ready(cx)
32    }
33
34    fn call(&mut self, mut req: Request) -> Self::Future {
35        let mut this = self.clone();
36        Box::pin(async move {
37            match this.inner.load_from_headers(req.headers()).await {
38                Ok(Some(session)) => {
39                    crate::debug!("cookie: session loaded");
40                    req.extensions_mut().insert(session);
41                }
42                Ok(None) => {
43                    crate::debug!("cookie: no session in request");
44                }
45                Err(e) => {
46                    crate::debug!("cookie: error loading session");
47                    return Ok(e);
48                }
49            }
50
51            this.rest.call(req).await.map(IntoResponse::into_response)
52        })
53    }
54}
55
56impl<SERV, T> Layer<SERV> for CookieContext<T>
57where
58    T: 'static,
59{
60    type Service = CookieService<T, SERV>;
61
62    fn layer(&self, inner: SERV) -> Self::Service {
63        CookieService {
64            inner: self.clone(),
65            rest: inner,
66        }
67    }
68}
69
70impl<T, SERV> Clone for CookieService<T, SERV>
71where
72    SERV: Clone,
73{
74    fn clone(&self) -> Self {
75        Self {
76            inner: self.inner.clone(),
77            rest: self.rest.clone(),
78        }
79    }
80}