systemprompt_security/extraction/
token.rs

1//! Bearer/proxy-header token extraction from requests.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use axum::http::HeaderMap;
7use std::error::Error;
8use std::fmt;
9
10use super::cookie::{CookieExtractionError, CookieExtractor};
11
12const DEFAULT_MCP_HEADER_NAME: &str = "x-mcp-proxy-auth";
13const BEARER_PREFIX: &str = "Bearer ";
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum ExtractionMethod {
17    AuthorizationHeader,
18    McpProxyHeader,
19    Cookie,
20}
21
22impl fmt::Display for ExtractionMethod {
23    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
24        match self {
25            Self::AuthorizationHeader => write!(f, "Authorization header"),
26            Self::McpProxyHeader => write!(f, "MCP proxy header"),
27            Self::Cookie => write!(f, "Cookie"),
28        }
29    }
30}
31
32#[derive(Debug, Clone)]
33pub struct TokenExtractor {
34    fallback_chain: Vec<ExtractionMethod>,
35    cookie_name: String,
36    mcp_header_name: String,
37}
38
39impl TokenExtractor {
40    #[must_use]
41    pub fn new(fallback_chain: Vec<ExtractionMethod>) -> Self {
42        Self {
43            fallback_chain,
44            cookie_name: CookieExtractor::DEFAULT_COOKIE_NAME.to_owned(),
45            mcp_header_name: DEFAULT_MCP_HEADER_NAME.to_owned(),
46        }
47    }
48
49    #[must_use]
50    pub fn with_cookie_name(mut self, name: String) -> Self {
51        self.cookie_name = name;
52        self
53    }
54
55    #[must_use]
56    pub fn with_mcp_header_name(mut self, name: String) -> Self {
57        self.mcp_header_name = name;
58        self
59    }
60
61    #[must_use]
62    pub fn standard() -> Self {
63        Self::new(vec![
64            ExtractionMethod::AuthorizationHeader,
65            ExtractionMethod::McpProxyHeader,
66            ExtractionMethod::Cookie,
67        ])
68    }
69
70    #[must_use]
71    pub fn browser_only() -> Self {
72        Self::new(vec![
73            ExtractionMethod::AuthorizationHeader,
74            ExtractionMethod::Cookie,
75        ])
76    }
77
78    #[must_use]
79    pub fn api_only() -> Self {
80        Self::new(vec![ExtractionMethod::AuthorizationHeader])
81    }
82
83    #[must_use]
84    pub fn chain(&self) -> &[ExtractionMethod] {
85        &self.fallback_chain
86    }
87
88    pub fn extract(&self, headers: &HeaderMap) -> Result<String, TokenExtractionError> {
89        for method in &self.fallback_chain {
90            match method {
91                ExtractionMethod::AuthorizationHeader => {
92                    if let Ok(token) = Self::extract_from_authorization(headers) {
93                        return Ok(token);
94                    }
95                },
96                ExtractionMethod::McpProxyHeader => {
97                    if let Ok(token) = self.extract_from_mcp_proxy(headers) {
98                        return Ok(token);
99                    }
100                },
101                ExtractionMethod::Cookie => {
102                    if let Ok(token) = self.extract_from_cookie(headers) {
103                        return Ok(token);
104                    }
105                },
106            }
107        }
108
109        Err(TokenExtractionError::NoTokenFound)
110    }
111
112    pub fn extract_from_authorization(headers: &HeaderMap) -> Result<String, TokenExtractionError> {
113        let auth_headers = headers.get_all("authorization");
114
115        if auth_headers.iter().count() == 0 {
116            return Err(TokenExtractionError::MissingAuthorizationHeader);
117        }
118
119        for auth_value in &auth_headers {
120            let Ok(auth_header) = auth_value.to_str().map_err(|e| {
121                tracing::debug!(error = %e, "Authorization header contains non-ASCII characters");
122                e
123            }) else {
124                continue;
125            };
126
127            if let Some(token) = auth_header.strip_prefix(BEARER_PREFIX) {
128                let token = token.trim();
129                if !token.is_empty() {
130                    return Ok(token.to_owned());
131                }
132            }
133        }
134
135        Err(TokenExtractionError::InvalidAuthorizationFormat)
136    }
137
138    pub fn extract_from_mcp_proxy(
139        &self,
140        headers: &HeaderMap,
141    ) -> Result<String, TokenExtractionError> {
142        let header_value = headers
143            .get(&self.mcp_header_name)
144            .ok_or(TokenExtractionError::MissingMcpProxyHeader)?;
145
146        let auth_header = header_value
147            .to_str()
148            .map_err(|_e| TokenExtractionError::InvalidMcpProxyFormat)?;
149
150        auth_header
151            .strip_prefix(BEARER_PREFIX)
152            .ok_or(TokenExtractionError::InvalidMcpProxyFormat)
153            .map(str::to_owned)
154    }
155
156    pub fn extract_from_cookie(&self, headers: &HeaderMap) -> Result<String, TokenExtractionError> {
157        CookieExtractor::new(self.cookie_name.clone())
158            .extract(headers)
159            .map_err(|err| match err {
160                CookieExtractionError::MissingCookie => TokenExtractionError::MissingCookie,
161                CookieExtractionError::InvalidCookieFormat => {
162                    TokenExtractionError::InvalidCookieFormat
163                },
164                CookieExtractionError::TokenNotFoundInCookie => {
165                    TokenExtractionError::TokenNotFoundInCookie
166                },
167            })
168    }
169}
170
171#[derive(Debug, Clone, Copy, PartialEq, Eq)]
172pub enum TokenExtractionError {
173    NoTokenFound,
174    MissingAuthorizationHeader,
175    InvalidAuthorizationFormat,
176    MissingMcpProxyHeader,
177    InvalidMcpProxyFormat,
178    MissingCookie,
179    InvalidCookieFormat,
180    TokenNotFoundInCookie,
181}
182
183impl fmt::Display for TokenExtractionError {
184    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
185        match self {
186            Self::NoTokenFound => write!(f, "No token found in request"),
187            Self::MissingAuthorizationHeader => {
188                write!(f, "Missing Authorization header")
189            },
190            Self::InvalidAuthorizationFormat => {
191                write!(
192                    f,
193                    "Invalid Authorization header format (expected 'Bearer <token>')"
194                )
195            },
196            Self::MissingMcpProxyHeader => {
197                write!(f, "Missing MCP proxy authorization header")
198            },
199            Self::InvalidMcpProxyFormat => {
200                write!(
201                    f,
202                    "Invalid MCP proxy header format (expected 'Bearer <token>')"
203                )
204            },
205            Self::MissingCookie => write!(f, "Missing cookie header"),
206            Self::InvalidCookieFormat => write!(f, "Invalid cookie format"),
207            Self::TokenNotFoundInCookie => write!(f, "Token not found in cookies"),
208        }
209    }
210}
211
212impl Error for TokenExtractionError {}
systemprompt_security/extraction/token.rs

systemprompt_security/extraction/
token.rs
