koan_server/auth/
middleware.rs

1//! Axum middleware for JWT authentication.
2//!
3//! Extracts a token from the `koan_access` cookie or `Authorization: Bearer`,
4//! validates it, and injects `AuthUser` into request extensions.
5
6use std::sync::Arc;
7
8use axum::extract::Request;
9use axum::http::{StatusCode, header};
10use axum::middleware::Next;
11use axum::response::{IntoResponse, Response};
12use subtle::ConstantTimeEq;
13
14use koan_core::auth::{self, Role};
15
16use super::AuthUser;
17
18/// Shared state for the auth middleware.
19#[derive(Clone)]
20pub struct AuthState {
21    /// Ed25519 public key PEM for JWT verification.
22    pub public_pem: Arc<Vec<u8>>,
23    /// Whether auth is enforced.
24    pub auth_enabled: bool,
25    /// Process-scoped introspection key. Bypasses auth when matched.
26    /// Generated randomly on server start, dies with the process.
27    pub introspection_key: Option<Arc<String>>,
28}
29
30/// Axum middleware: validate JWT and inject `AuthUser`.
31///
32/// When `auth_enabled = false`, injects anonymous admin and passes through.
33/// When `auth_enabled = true`, requires a valid token.
34pub async fn auth_middleware(
35    axum::extract::State(state): axum::extract::State<AuthState>,
36    mut request: Request,
37    next: Next,
38) -> Response {
39    if !state.auth_enabled {
40        request.extensions_mut().insert(AuthUser::anonymous_admin());
41        return next.run(request).await;
42    }
43
44    // Check for introspection key (playground bypass).
45    if let Some(ref expected_key) = state.introspection_key
46        && let Some(provided) = request
47            .headers()
48            .get("X-Introspection-Key")
49            .and_then(|v| v.to_str().ok())
50        && provided
51            .as_bytes()
52            .ct_eq(expected_key.as_bytes())
53            .unwrap_u8()
54            == 1
55    {
56        request.extensions_mut().insert(AuthUser::anonymous_admin());
57        return next.run(request).await;
58    }
59
60    let Some(token) = extract_token(&request) else {
61        return (
62            StatusCode::UNAUTHORIZED,
63            [("WWW-Authenticate", "Bearer")],
64            "missing or invalid Authorization header",
65        )
66            .into_response();
67    };
68
69    match auth::validate_access_token(&state.public_pem, &token) {
70        Ok(claims) => {
71            let role = claims.role.parse().unwrap_or(Role::Readonly);
72            let user = AuthUser {
73                user_id: claims.sub,
74                username: claims.username,
75                role,
76            };
77            request.extensions_mut().insert(user);
78            next.run(request).await
79        }
80        Err(_) => (
81            StatusCode::UNAUTHORIZED,
82            [("WWW-Authenticate", "Bearer")],
83            "invalid or expired token",
84        )
85            .into_response(),
86    }
87}
88
89/// Priority: `koan_access` cookie, then `Authorization: Bearer`, then `?token=`.
90///
91/// The query parameter is confined to the WebSocket route, which is the only one
92/// that cannot carry a header. A token in a URL survives in shell history, proxy
93/// logs and `Referer`.
94fn extract_token(request: &Request) -> Option<String> {
95    request
96        .headers()
97        .get(header::COOKIE)
98        .and_then(|v| v.to_str().ok())
99        .and_then(|cookies| {
100            cookies
101                .split(';')
102                .find_map(|c| c.trim().strip_prefix("koan_access=").map(String::from))
103        })
104        .or_else(|| {
105            request
106                .headers()
107                .get(header::AUTHORIZATION)
108                .and_then(|v| v.to_str().ok())
109                .and_then(|v| v.strip_prefix("Bearer "))
110                .map(String::from)
111        })
112        .or_else(|| {
113            if request.uri().path() != "/graphql/ws" {
114                return None;
115            }
116            request.uri().query().and_then(|q| {
117                q.split('&')
118                    .find_map(|pair| pair.strip_prefix("token=").map(String::from))
119            })
120        })
121}
122
123// ---------------------------------------------------------------------------
124// Tests
125// ---------------------------------------------------------------------------
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130    use axum::body::Body;
131    use axum::http::Request as HttpRequest;
132    use axum::routing::get;
133    use tower::ServiceExt as _;
134
135    /// Echoes the `AuthUser` the middleware injected, so tests can assert on it.
136    async fn echo_user(axum::Extension(user): axum::Extension<AuthUser>) -> String {
137        format!("{}:{}", user.username, user.role.as_str())
138    }
139
140    async fn call(state: AuthState, req: HttpRequest<Body>) -> (StatusCode, String) {
141        let app = axum::Router::new()
142            .route("/graphql", get(echo_user))
143            .route("/graphql/ws", get(echo_user))
144            .layer(axum::middleware::from_fn_with_state(state, auth_middleware));
145        let resp = app.oneshot(req).await.unwrap();
146        let status = resp.status();
147        let bytes = axum::body::to_bytes(resp.into_body(), 64 * 1024)
148            .await
149            .unwrap();
150        (status, String::from_utf8_lossy(&bytes).into_owned())
151    }
152
153    /// A live keypair plus a matching token for `role`.
154    fn keys_and_token(role: Role) -> (Vec<u8>, String) {
155        let (private_pem, public_pem) = auth::generate_keypair_pem().unwrap();
156        let token = auth::mint_access_token(private_pem.as_bytes(), 7, "alice", role, 900).unwrap();
157        (public_pem.into_bytes(), token)
158    }
159
160    fn enforcing(public_pem: Vec<u8>, key: Option<&str>) -> AuthState {
161        AuthState {
162            public_pem: Arc::new(public_pem),
163            auth_enabled: true,
164            introspection_key: key.map(|k| Arc::new(k.to_string())),
165        }
166    }
167
168    #[tokio::test]
169    async fn auth_disabled_grants_anonymous_admin() {
170        let state = AuthState {
171            public_pem: Arc::new(Vec::new()),
172            auth_enabled: false,
173            introspection_key: None,
174        };
175        let req = HttpRequest::get("/graphql").body(Body::empty()).unwrap();
176        let (status, body) = call(state, req).await;
177        assert_eq!(status, StatusCode::OK);
178        assert_eq!(body, "anonymous:admin");
179    }
180
181    #[tokio::test]
182    async fn missing_token_is_unauthorized() {
183        let (public_pem, _) = keys_and_token(Role::Admin);
184        let req = HttpRequest::get("/graphql").body(Body::empty()).unwrap();
185        let (status, _) = call(enforcing(public_pem, None), req).await;
186        assert_eq!(status, StatusCode::UNAUTHORIZED);
187    }
188
189    #[tokio::test]
190    async fn bearer_token_authenticates() {
191        let (public_pem, token) = keys_and_token(Role::User);
192        let req = HttpRequest::get("/graphql")
193            .header(header::AUTHORIZATION, format!("Bearer {token}"))
194            .body(Body::empty())
195            .unwrap();
196        let (status, body) = call(enforcing(public_pem, None), req).await;
197        assert_eq!(status, StatusCode::OK);
198        assert_eq!(body, "alice:user");
199    }
200
201    #[tokio::test]
202    async fn cookie_takes_precedence_over_bearer() {
203        let (public_pem, cookie_token) = keys_and_token(Role::Readonly);
204        let req = HttpRequest::get("/graphql")
205            .header(
206                header::COOKIE,
207                format!("other=1; koan_access={cookie_token}"),
208            )
209            .header(header::AUTHORIZATION, "Bearer garbage")
210            .body(Body::empty())
211            .unwrap();
212        let (status, body) = call(enforcing(public_pem, None), req).await;
213        assert_eq!(status, StatusCode::OK);
214        assert_eq!(body, "alice:readonly");
215    }
216
217    #[tokio::test]
218    async fn query_param_token_only_works_on_the_ws_route() {
219        let (public_pem, token) = keys_and_token(Role::Admin);
220
221        let req = HttpRequest::get(format!("/graphql?token={token}"))
222            .body(Body::empty())
223            .unwrap();
224        let (status, _) = call(enforcing(public_pem.clone(), None), req).await;
225        assert_eq!(status, StatusCode::UNAUTHORIZED);
226
227        let req = HttpRequest::get(format!("/graphql/ws?token={token}"))
228            .body(Body::empty())
229            .unwrap();
230        let (status, body) = call(enforcing(public_pem, None), req).await;
231        assert_eq!(status, StatusCode::OK);
232        assert_eq!(body, "alice:admin");
233    }
234
235    #[tokio::test]
236    async fn introspection_key_bypasses_auth_only_when_it_matches() {
237        let (public_pem, _) = keys_and_token(Role::Admin);
238
239        let req = HttpRequest::get("/graphql")
240            .header("X-Introspection-Key", "sekrit")
241            .body(Body::empty())
242            .unwrap();
243        let (status, body) = call(enforcing(public_pem.clone(), Some("sekrit")), req).await;
244        assert_eq!(status, StatusCode::OK);
245        assert_eq!(body, "anonymous:admin");
246
247        let req = HttpRequest::get("/graphql")
248            .header("X-Introspection-Key", "sekrjt")
249            .body(Body::empty())
250            .unwrap();
251        let (status, _) = call(enforcing(public_pem, Some("sekrit")), req).await;
252        assert_eq!(status, StatusCode::UNAUTHORIZED);
253    }
254
255    #[tokio::test]
256    async fn tampered_token_is_rejected() {
257        let (public_pem, token) = keys_and_token(Role::Admin);
258        let req = HttpRequest::get("/graphql")
259            .header(header::AUTHORIZATION, format!("Bearer {token}x"))
260            .body(Body::empty())
261            .unwrap();
262        let (status, _) = call(enforcing(public_pem, None), req).await;
263        assert_eq!(status, StatusCode::UNAUTHORIZED);
264    }
265
266    #[tokio::test]
267    async fn token_signed_by_another_key_is_rejected() {
268        let (_, token) = keys_and_token(Role::Admin);
269        let (other_public, _) = keys_and_token(Role::Admin);
270        let req = HttpRequest::get("/graphql")
271            .header(header::AUTHORIZATION, format!("Bearer {token}"))
272            .body(Body::empty())
273            .unwrap();
274        let (status, _) = call(enforcing(other_public, None), req).await;
275        assert_eq!(status, StatusCode::UNAUTHORIZED);
276    }
277
278    #[tokio::test]
279    async fn unparseable_role_claim_falls_back_to_readonly() {
280        let (private_pem, public_pem) = auth::generate_keypair_pem().unwrap();
281        let token = auth::mint_access_token_with_role_str(
282            private_pem.as_bytes(),
283            7,
284            "alice",
285            "wizard",
286            900,
287        )
288        .unwrap();
289        let req = HttpRequest::get("/graphql")
290            .header(header::AUTHORIZATION, format!("Bearer {token}"))
291            .body(Body::empty())
292            .unwrap();
293        let (status, body) = call(enforcing(public_pem.into_bytes(), None), req).await;
294        assert_eq!(status, StatusCode::OK);
295        assert_eq!(body, "alice:readonly");
296    }
297}
koan_server/auth/middleware.rs

koan_server/auth/
middleware.rs
