systemprompt_security/extraction/
cookie.rs

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

systemprompt_security/extraction/
cookie.rs
