Skip to main content

mas_axum_utils/
cookies.rs

1// Copyright 2025, 2026 Element Creations Ltd.
2// Copyright 2024, 2025 New Vector Ltd.
3// Copyright 2022-2024 The Matrix.org Foundation C.I.C.
4//
5// SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
6// Please see LICENSE files in the repository root for full details.
7
8//! Private (encrypted) cookie jar, based on axum-extra's cookie jar
9
10use std::{convert::Infallible, future::ready};
11
12use axum::{
13    extract::{FromRef, FromRequestParts},
14    response::{IntoResponseParts, ResponseParts},
15};
16use axum_extra::extract::cookie::{Cookie, Key, PrivateCookieJar, SameSite};
17use http::request::Parts;
18use serde::{Serialize, de::DeserializeOwned};
19use thiserror::Error;
20use url::Url;
21
22#[derive(Debug, Error)]
23#[error("could not decode cookie")]
24pub enum CookieDecodeError {
25    Deserialize(#[from] serde_json::Error),
26}
27
28/// Manages cookie options and encryption key
29///
30/// This is meant to be accessible through axum's state via the [`FromRef`]
31/// trait
32#[derive(Clone)]
33pub struct CookieManager {
34    options: CookieOption,
35    key: Key,
36}
37
38impl CookieManager {
39    #[must_use]
40    pub const fn new(base_url: Url, key: Key) -> Self {
41        let options = CookieOption::new(base_url);
42        Self { options, key }
43    }
44
45    #[must_use]
46    pub fn derive_from(base_url: Url, key: &[u8]) -> Self {
47        let key = Key::derive_from(key);
48        Self::new(base_url, key)
49    }
50
51    #[must_use]
52    pub fn cookie_jar(&self) -> CookieJar {
53        let inner = PrivateCookieJar::new(self.key.clone());
54        let options = self.options.clone();
55
56        CookieJar { inner, options }
57    }
58
59    #[must_use]
60    pub fn cookie_jar_from_headers(&self, headers: &http::HeaderMap) -> CookieJar {
61        let inner = PrivateCookieJar::from_headers(headers, self.key.clone());
62        let options = self.options.clone();
63
64        CookieJar { inner, options }
65    }
66}
67
68impl<S> FromRequestParts<S> for CookieJar
69where
70    CookieManager: FromRef<S>,
71    S: Send + Sync,
72{
73    type Rejection = Infallible;
74
75    fn from_request_parts(
76        parts: &mut Parts,
77        state: &S,
78    ) -> impl Future<Output = Result<Self, Self::Rejection>> + Send {
79        let cookie_manager = CookieManager::from_ref(state);
80        ready(Ok(cookie_manager.cookie_jar_from_headers(&parts.headers)))
81    }
82}
83
84#[derive(Debug, Clone)]
85struct CookieOption {
86    base_url: Url,
87}
88
89impl CookieOption {
90    const fn new(base_url: Url) -> Self {
91        Self { base_url }
92    }
93
94    fn secure(&self) -> bool {
95        self.base_url.scheme() == "https"
96    }
97
98    fn path(&self) -> &str {
99        self.base_url.path()
100    }
101
102    fn apply<'a>(&self, mut cookie: Cookie<'a>) -> Cookie<'a> {
103        cookie.set_http_only(true);
104        cookie.set_secure(self.secure());
105        cookie.set_path(self.path().to_owned());
106        cookie.set_same_site(SameSite::Lax);
107        cookie
108    }
109}
110
111/// A cookie jar which encrypts cookies & sets secure options
112pub struct CookieJar {
113    inner: PrivateCookieJar<Key>,
114    options: CookieOption,
115}
116
117impl CookieJar {
118    /// Save the given payload in a cookie
119    ///
120    /// If `permanent` is true, the cookie will be valid for 10 years
121    ///
122    /// # Panics
123    ///
124    /// Panics if the payload cannot be serialized
125    #[must_use]
126    pub fn save<T: Serialize>(mut self, key: &str, payload: &T, permanent: bool) -> Self {
127        let serialized =
128            serde_json::to_string(payload).expect("failed to serialize cookie payload");
129
130        let cookie = Cookie::new(key.to_owned(), serialized);
131        let mut cookie = self.options.apply(cookie);
132
133        if permanent {
134            // XXX: this should use a clock
135            cookie.make_permanent();
136        }
137
138        self.inner = self.inner.add(cookie);
139
140        self
141    }
142
143    /// Remove a cookie from the jar
144    #[must_use]
145    pub fn remove(mut self, key: &str) -> Self {
146        self.inner = self.inner.remove(key.to_owned());
147        self
148    }
149
150    /// Load and deserialize a cookie from the jar
151    ///
152    /// Returns `None` if the cookie is not present
153    ///
154    /// # Errors
155    ///
156    /// Returns an error if the cookie cannot be deserialized
157    pub fn load<T: DeserializeOwned>(&self, key: &str) -> Result<Option<T>, CookieDecodeError> {
158        let Some(cookie) = self.inner.get(key) else {
159            return Ok(None);
160        };
161
162        let decoded = serde_json::from_str(cookie.value())?;
163        Ok(Some(decoded))
164    }
165}
166
167impl IntoResponseParts for CookieJar {
168    type Error = Infallible;
169
170    fn into_response_parts(self, res: ResponseParts) -> Result<ResponseParts, Self::Error> {
171        self.inner.into_response_parts(res)
172    }
173}