ws_kit/
extractor.rs

1//! Token extraction for WebSocket authentication.
2//!
3//! Extracts bearer tokens from HTTP upgrade requests using a configurable set
4//! of [`TokenSourceKind`] strategies. Supports `Authorization: Bearer <token>`,
5//! query parameters (`?token=` / `?access_token=`), and `Cookie` headers.
6
7use thiserror::Error;
8
9#[cfg(feature = "axum")]
10use http::request::Parts;
11
12/// Where to look for a token.
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub enum TokenSourceKind {
15    /// `Authorization: Bearer <token>` header.
16    AuthorizationHeader,
17    /// Query parameter with the given name (e.g. `token`, `access_token`).
18    QueryParam(String),
19    /// Cookie with the given name.
20    Cookie(String),
21}
22
23/// Error during authentication extraction/validation.
24#[derive(Debug, Error, Clone, PartialEq, Eq)]
25pub enum WsAuthError {
26    /// No token found in any configured source.
27    #[error("authentication token missing")]
28    Missing,
29
30    /// Token was found but is invalid.
31    #[error("authentication token invalid")]
32    Invalid,
33
34    /// Token is expired (if expiry is checked downstream).
35    #[error("authentication token expired")]
36    Expired,
37}
38
39/// Extracts tokens from HTTP request parts according to configured sources.
40#[derive(Debug, Clone)]
41pub struct TokenExtractor {
42    sources: Vec<TokenSourceKind>,
43}
44
45impl Default for TokenExtractor {
46    fn default() -> Self {
47        Self {
48            sources: vec![
49                TokenSourceKind::AuthorizationHeader,
50                TokenSourceKind::QueryParam("token".to_string()),
51                TokenSourceKind::QueryParam("access_token".to_string()),
52            ],
53        }
54    }
55}
56
57impl TokenExtractor {
58    /// Create a new extractor with the given ordered sources.
59    pub fn new(sources: Vec<TokenSourceKind>) -> Self {
60        Self { sources }
61    }
62
63    /// Bearer-only extractor.
64    pub fn bearer_only() -> Self {
65        Self {
66            sources: vec![TokenSourceKind::AuthorizationHeader],
67        }
68    }
69
70    /// Add a source.
71    pub fn with_source(mut self, src: TokenSourceKind) -> Self {
72        self.sources.push(src);
73        self
74    }
75
76    /// Ordered sources.
77    pub fn sources(&self) -> &[TokenSourceKind] {
78        &self.sources
79    }
80
81    /// Extract the first token found according to source order.
82    ///
83    /// When the `axum` feature is enabled the first argument is `&http::request::Parts`
84    /// (from `axum`'s request). The second argument is the raw query string
85    /// (e.g. `"token=abc&foo=bar"` without leading `?`).
86    ///
87    /// Returns `None` if no token was found in any source.
88    #[cfg(feature = "axum")]
89    pub fn extract_token(&self, parts: &Parts, query: &str) -> Option<String> {
90        for src in &self.sources {
91            match src {
92                TokenSourceKind::AuthorizationHeader => {
93                    if let Some(v) = Self::extract_bearer(parts) {
94                        return Some(v);
95                    }
96                }
97                TokenSourceKind::QueryParam(name) => {
98                    if let Some(v) = Self::extract_query_param(query, name) {
99                        return Some(v);
100                    }
101                }
102                TokenSourceKind::Cookie(name) => {
103                    if let Some(v) = Self::extract_cookie(parts, name) {
104                        return Some(v);
105                    }
106                }
107            }
108        }
109        None
110    }
111
112    /// Extract token without `Parts` (query and cookie header strings only).
113    ///
114    /// Useful when `axum` is not enabled or for tests.
115    pub fn extract_from_parts(
116        &self,
117        auth_header: Option<&str>,
118        cookie_header: Option<&str>,
119        query: &str,
120    ) -> Option<String> {
121        for src in &self.sources {
122            match src {
123                TokenSourceKind::AuthorizationHeader => {
124                    if let Some(v) = auth_header.and_then(Self::parse_bearer) {
125                        if !v.is_empty() {
126                            return Some(v);
127                        }
128                    }
129                }
130                TokenSourceKind::QueryParam(name) => {
131                    if let Some(v) = Self::extract_query_param(query, name) {
132                        return Some(v);
133                    }
134                }
135                TokenSourceKind::Cookie(name) => {
136                    if let Some(v) = cookie_header.and_then(|c| Self::parse_cookie(c, name)) {
137                        return Some(v);
138                    }
139                }
140            }
141        }
142        None
143    }
144
145    /// Validate presence of token, returning a typed error.
146    #[cfg(feature = "axum")]
147    pub fn require_token(&self, parts: &Parts, query: &str) -> Result<String, WsAuthError> {
148        self.extract_token(parts, query).ok_or(WsAuthError::Missing)
149    }
150
151    #[cfg(feature = "axum")]
152    fn extract_bearer(parts: &Parts) -> Option<String> {
153        let val = parts.headers.get(http::header::AUTHORIZATION)?;
154        let s = val.to_str().ok()?;
155        Self::parse_bearer(s)
156    }
157
158    fn parse_bearer(s: &str) -> Option<String> {
159        // Case-insensitive "Bearer " prefix
160        if s.len() < 7 {
161            return None;
162        }
163        let prefix = &s[..7];
164        if !prefix.eq_ignore_ascii_case("bearer ") {
165            return None;
166        }
167        let token = s[7..].trim();
168        if token.is_empty() {
169            None
170        } else {
171            Some(token.to_string())
172        }
173    }
174
175    fn extract_query_param(query: &str, name: &str) -> Option<String> {
176        let query = query.trim_start_matches('?');
177        if query.is_empty() {
178            return None;
179        }
180        for pair in query.split('&') {
181            if pair.is_empty() {
182                continue;
183            }
184            let mut split = pair.splitn(2, '=');
185            let k = split.next().unwrap_or("");
186            let v = split.next().unwrap_or("");
187            // URL-decode k and v minimally: '+' -> ' ', %XX
188            let k_decoded = url_decode(k);
189            if k_decoded == name {
190                let v_decoded = url_decode(v);
191                if !v_decoded.is_empty() {
192                    return Some(v_decoded);
193                }
194            }
195        }
196        None
197    }
198
199    #[cfg(feature = "axum")]
200    fn extract_cookie(parts: &Parts, name: &str) -> Option<String> {
201        let val = parts.headers.get(http::header::COOKIE)?;
202        let s = val.to_str().ok()?;
203        Self::parse_cookie(s, name)
204    }
205
206    fn parse_cookie(header: &str, name: &str) -> Option<String> {
207        for part in header.split(';') {
208            let part = part.trim();
209            let mut split = part.splitn(2, '=');
210            let k = split.next().unwrap_or("").trim();
211            let v = split.next().unwrap_or("").trim();
212            // Trim optional quotes
213            let v = v.trim_matches('"');
214            if k == name && !v.is_empty() {
215                return Some(url_decode(v));
216            }
217        }
218        None
219    }
220}
221
222fn url_decode(s: &str) -> String {
223    let mut out = String::with_capacity(s.len());
224    let mut chars = s.chars().peekable();
225    while let Some(c) = chars.next() {
226        match c {
227            '+' => out.push(' '),
228            '%' => {
229                let hi = chars.next();
230                let lo = chars.next();
231                if let (Some(h), Some(l)) = (hi, lo) {
232                    if let (Some(hv), Some(lv)) = (hex_val(h), hex_val(l)) {
233                        out.push((hv * 16 + lv) as char);
234                    } else {
235                        out.push('%');
236                        out.push(h);
237                        out.push(l);
238                    }
239                } else {
240                    out.push('%');
241                    if let Some(h) = hi {
242                        out.push(h);
243                    }
244                    if let Some(l) = lo {
245                        out.push(l);
246                    }
247                }
248            }
249            other => out.push(other),
250        }
251    }
252    out
253}
254
255fn hex_val(c: char) -> Option<u8> {
256    match c {
257        '0'..='9' => Some(c as u8 - b'0'),
258        'a'..='f' => Some(c as u8 - b'a' + 10),
259        'A'..='F' => Some(c as u8 - b'A' + 10),
260        _ => None,
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267
268    #[test]
269    fn parse_bearer_ok() {
270        assert_eq!(
271            TokenExtractor::parse_bearer("Bearer abc123"),
272            Some("abc123".to_string())
273        );
274        assert_eq!(
275            TokenExtractor::parse_bearer("bearer abc123"),
276            Some("abc123".to_string())
277        );
278        assert_eq!(
279            TokenExtractor::parse_bearer("BEARER abc123"),
280            Some("abc123".to_string())
281        );
282        assert_eq!(
283            TokenExtractor::parse_bearer("Bearer   spaced  "),
284            Some("spaced".to_string())
285        );
286    }
287
288    #[test]
289    fn parse_bearer_fail() {
290        assert_eq!(TokenExtractor::parse_bearer("Basic abc"), None);
291        assert_eq!(TokenExtractor::parse_bearer("Bearer "), None);
292        assert_eq!(TokenExtractor::parse_bearer("Bearer"), None);
293        assert_eq!(TokenExtractor::parse_bearer(""), None);
294    }
295
296    #[test]
297    fn query_param_extraction() {
298        assert_eq!(
299            TokenExtractor::extract_query_param("token=abc123&foo=bar", "token"),
300            Some("abc123".to_string())
301        );
302        assert_eq!(
303            TokenExtractor::extract_query_param("?access_token=xyz", "access_token"),
304            Some("xyz".to_string())
305        );
306        assert_eq!(
307            TokenExtractor::extract_query_param("foo=1&token=second", "token"),
308            Some("second".to_string())
309        );
310        assert_eq!(TokenExtractor::extract_query_param("", "token"), None);
311        assert_eq!(
312            TokenExtractor::extract_query_param("token=&foo=bar", "token"),
313            None
314        );
315    }
316
317    #[test]
318    fn query_url_decode() {
319        assert_eq!(
320            TokenExtractor::extract_query_param("token=hello%20world", "token"),
321            Some("hello world".to_string())
322        );
323        assert_eq!(
324            TokenExtractor::extract_query_param("token=a%2Bb", "token"),
325            Some("a+b".to_string())
326        );
327    }
328
329    #[test]
330    fn cookie_extraction() {
331        assert_eq!(
332            TokenExtractor::parse_cookie("session=abc; token=xyz; foo=bar", "token"),
333            Some("xyz".to_string())
334        );
335        assert_eq!(
336            TokenExtractor::parse_cookie("token=hello%20world", "token"),
337            Some("hello world".to_string())
338        );
339        assert_eq!(TokenExtractor::parse_cookie("a=b; c=d", "token"), None);
340    }
341
342    #[test]
343    fn extractor_order_header_first() {
344        let ex = TokenExtractor::default();
345        let token = ex.extract_from_parts(Some("Bearer header_tok"), None, "token=query_tok");
346        assert_eq!(token, Some("header_tok".to_string()));
347    }
348
349    #[test]
350    fn extractor_fallback_to_query() {
351        let ex = TokenExtractor::default();
352        let token = ex.extract_from_parts(None, None, "token=query_tok");
353        assert_eq!(token, Some("query_tok".to_string()));
354    }
355
356    #[test]
357    fn extractor_fallback_access_token() {
358        let ex = TokenExtractor::default();
359        let token = ex.extract_from_parts(None, None, "access_token=at");
360        assert_eq!(token, Some("at".to_string()));
361    }
362
363    #[test]
364    fn extractor_cookie_source() {
365        let ex = TokenExtractor::new(vec![TokenSourceKind::Cookie("session".to_string())]);
366        let token = ex.extract_from_parts(None, Some("session=mycookie"), "");
367        assert_eq!(token, Some("mycookie".to_string()));
368    }
369
370    #[test]
371    fn extractor_none_when_missing() {
372        let ex = TokenExtractor::default();
373        assert_eq!(ex.extract_from_parts(None, None, ""), None);
374    }
375
376    #[cfg(feature = "axum")]
377    mod axum_tests {
378        use super::*;
379        use http::{HeaderMap, Request};
380
381        fn parts_with(headers: HeaderMap, uri: &str) -> http::request::Parts {
382            let req = Request::builder().uri(uri).body(()).unwrap();
383            let (mut parts, _) = req.into_parts();
384            parts.headers = headers;
385            parts
386        }
387
388        #[test]
389        fn axum_bearer_extraction() {
390            let mut headers = HeaderMap::new();
391            headers.insert(
392                http::header::AUTHORIZATION,
393                "Bearer axum_tok".parse().unwrap(),
394            );
395            let parts = parts_with(headers, "/ws?token=ignored");
396            let ex = TokenExtractor::default();
397            assert_eq!(
398                ex.extract_token(&parts, "token=ignored"),
399                Some("axum_tok".to_string())
400            );
401        }
402
403        #[test]
404        fn axum_query_fallback() {
405            let parts = parts_with(HeaderMap::new(), "/ws?token=qtoken");
406            let ex = TokenExtractor::default();
407            assert_eq!(
408                ex.extract_token(&parts, "token=qtoken"),
409                Some("qtoken".to_string())
410            );
411        }
412
413        #[test]
414        fn axum_cookie() {
415            let mut headers = HeaderMap::new();
416            headers.insert(http::header::COOKIE, "tok=cookie_val".parse().unwrap());
417            let parts = parts_with(headers, "/ws");
418            let ex = TokenExtractor::new(vec![TokenSourceKind::Cookie("tok".to_string())]);
419            assert_eq!(ex.extract_token(&parts, ""), Some("cookie_val".to_string()));
420        }
421    }
422}
ws_kit/extractor.rs

ws_kit/
extractor.rs
