systemprompt_security/extraction/
cookie.rs

1//! Cookie-based credential extraction.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use axum::http::HeaderMap;
7
8#[derive(Debug, Clone)]
9pub struct CookieExtractor {
10    cookie_name: String,
11}
12
13impl Default for CookieExtractor {
14    fn default() -> Self {
15        Self::new(Self::DEFAULT_COOKIE_NAME)
16    }
17}
18
19impl CookieExtractor {
20    pub const DEFAULT_COOKIE_NAME: &'static str = "access_token";
21
22    pub fn new(cookie_name: impl Into<String>) -> Self {
23        Self {
24            cookie_name: cookie_name.into(),
25        }
26    }
27
28    pub fn extract(&self, headers: &HeaderMap) -> Result<String, CookieExtractionError> {
29        self.extract_internal(headers)
30    }
31
32    pub fn extract_access_token(headers: &HeaderMap) -> Result<String, CookieExtractionError> {
33        Self::default().extract(headers)
34    }
35
36    fn extract_internal(&self, headers: &HeaderMap) -> Result<String, CookieExtractionError> {
37        let cookie_header = headers
38            .get("cookie")
39            .ok_or(CookieExtractionError::MissingCookie)?
40            .to_str()
41            .map_err(|_e| CookieExtractionError::InvalidCookieFormat)?;
42
43        for cookie in cookie_header.split(';') {
44            let cookie = cookie.trim();
45            let cookie_prefix = format!("{}=", self.cookie_name);
46            if let Some(value) = cookie.strip_prefix(&cookie_prefix)
47                && !value.is_empty()
48            {
49                return Ok(value.to_owned());
50            }
51        }
52
53        Err(CookieExtractionError::TokenNotFoundInCookie)
54    }
55}
56
57#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub enum CookieExtractionError {
59    MissingCookie,
60    InvalidCookieFormat,
61    TokenNotFoundInCookie,
62}
63
64impl std::fmt::Display for CookieExtractionError {
65    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
66        match self {
67            Self::MissingCookie => write!(f, "Missing cookie header"),
68            Self::InvalidCookieFormat => write!(f, "Invalid cookie format"),
69            Self::TokenNotFoundInCookie => {
70                write!(f, "Access token not found in cookies")
71            },
72        }
73    }
74}
75
76impl std::error::Error for CookieExtractionError {}
systemprompt_security/extraction/cookie.rs

systemprompt_security/extraction/
cookie.rs
