assay_workflow/
auth_mode.rs

1//! Auth configuration types — `AuthMode`, `JwtConfig`, `JwksCache`.
2//! Lives at crate root (not under `api/`) so `ctx.rs` can import
3//! `AuthMode` without creating a dependency cycle with `api/auth.rs`.
4//!
5//! `api/auth.rs` re-exports these types so the public path
6//! `assay_workflow::api::auth::AuthMode` continues to work.
7
8use std::sync::Arc;
9use std::time::{Duration, Instant};
10
11use jsonwebtoken::jwk::JwkSet;
12use tokio::sync::RwLock;
13use tracing::{debug, info};
14
15const JWKS_CACHE_TTL: Duration = Duration::from_secs(300); // 5 minutes
16
17// ── Auth Mode ───────────────────────────────────────────────
18
19/// Auth configuration for the engine's HTTP API.
20///
21/// Both authentication methods (JWT and API key) can be enabled at the same time.
22/// When both are enabled, the middleware dispatches on token shape — tokens that
23/// parse as a JWS header are validated as JWTs, everything else is validated as
24/// an API key. This lets the same server accept long-lived machine API keys
25/// alongside short-lived OIDC-issued user tokens without the caller picking a
26/// mode up front.
27#[derive(Clone, Debug, Default)]
28pub struct AuthMode {
29    /// API-key authentication enabled. When true, Bearer tokens that are not
30    /// JWT-shaped are validated against the `api_keys` table.
31    pub api_key: bool,
32    /// JWT authentication enabled. When set, Bearer tokens that parse as a
33    /// JWS header are validated against the issuer's JWKS.
34    pub jwt: Option<JwtConfig>,
35}
36
37/// JWT validation configuration.
38#[derive(Clone, Debug)]
39pub struct JwtConfig {
40    pub issuer: String,
41    pub audience: Option<String>,
42    pub jwks_cache: Arc<JwksCache>,
43}
44
45impl AuthMode {
46    /// Open access — no authentication. All requests allowed.
47    pub fn no_auth() -> Self {
48        Self::default()
49    }
50
51    /// JWT/OIDC only. Tokens are validated against the issuer's JWKS.
52    pub fn jwt(issuer: String, audience: Option<String>) -> Self {
53        Self {
54            api_key: false,
55            jwt: Some(JwtConfig::new(issuer, audience)),
56        }
57    }
58
59    /// API key only. Bearer tokens are hashed and looked up in the store.
60    pub fn api_key() -> Self {
61        Self {
62            api_key: true,
63            jwt: None,
64        }
65    }
66
67    /// Both JWT and API key. Tokens that parse as JWTs take the JWT path;
68    /// everything else takes the API-key path.
69    pub fn combined(issuer: String, audience: Option<String>) -> Self {
70        Self {
71            api_key: true,
72            jwt: Some(JwtConfig::new(issuer, audience)),
73        }
74    }
75
76    /// True if any authentication method is enabled.
77    pub fn is_enabled(&self) -> bool {
78        self.api_key || self.jwt.is_some()
79    }
80
81    /// Human-readable summary for startup logging.
82    pub fn describe(&self) -> String {
83        match (self.jwt.as_ref(), self.api_key) {
84            (None, false) => "no-auth (open access)".to_string(),
85            (None, true) => "api-key".to_string(),
86            (Some(c), false) => format!("jwt (issuer: {})", c.issuer),
87            (Some(c), true) => format!("jwt (issuer: {}) + api-key", c.issuer),
88        }
89    }
90}
91
92impl JwtConfig {
93    /// Build a JwtConfig with a fresh JWKS cache pointed at `issuer`'s OIDC discovery endpoint.
94    pub fn new(issuer: String, audience: Option<String>) -> Self {
95        Self {
96            jwks_cache: Arc::new(JwksCache::new(issuer.clone())),
97            issuer,
98            audience,
99        }
100    }
101}
102
103// ── JWKS Cache ──────────────────────────────────────────────
104
105/// Caches JWKS keys fetched from the OIDC provider.
106/// Keys are refreshed after `JWKS_CACHE_TTL` or on cache miss for a specific `kid`.
107pub struct JwksCache {
108    issuer: String,
109    cache: RwLock<Option<CachedJwks>>,
110    http: reqwest::Client,
111}
112
113struct CachedJwks {
114    jwks: JwkSet,
115    fetched_at: Instant,
116}
117
118impl std::fmt::Debug for JwksCache {
119    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
120        f.debug_struct("JwksCache")
121            .field("issuer", &self.issuer)
122            .finish()
123    }
124}
125
126impl JwksCache {
127    pub fn new(issuer: String) -> Self {
128        Self {
129            issuer,
130            cache: RwLock::new(None),
131            http: reqwest::Client::builder()
132                .timeout(Duration::from_secs(10))
133                .build()
134                .expect("building JWKS HTTP client"),
135        }
136    }
137
138    /// For testing: create a cache pre-loaded with keys (no HTTP fetching needed).
139    pub fn with_jwks(issuer: String, jwks: JwkSet) -> Self {
140        Self {
141            issuer,
142            cache: RwLock::new(Some(CachedJwks {
143                jwks,
144                fetched_at: Instant::now(),
145            })),
146            http: reqwest::Client::new(),
147        }
148    }
149
150    /// Get the JWKS, fetching from the provider if the cache is stale or empty.
151    pub async fn get_jwks(&self) -> anyhow::Result<JwkSet> {
152        // Check cache
153        {
154            let cache = self.cache.read().await;
155            if let Some(ref cached) = *cache
156                && cached.fetched_at.elapsed() < JWKS_CACHE_TTL
157            {
158                return Ok(cached.jwks.clone());
159            }
160        }
161
162        // Cache miss or stale — fetch fresh JWKS
163        self.refresh().await
164    }
165
166    /// Force-refresh the JWKS (e.g., when a kid is not found in the current set).
167    pub async fn refresh(&self) -> anyhow::Result<JwkSet> {
168        let jwks_uri = self.discover_jwks_uri().await?;
169        debug!("Fetching JWKS from {jwks_uri}");
170
171        let jwks: JwkSet = self.http.get(&jwks_uri).send().await?.json().await?;
172        info!(
173            "Fetched {} keys from JWKS endpoint",
174            jwks.keys.len()
175        );
176
177        let mut cache = self.cache.write().await;
178        *cache = Some(CachedJwks {
179            jwks: jwks.clone(),
180            fetched_at: Instant::now(),
181        });
182
183        Ok(jwks)
184    }
185
186    /// Discover the JWKS URI from the OIDC discovery endpoint.
187    async fn discover_jwks_uri(&self) -> anyhow::Result<String> {
188        let discovery_url = format!(
189            "{}/.well-known/openid-configuration",
190            self.issuer.trim_end_matches('/')
191        );
192
193        let resp: serde_json::Value = self
194            .http
195            .get(&discovery_url)
196            .send()
197            .await?
198            .json()
199            .await?;
200
201        resp.get("jwks_uri")
202            .and_then(|v| v.as_str())
203            .map(String::from)
204            .ok_or_else(|| anyhow::anyhow!("OIDC discovery response missing jwks_uri"))
205    }
206
207    /// Find a decoding key by `kid` (key ID) from the cached JWKS.
208    /// If the kid isn't found, refreshes the cache once and retries.
209    pub async fn find_key(&self, kid: &str) -> anyhow::Result<jsonwebtoken::DecodingKey> {
210        let jwks = self.get_jwks().await?;
211
212        // Try to find by kid
213        if let Some(key) = find_key_in_set(&jwks, kid) {
214            return Ok(key);
215        }
216
217        // kid not found — refresh and retry (key rotation)
218        debug!("kid '{kid}' not in JWKS cache, refreshing");
219        let jwks = self.refresh().await?;
220
221        find_key_in_set(&jwks, kid)
222            .ok_or_else(|| anyhow::anyhow!("No key with kid '{kid}' in JWKS"))
223    }
224
225    /// Find a decoding key when no kid is provided (use the first matching key).
226    pub async fn find_any_key(
227        &self,
228        alg: jsonwebtoken::Algorithm,
229    ) -> anyhow::Result<jsonwebtoken::DecodingKey> {
230        let jwks = self.get_jwks().await?;
231
232        for key in &jwks.keys {
233            if let Ok(dk) = jsonwebtoken::DecodingKey::from_jwk(key) {
234                let _ = alg;
235                return Ok(dk);
236            }
237        }
238
239        anyhow::bail!("No suitable key found in JWKS for algorithm {alg:?}")
240    }
241}
242
243fn find_key_in_set(
244    jwks: &JwkSet,
245    kid: &str,
246) -> Option<jsonwebtoken::DecodingKey> {
247    jwks.keys
248        .iter()
249        .find(|k| k.common.key_id.as_deref() == Some(kid))
250        .and_then(|k| jsonwebtoken::DecodingKey::from_jwk(k).ok())
251}
assay_workflow/auth_mode.rs

assay_workflow/
auth_mode.rs
