viz_core/middleware/
cookie.rs

1//! Cookie Middleware.
2
3use std::fmt;
4
5use crate::{
6    header::{HeaderValue, COOKIE, SET_COOKIE},
7    types::{Cookie, CookieJar, Cookies},
8    Handler, IntoResponse, Request, Response, Result, Transform,
9};
10
11#[cfg(any(feature = "cookie-signed", feature = "cookie-private"))]
12use crate::types::CookieKey;
13
14/// A configure for [`CookieMiddleware`].
15pub struct Config {
16    #[cfg(any(feature = "cookie-signed", feature = "cookie-private"))]
17    key: std::sync::Arc<CookieKey>,
18}
19
20impl Config {
21    /// Creates a new config with the [`Key`][CookieKey].
22    #[cfg(any(feature = "cookie-signed", feature = "cookie-private"))]
23    #[must_use]
24    pub fn with_key(key: CookieKey) -> Self {
25        Self {
26            key: std::sync::Arc::new(key),
27        }
28    }
29}
30
31impl Default for Config {
32    fn default() -> Self {
33        Self {
34            #[cfg(any(feature = "cookie-signed", feature = "cookie-private"))]
35            key: std::sync::Arc::new(CookieKey::generate()),
36        }
37    }
38}
39
40impl fmt::Debug for Config {
41    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
42        let mut d = f.debug_struct("CookieConfig");
43
44        #[cfg(any(feature = "cookie-signed", feature = "cookie-private"))]
45        d.field("key", &[..]);
46
47        d.finish()
48    }
49}
50
51impl<H> Transform<H> for Config
52where
53    H: Clone,
54{
55    type Output = CookieMiddleware<H>;
56
57    fn transform(&self, h: H) -> Self::Output {
58        CookieMiddleware {
59            h,
60            #[cfg(any(feature = "cookie-signed", feature = "cookie-private"))]
61            key: self.key.clone(),
62        }
63    }
64}
65
66/// Cookie middleware.
67#[derive(Clone)]
68pub struct CookieMiddleware<H> {
69    h: H,
70    #[cfg(any(feature = "cookie-signed", feature = "cookie-private"))]
71    key: std::sync::Arc<CookieKey>,
72}
73
74impl<H> fmt::Debug for CookieMiddleware<H> {
75    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
76        let mut d = f.debug_struct("CookieMiddleware");
77
78        #[cfg(any(feature = "cookie-signed", feature = "cookie-private"))]
79        d.field("key", &[..]);
80
81        d.finish()
82    }
83}
84
85#[crate::async_trait]
86impl<H, O> Handler<Request> for CookieMiddleware<H>
87where
88    H: Handler<Request, Output = Result<O>>,
89    O: IntoResponse,
90{
91    type Output = Result<Response>;
92
93    async fn call(&self, mut req: Request) -> Self::Output {
94        let jar = req
95            .headers()
96            .get_all(COOKIE)
97            .iter()
98            .map(HeaderValue::to_str)
99            .filter_map(Result::ok)
100            .fold(CookieJar::new(), add_cookie);
101
102        let cookies = Cookies::new(jar);
103        #[cfg(any(feature = "cookie-signed", feature = "cookie-private"))]
104        let cookies = cookies.with_key(self.key.clone());
105
106        req.extensions_mut().insert::<Cookies>(cookies.clone());
107
108        self.h
109            .call(req)
110            .await
111            .map(IntoResponse::into_response)
112            .map(|mut res| {
113                if let Ok(c) = cookies.jar().lock() {
114                    c.delta()
115                        .map(Cookie::encoded)
116                        .map(|cookie| HeaderValue::from_str(&cookie.to_string()))
117                        .filter_map(Result::ok)
118                        .fold(res.headers_mut(), |headers, cookie| {
119                            headers.append(SET_COOKIE, cookie);
120                            headers
121                        });
122                }
123                res
124            })
125    }
126}
127
128#[inline]
129fn add_cookie(mut jar: CookieJar, value: &str) -> CookieJar {
130    Cookie::split_parse_encoded(value)
131        .filter_map(Result::ok)
132        .map(Cookie::into_owned)
133        .for_each(|cookie| jar.add_original(cookie));
134    jar
135}
viz_core/middleware/cookie.rs

viz_core/middleware/
cookie.rs
