Skip to main content

trillium_cookies/
cookies_handler.rs

1use cookie::{Cookie, CookieJar};
2use trillium::{Conn, Handler, HeaderValue, HeaderValues, KnownHeaderName};
3
4/// The trillium cookie handler. See crate level docs for an example. This
5/// must run before any handlers access the cookie jar.
6#[derive(Clone, Copy, Debug, Default)]
7pub struct CookiesHandler {
8    // this is in order to force users to call CookiesHandler::new or
9    // CookiesHandler::default, allowing us to add
10    // customization/settings later without breaking existing usage
11    _priv: (),
12}
13
14impl CookiesHandler {
15    /// constructs a new cookies handler
16    pub fn new() -> Self {
17        Self::default()
18    }
19}
20
21impl Handler for CookiesHandler {
22    async fn run(&self, mut conn: Conn) -> Conn {
23        let mut jar: CookieJar = conn.take_state().unwrap_or_default();
24
25        if let Some(cookies) = conn.request_headers().get_values(KnownHeaderName::Cookie) {
26            for cookie in cookies.iter().filter_map(HeaderValue::as_str) {
27                for pair in cookie.split(';') {
28                    if let Ok(cookie) = Cookie::parse_encoded(String::from(pair)) {
29                        jar.add_original(cookie);
30                    }
31                }
32            }
33        }
34
35        conn.with_state(jar)
36    }
37
38    async fn before_send(&self, mut conn: Conn) -> Conn {
39        if let Some(jar) = conn.take_state::<CookieJar>() {
40            conn.response_headers_mut().append(
41                KnownHeaderName::SetCookie,
42                jar.delta()
43                    .map(|cookie| cookie.encoded().to_string())
44                    .collect::<HeaderValues>(),
45            );
46            conn.with_state(jar)
47        } else {
48            conn
49        }
50    }
51}
52
53/// Alias for CookiesHandler::new()
54pub fn cookies() -> CookiesHandler {
55    CookiesHandler::new()
56}