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;
15use koan_core::db::pool::Pool;
16
17use super::AuthUser;
18
19/// Shared state for the auth middleware.
20#[derive(Clone)]
21pub struct AuthState {
22    /// Ed25519 public key PEM for JWT verification.
23    pub public_pem: Arc<Vec<u8>>,
24    /// Whether auth is enforced.
25    pub auth_enabled: bool,
26    /// Process-scoped introspection key. Bypasses auth when matched.
27    /// Generated randomly on server start, dies with the process.
28    pub introspection_key: Option<Arc<String>>,
29    /// Where a token's account is looked up; see `super::current_user`.
30    pub pool: Arc<Pool>,
31}
32
33/// Axum middleware: validate JWT and inject `AuthUser`.
34///
35/// When `auth_enabled = false`, injects anonymous admin and passes through.
36/// When `auth_enabled = true`, requires a valid token.
37pub async fn auth_middleware(
38    axum::extract::State(state): axum::extract::State<AuthState>,
39    mut request: Request,
40    next: Next,
41) -> Response {
42    if !state.auth_enabled {
43        request.extensions_mut().insert(AuthUser::anonymous_admin());
44        return next.run(request).await;
45    }
46
47    // Check for introspection key (playground bypass).
48    if let Some(ref expected_key) = state.introspection_key
49        && let Some(provided) = request
50            .headers()
51            .get("X-Introspection-Key")
52            .and_then(|v| v.to_str().ok())
53        && provided
54            .as_bytes()
55            .ct_eq(expected_key.as_bytes())
56            .unwrap_u8()
57            == 1
58    {
59        request.extensions_mut().insert(AuthUser::anonymous_admin());
60        return next.run(request).await;
61    }
62
63    let Some(token) = extract_token(&request) else {
64        return (
65            StatusCode::UNAUTHORIZED,
66            [("WWW-Authenticate", "Bearer")],
67            "missing or invalid Authorization header",
68        )
69            .into_response();
70    };
71
72    let mark = auth::account_mark();
73    // An access token, or failing that an API key: the key a device or a
74    // script was given, which acts as its account until revoked. A key has no
75    // expiry of its own, so a socket opened with one closes only when the
76    // account changes, a revocation included.
77    let user = match auth::validate_access_token(&state.public_pem, &token) {
78        Ok(claims) => {
79            let expires = claims.exp;
80            super::current_user(&state.pool, claims)
81                .await
82                .map(|user| (user, Some(expires)))
83        }
84        Err(_) => super::key_user(&state.pool, token)
85            .await
86            .map(|user| (user, None)),
87    };
88    match user {
89        Some((user, expires)) => {
90            request.extensions_mut().insert(super::Lease {
91                user_id: user.user_id,
92                mark,
93                expires,
94            });
95            request.extensions_mut().insert(user);
96            next.run(request).await
97        }
98        None => (
99            StatusCode::UNAUTHORIZED,
100            [("WWW-Authenticate", "Bearer")],
101            "invalid or expired token",
102        )
103            .into_response(),
104    }
105}
106
107/// Priority: `koan_access` cookie, then `Authorization: Bearer`, then `?token=`.
108///
109/// The query parameter is confined to the WebSocket route, which is the only one
110/// that cannot carry a header. A token in a URL survives in shell history, proxy
111/// logs and `Referer`.
112fn extract_token(request: &Request) -> Option<String> {
113    request
114        .headers()
115        .get(header::COOKIE)
116        .and_then(|v| v.to_str().ok())
117        .and_then(|cookies| {
118            cookies
119                .split(';')
120                .find_map(|c| c.trim().strip_prefix("koan_access=").map(String::from))
121        })
122        .or_else(|| {
123            request
124                .headers()
125                .get(header::AUTHORIZATION)
126                .and_then(|v| v.to_str().ok())
127                .and_then(|v| v.strip_prefix("Bearer "))
128                .map(String::from)
129        })
130        .or_else(|| {
131            if request.uri().path() != "/graphql/ws" {
132                return None;
133            }
134            request.uri().query().and_then(|q| {
135                q.split('&')
136                    .find_map(|pair| pair.strip_prefix("token=").map(String::from))
137            })
138        })
139}
140
141// ---------------------------------------------------------------------------
142// Tests
143// ---------------------------------------------------------------------------
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148    use axum::body::Body;
149    use axum::http::Request as HttpRequest;
150    use axum::routing::get;
151    use koan_core::auth::Role;
152    use tower::ServiceExt as _;
153
154    /// Echoes the `AuthUser` the middleware injected, so tests can assert on it.
155    async fn echo_user(axum::Extension(user): axum::Extension<AuthUser>) -> String {
156        format!("{}:{}", user.username, user.role.as_str())
157    }
158
159    async fn call(state: AuthState, req: HttpRequest<Body>) -> (StatusCode, String) {
160        let app = axum::Router::new()
161            .route("/graphql", get(echo_user))
162            .route("/graphql/ws", get(echo_user))
163            .layer(axum::middleware::from_fn_with_state(state, auth_middleware));
164        let resp = app.oneshot(req).await.unwrap();
165        let status = resp.status();
166        let bytes = axum::body::to_bytes(resp.into_body(), 64 * 1024)
167            .await
168            .unwrap();
169        (status, String::from_utf8_lossy(&bytes).into_owned())
170    }
171
172    /// A live keypair plus a matching token for `alice` at `role`.
173    fn keys_and_token(role: Role) -> (Vec<u8>, String) {
174        let (private_pem, public_pem) = auth::generate_keypair_pem().unwrap();
175        let token = auth::mint_access_token(private_pem.as_bytes(), 1, "alice", role, 900).unwrap();
176        (public_pem.into_bytes(), token)
177    }
178
179    /// Enforcing auth over a database whose first account is `username` at
180    /// `role`: id 1, which the tokens above name.
181    fn enforcing_with(
182        public_pem: Vec<u8>,
183        key: Option<&str>,
184        username: &str,
185        role: Role,
186    ) -> (AuthState, tempfile::TempDir) {
187        let dir = tempfile::tempdir().unwrap();
188        let path = dir.path().join("koan.db");
189        let db = koan_core::db::connection::Database::open(&path).unwrap();
190        koan_core::db::queries::auth::create_user(&db.conn, username, "pw", role).unwrap();
191        let state = AuthState {
192            public_pem: Arc::new(public_pem),
193            auth_enabled: true,
194            introspection_key: key.map(|k| Arc::new(k.to_string())),
195            pool: Arc::new(Pool::new(path)),
196        };
197        (state, dir)
198    }
199
200    fn enforcing(
201        public_pem: Vec<u8>,
202        key: Option<&str>,
203        role: Role,
204    ) -> (AuthState, tempfile::TempDir) {
205        enforcing_with(public_pem, key, "alice", role)
206    }
207
208    #[tokio::test]
209    async fn auth_disabled_grants_anonymous_admin() {
210        let state = AuthState {
211            public_pem: Arc::new(Vec::new()),
212            auth_enabled: false,
213            introspection_key: None,
214            pool: Arc::new(Pool::new("/nonexistent/koan.db".into())),
215        };
216        let req = HttpRequest::get("/graphql").body(Body::empty()).unwrap();
217        let (status, body) = call(state, req).await;
218        assert_eq!(status, StatusCode::OK);
219        assert_eq!(body, "anonymous:admin");
220    }
221
222    #[tokio::test]
223    async fn missing_token_is_unauthorized() {
224        let (public_pem, _) = keys_and_token(Role::Admin);
225        let req = HttpRequest::get("/graphql").body(Body::empty()).unwrap();
226        let (state, _dir) = enforcing(public_pem, None, Role::Admin);
227        let (status, _) = call(state, req).await;
228        assert_eq!(status, StatusCode::UNAUTHORIZED);
229    }
230
231    #[tokio::test]
232    async fn bearer_token_authenticates() {
233        let (public_pem, token) = keys_and_token(Role::User);
234        let req = HttpRequest::get("/graphql")
235            .header(header::AUTHORIZATION, format!("Bearer {token}"))
236            .body(Body::empty())
237            .unwrap();
238        let (state, _dir) = enforcing(public_pem, None, Role::User);
239        let (status, body) = call(state, req).await;
240        assert_eq!(status, StatusCode::OK);
241        assert_eq!(body, "alice:user");
242    }
243
244    #[tokio::test]
245    async fn cookie_takes_precedence_over_bearer() {
246        let (public_pem, cookie_token) = keys_and_token(Role::Readonly);
247        let req = HttpRequest::get("/graphql")
248            .header(
249                header::COOKIE,
250                format!("other=1; koan_access={cookie_token}"),
251            )
252            .header(header::AUTHORIZATION, "Bearer garbage")
253            .body(Body::empty())
254            .unwrap();
255        let (state, _dir) = enforcing(public_pem, None, Role::Readonly);
256        let (status, body) = call(state, req).await;
257        assert_eq!(status, StatusCode::OK);
258        assert_eq!(body, "alice:readonly");
259    }
260
261    #[tokio::test]
262    async fn query_param_token_only_works_on_the_ws_route() {
263        let (public_pem, token) = keys_and_token(Role::Admin);
264
265        let req = HttpRequest::get(format!("/graphql?token={token}"))
266            .body(Body::empty())
267            .unwrap();
268        let (state, _dir) = enforcing(public_pem, None, Role::Admin);
269        let (status, _) = call(state.clone(), req).await;
270        assert_eq!(status, StatusCode::UNAUTHORIZED);
271
272        let req = HttpRequest::get(format!("/graphql/ws?token={token}"))
273            .body(Body::empty())
274            .unwrap();
275        let (status, body) = call(state, req).await;
276        assert_eq!(status, StatusCode::OK);
277        assert_eq!(body, "alice:admin");
278    }
279
280    #[tokio::test]
281    async fn introspection_key_bypasses_auth_only_when_it_matches() {
282        let (public_pem, _) = keys_and_token(Role::Admin);
283
284        let req = HttpRequest::get("/graphql")
285            .header("X-Introspection-Key", "sekrit")
286            .body(Body::empty())
287            .unwrap();
288        let (state, _dir) = enforcing(public_pem, Some("sekrit"), Role::Admin);
289        let (status, body) = call(state.clone(), req).await;
290        assert_eq!(status, StatusCode::OK);
291        assert_eq!(body, "anonymous:admin");
292
293        let req = HttpRequest::get("/graphql")
294            .header("X-Introspection-Key", "sekrjt")
295            .body(Body::empty())
296            .unwrap();
297        let (status, _) = call(state, req).await;
298        assert_eq!(status, StatusCode::UNAUTHORIZED);
299    }
300
301    #[tokio::test]
302    async fn api_key_authenticates_as_its_account_until_revoked() {
303        let (public_pem, _) = keys_and_token(Role::Admin);
304        let (state, _dir) = enforcing(public_pem, None, Role::User);
305        let key = {
306            let db = state.pool.get().unwrap();
307            koan_core::db::queries::api_keys::create_api_key(&db.conn, 1, "script")
308                .unwrap()
309                .1
310        };
311        let bearer = |key: &str| {
312            HttpRequest::get("/graphql")
313                .header(header::AUTHORIZATION, format!("Bearer {key}"))
314                .body(Body::empty())
315                .unwrap()
316        };
317
318        let (status, body) = call(state.clone(), bearer(&key)).await;
319        assert_eq!(status, StatusCode::OK);
320        assert_eq!(body, "alice:user");
321
322        let (status, _) = call(state.clone(), bearer(&format!("{key}x"))).await;
323        assert_eq!(status, StatusCode::UNAUTHORIZED);
324
325        {
326            let db = state.pool.get().unwrap();
327            assert!(
328                koan_core::db::queries::api_keys::revoke_api_key_value(&db.conn, &key).unwrap()
329            );
330        }
331        let (status, _) = call(state, bearer(&key)).await;
332        assert_eq!(status, StatusCode::UNAUTHORIZED);
333    }
334
335    #[tokio::test]
336    async fn tampered_token_is_rejected() {
337        let (public_pem, token) = keys_and_token(Role::Admin);
338        let req = HttpRequest::get("/graphql")
339            .header(header::AUTHORIZATION, format!("Bearer {token}x"))
340            .body(Body::empty())
341            .unwrap();
342        let (state, _dir) = enforcing(public_pem, None, Role::Admin);
343        let (status, _) = call(state, req).await;
344        assert_eq!(status, StatusCode::UNAUTHORIZED);
345    }
346
347    #[tokio::test]
348    async fn token_signed_by_another_key_is_rejected() {
349        let (_, token) = keys_and_token(Role::Admin);
350        let (other_public, _) = keys_and_token(Role::Admin);
351        let req = HttpRequest::get("/graphql")
352            .header(header::AUTHORIZATION, format!("Bearer {token}"))
353            .body(Body::empty())
354            .unwrap();
355        let (state, _dir) = enforcing(other_public, None, Role::Admin);
356        let (status, _) = call(state, req).await;
357        assert_eq!(status, StatusCode::UNAUTHORIZED);
358    }
359
360    #[tokio::test]
361    async fn the_role_is_the_accounts_now_not_the_tokens() {
362        let (public_pem, token) = keys_and_token(Role::Admin);
363        let req = || {
364            HttpRequest::get("/graphql")
365                .header(header::AUTHORIZATION, format!("Bearer {token}"))
366                .body(Body::empty())
367                .unwrap()
368        };
369        // Demoted since the token was minted.
370        let (state, _dir) = enforcing(public_pem.clone(), None, Role::Readonly);
371        assert_eq!(
372            call(state, req()).await,
373            (StatusCode::OK, "alice:readonly".into())
374        );
375
376        // Deleted since.
377        let (state, dir) = enforcing(public_pem.clone(), None, Role::Admin);
378        let db = koan_core::db::connection::Database::open(&dir.path().join("koan.db")).unwrap();
379        koan_core::db::queries::auth::delete_user(&db.conn, 1).unwrap();
380        assert_eq!(call(state, req()).await.0, StatusCode::UNAUTHORIZED);
381
382        // Its id now belongs to someone else.
383        let (state, _dir) = enforcing_with(public_pem, None, "bob", Role::Admin);
384        assert_eq!(call(state, req()).await.0, StatusCode::UNAUTHORIZED);
385    }
386}
koan_server/auth/middleware.rs

koan_server/auth/
middleware.rs
