Skip to main content

trillium_cookies/
cookies_handler.rs

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