Skip to main content

assay_workflow/api/
auth.rs

1//! Auth middleware + API-key helpers.
2//!
3//! `AuthMode`, `JwtConfig`, and `JwksCache` are defined in
4//! `crate::auth_mode` (no crate-internal deps) and re-exported here so
5//! the public path `assay_workflow::api::auth::AuthMode` keeps working.
6
7use std::sync::Arc;
8
9use axum::extract::{Request, State};
10use axum::http::StatusCode;
11use axum::middleware::Next;
12use axum::response::{IntoResponse, Response};
13use axum::Json;
14use sha2::{Digest, Sha256};
15use tracing::{info, warn};
16
17use crate::ctx::WorkflowCtx;
18use crate::store::WorkflowStore;
19
20// Re-export the config types from their cycle-free home.
21pub use crate::auth_mode::{AuthMode, JwksCache, JwtConfig};
22
23// ── Middleware ───────────────────────────────────────────────
24
25/// Axum middleware that enforces authentication based on the configured mode.
26///
27/// When both JWT and API-key auth are enabled, dispatch is based on token shape:
28/// if the Bearer token parses as a JWS header it takes the JWT path, otherwise
29/// the API-key path. A semantically-invalid JWT (expired, forged signature, wrong
30/// audience) is rejected and is *not* retried as an API key — a token that looks
31/// like a JWT is treated as a JWT.
32///
33/// **Bootstrap window:** `POST /api/v1/api-keys` is accepted without a Bearer
34/// token iff the `api_keys` table is empty. This is the only way a freshly
35/// deployed server running in API-key or combined mode can receive its first
36/// credential without operator shell access. The window closes the moment any
37/// key exists.
38pub async fn auth_middleware<S: WorkflowStore>(
39    State(state): State<Arc<WorkflowCtx<S>>>,
40    request: Request,
41    next: Next,
42) -> Response {
43    let auth = &state.auth_mode;
44
45    if !auth.is_enabled() {
46        return next.run(request).await;
47    }
48
49    if is_bootstrap_request(&request) {
50        match state.store().api_keys_empty().await {
51            Ok(true) => {
52                info!(
53                    "Allowing unauthenticated POST /api/v1/api-keys — api_keys table is empty (bootstrap window)"
54                );
55                return next.run(request).await;
56            }
57            Ok(false) => {
58                // Fall through to normal auth — bootstrap window is closed.
59            }
60            Err(e) => {
61                warn!("api_keys_empty check failed: {e}");
62                return (
63                    StatusCode::INTERNAL_SERVER_ERROR,
64                    Json(serde_json::json!({"error": "auth bootstrap check failed"})),
65                )
66                    .into_response();
67            }
68        }
69    }
70
71    let token = match extract_bearer(&request) {
72        Some(t) => t,
73        None => return auth_error("Missing Authorization: Bearer <token>"),
74    };
75
76    if jsonwebtoken::decode_header(token).is_ok() {
77        match &auth.jwt {
78            Some(jwt) => {
79                validate_jwt(
80                    &jwt.issuer,
81                    jwt.audience.as_deref(),
82                    &jwt.jwks_cache,
83                    request,
84                    next,
85                )
86                .await
87            }
88            None => auth_error("JWT authentication is not enabled on this server"),
89        }
90    } else if auth.api_key {
91        validate_api_key(state, request, next).await
92    } else {
93        auth_error("Token is not a valid JWT and API-key authentication is not enabled")
94    }
95}
96
97async fn validate_api_key<S: WorkflowStore>(
98    state: Arc<WorkflowCtx<S>>,
99    request: Request,
100    next: Next,
101) -> Response {
102    let token = match extract_bearer(&request) {
103        Some(t) => t,
104        None => return auth_error("Missing Authorization: Bearer <api-key>"),
105    };
106
107    let hash = hash_api_key(token);
108    match state.store().validate_api_key(&hash).await {
109        Ok(true) => next.run(request).await,
110        Ok(false) => {
111            warn!(
112                "Invalid API key (prefix: {}...)",
113                &token[..8.min(token.len())]
114            );
115            auth_error("Invalid API key")
116        }
117        Err(e) => {
118            warn!("API key validation error: {e}");
119            (
120                StatusCode::INTERNAL_SERVER_ERROR,
121                Json(serde_json::json!({"error": "auth check failed"})),
122            )
123                .into_response()
124        }
125    }
126}
127
128async fn validate_jwt(
129    issuer: &str,
130    audience: Option<&str>,
131    jwks_cache: &JwksCache,
132    request: Request,
133    next: Next,
134) -> Response {
135    use jsonwebtoken::Validation;
136
137    let token = match extract_bearer(&request) {
138        Some(t) => t,
139        None => return auth_error("Missing Authorization: Bearer <jwt>"),
140    };
141
142    // Decode header to get algorithm and kid
143    let header = match jsonwebtoken::decode_header(token) {
144        Ok(h) => h,
145        Err(e) => {
146            warn!("Invalid JWT header: {e}");
147            return auth_error("Invalid JWT");
148        }
149    };
150
151    // Find the decoding key from JWKS
152    let decoding_key = match &header.kid {
153        Some(kid) => match jwks_cache.find_key(kid).await {
154            Ok(key) => key,
155            Err(e) => {
156                warn!("JWKS key lookup failed: {e}");
157                return auth_error("JWT validation failed: key not found");
158            }
159        },
160        None => match jwks_cache.find_any_key(header.alg).await {
161            Ok(key) => key,
162            Err(e) => {
163                warn!("JWKS key lookup failed (no kid): {e}");
164                return auth_error("JWT validation failed: no suitable key");
165            }
166        },
167    };
168
169    // Build validation rules
170    let mut validation = Validation::new(header.alg);
171    validation.set_issuer(&[issuer]);
172    if let Some(aud) = audience {
173        validation.set_audience(&[aud]);
174    } else {
175        validation.validate_aud = false;
176    }
177
178    // Validate signature + claims
179    match jsonwebtoken::decode::<serde_json::Value>(token, &decoding_key, &validation) {
180        Ok(_) => next.run(request).await,
181        Err(e) => {
182            warn!("JWT validation failed: {e}");
183            auth_error(&format!("JWT validation failed: {e}"))
184        }
185    }
186}
187
188fn extract_bearer(request: &Request) -> Option<&str> {
189    request
190        .headers()
191        .get("authorization")
192        .and_then(|v| v.to_str().ok())
193        .and_then(|v| v.strip_prefix("Bearer "))
194}
195
196/// True iff the request is the bootstrap-window endpoint
197/// (`POST /api/v1/api-keys`). The caller is still responsible for checking
198/// that the `api_keys` table is empty before actually allowing unauth access.
199fn is_bootstrap_request(request: &Request) -> bool {
200    request.method() == axum::http::Method::POST
201        && request.uri().path() == "/api/v1/api-keys"
202}
203
204fn auth_error(msg: &str) -> Response {
205    (
206        StatusCode::UNAUTHORIZED,
207        Json(serde_json::json!({"error": msg})),
208    )
209        .into_response()
210}
211
212// ── API Key Helpers ─────────────────────────────────────────
213
214/// Hash an API key with SHA-256 for storage/lookup.
215pub fn hash_api_key(key: &str) -> String {
216    let mut hasher = Sha256::new();
217    hasher.update(key.as_bytes());
218    data_encoding::HEXLOWER.encode(&hasher.finalize())
219}
220
221/// Generate a new random API key (32 bytes, hex-encoded).
222pub fn generate_api_key() -> String {
223    use rand::Rng;
224    let bytes: [u8; 32] = rand::rng().random();
225    format!("assay_{}", data_encoding::HEXLOWER.encode(&bytes))
226}
227
228/// Extract the prefix (first 8 chars after "assay_") for display.
229pub fn key_prefix(key: &str) -> String {
230    let stripped = key.strip_prefix("assay_").unwrap_or(key);
231    format!("assay_{}...", &stripped[..8.min(stripped.len())])
232}