slim_config/websocket/
query_token_layer.rs

1// Copyright AGNTCY Contributors (https://github.com/agntcy)
2// SPDX-License-Identifier: Apache-2.0
3
4//! Promotes a `?token=<jwt>` query parameter to an `Authorization: Bearer …`
5//! header so the shared `ValidateJwtLayer` (which only inspects the
6//! Authorization header) can authenticate browser WebSocket clients —
7//! browsers cannot set custom headers when constructing `new WebSocket(url)`.
8//!
9//! If the request already carries an Authorization header, it wins and the
10//! query parameter is ignored. If no token is present anywhere, the request
11//! passes through unchanged and the downstream JWT layer will reject it.
12
13use std::task::{Context, Poll};
14
15use http::{Request, header};
16use tower::{Layer, Service};
17
18use crate::websocket::common::extract_query_param;
19
20/// Rebuild `uri` with the `token` query parameter removed.
21fn strip_token_param(uri: &http::Uri) -> http::Uri {
22    let query = match uri.query() {
23        Some(q) => q,
24        None => return uri.clone(),
25    };
26    let filtered: String = query
27        .split('&')
28        .filter(|seg| {
29            let key = seg.split_once('=').map(|(k, _)| k).unwrap_or(seg);
30            key != "token"
31        })
32        .collect::<Vec<_>>()
33        .join("&");
34    let path = uri.path();
35    let new_pq: http::uri::PathAndQuery = if filtered.is_empty() {
36        path.parse().unwrap_or_else(|_| "/".parse().unwrap())
37    } else {
38        format!("{}?{}", path, filtered)
39            .parse()
40            .unwrap_or_else(|_| path.parse().unwrap())
41    };
42    let mut parts = uri.clone().into_parts();
43    parts.path_and_query = Some(new_pq);
44    http::Uri::from_parts(parts).unwrap_or_else(|_| uri.clone())
45}
46
47#[derive(Clone, Default)]
48pub struct QueryTokenToAuthHeaderLayer;
49
50impl QueryTokenToAuthHeaderLayer {
51    pub fn new() -> Self {
52        Self
53    }
54}
55
56impl<S> Layer<S> for QueryTokenToAuthHeaderLayer {
57    type Service = QueryTokenToAuthHeader<S>;
58
59    fn layer(&self, inner: S) -> Self::Service {
60        QueryTokenToAuthHeader { inner }
61    }
62}
63
64#[derive(Clone)]
65pub struct QueryTokenToAuthHeader<S> {
66    inner: S,
67}
68
69impl<S, B> Service<Request<B>> for QueryTokenToAuthHeader<S>
70where
71    S: Service<Request<B>>,
72{
73    type Response = S::Response;
74    type Error = S::Error;
75    type Future = S::Future;
76
77    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
78        self.inner.poll_ready(cx)
79    }
80
81    fn call(&mut self, mut req: Request<B>) -> Self::Future {
82        if !req.headers().contains_key(header::AUTHORIZATION)
83            && let Some(token) = extract_query_param(req.uri().query(), "token")
84            && let Ok(value) = header::HeaderValue::from_str(&format!("Bearer {token}"))
85        {
86            req.headers_mut().insert(header::AUTHORIZATION, value);
87            *req.uri_mut() = strip_token_param(req.uri());
88        }
89        self.inner.call(req)
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96
97    use std::convert::Infallible;
98
99    use http::{Request, Response, StatusCode};
100    use tower::{ServiceBuilder, ServiceExt, service_fn};
101
102    async fn echo_auth(req: Request<()>) -> Result<Response<String>, Infallible> {
103        let auth = req
104            .headers()
105            .get(header::AUTHORIZATION)
106            .and_then(|v| v.to_str().ok())
107            .unwrap_or("")
108            .to_string();
109        // Return "<auth>|<uri>" so tests can assert both.
110        let body = format!("{}|{}", auth, req.uri());
111        Ok(Response::builder()
112            .status(StatusCode::OK)
113            .body(body)
114            .unwrap())
115    }
116
117    async fn call_with(req: Request<()>) -> (String, String) {
118        let mut s = ServiceBuilder::new()
119            .layer(QueryTokenToAuthHeaderLayer::new())
120            .service(service_fn(echo_auth));
121        let body = s
122            .ready()
123            .await
124            .unwrap()
125            .call(req)
126            .await
127            .unwrap()
128            .into_body();
129        let (auth, uri) = body.split_once('|').unwrap_or(("", &body));
130        (auth.to_string(), uri.to_string())
131    }
132
133    #[tokio::test]
134    async fn no_header_no_query_passes_through() {
135        let req = Request::builder().uri("/").body(()).unwrap();
136        let (auth, uri) = call_with(req).await;
137        assert_eq!(auth, "");
138        assert_eq!(uri, "/");
139    }
140
141    #[tokio::test]
142    async fn query_token_promoted_to_header_and_stripped_from_uri() {
143        let req = Request::builder().uri("/?token=abc").body(()).unwrap();
144        let (auth, uri) = call_with(req).await;
145        assert_eq!(auth, "Bearer abc");
146        assert_eq!(uri, "/");
147    }
148
149    #[tokio::test]
150    async fn token_stripped_other_params_preserved() {
151        let req = Request::builder()
152            .uri("/?foo=1&token=abc&bar=2")
153            .body(())
154            .unwrap();
155        let (auth, uri) = call_with(req).await;
156        assert_eq!(auth, "Bearer abc");
157        assert_eq!(uri, "/?foo=1&bar=2");
158    }
159
160    #[tokio::test]
161    async fn existing_header_wins_over_query() {
162        let req = Request::builder()
163            .uri("/?token=fromquery")
164            .header(header::AUTHORIZATION, "Bearer fromheader")
165            .body(())
166            .unwrap();
167        let (auth, uri) = call_with(req).await;
168        assert_eq!(auth, "Bearer fromheader");
169        // URI unchanged when header already present
170        assert_eq!(uri, "/?token=fromquery");
171    }
172
173    #[tokio::test]
174    async fn percent_encoded_token_is_decoded() {
175        let req = Request::builder().uri("/?token=ab%20c").body(()).unwrap();
176        let (auth, uri) = call_with(req).await;
177        assert_eq!(auth, "Bearer ab c");
178        assert_eq!(uri, "/");
179    }
180}
slim_config/websocket/query_token_layer.rs

slim_config/websocket/
query_token_layer.rs
