Skip to main content

mold_server/
auth.rs

1use axum::{
2    extract::Request,
3    http::{HeaderValue, StatusCode},
4    middleware::Next,
5    response::{IntoResponse, Response},
6    Json,
7};
8use serde::Serialize;
9use std::collections::HashSet;
10use std::sync::Arc;
11use subtle::ConstantTimeEq;
12use tracing::warn;
13
14/// Shared auth state — `None` means authentication is disabled.
15pub type AuthState = Option<Arc<ApiKeySet>>;
16
17/// Set of valid API keys loaded from `MOLD_API_KEY`.
18pub struct ApiKeySet {
19    keys: HashSet<String>,
20}
21
22impl ApiKeySet {
23    pub fn new(keys: HashSet<String>) -> Self {
24        Self { keys }
25    }
26
27    pub fn contains(&self, candidate: &str) -> bool {
28        // Check ALL keys unconditionally to avoid leaking which key matched
29        // via timing side-channel (`.any()` would short-circuit on first match).
30        let candidate_bytes = candidate.as_bytes();
31        let mut found = subtle::Choice::from(0u8);
32        for k in &self.keys {
33            found |= k.as_bytes().ct_eq(candidate_bytes);
34        }
35        found.into()
36    }
37}
38
39#[derive(Debug, Serialize)]
40struct AuthError {
41    error: String,
42    code: String,
43}
44
45/// Load API keys from the `MOLD_API_KEY` environment variable.
46///
47/// Formats:
48/// - Single key: `MOLD_API_KEY=my-secret`
49/// - Comma-separated: `MOLD_API_KEY=key1,key2,key3`
50/// - File reference: `MOLD_API_KEY=@/path/to/keys.txt` (one key per line)
51///
52/// Returns `None` when the variable is unset or empty (auth disabled).
53pub fn load_api_keys() -> anyhow::Result<AuthState> {
54    let raw = match std::env::var("MOLD_API_KEY") {
55        Ok(v) if !v.is_empty() => v,
56        _ => return Ok(None),
57    };
58
59    let keys = if let Some(path) = raw.strip_prefix('@') {
60        let contents = std::fs::read_to_string(path)
61            .map_err(|e| anyhow::anyhow!("failed to read API key file {path}: {e}"))?;
62        contents
63            .lines()
64            .map(str::trim)
65            .filter(|l| !l.is_empty() && !l.starts_with('#'))
66            .map(String::from)
67            .collect::<HashSet<_>>()
68    } else {
69        raw.split(',')
70            .map(str::trim)
71            .filter(|s| !s.is_empty())
72            .map(String::from)
73            .collect::<HashSet<_>>()
74    };
75
76    if keys.is_empty() {
77        return Ok(None);
78    }
79
80    tracing::info!(num_keys = keys.len(), "API key authentication enabled");
81    Ok(Some(Arc::new(ApiKeySet { keys })))
82}
83
84/// Paths that are exempt from API key authentication.
85const EXEMPT_PATHS: &[&str] = &["/health", "/api/docs", "/api/openapi.json"];
86
87/// Axum middleware that enforces API key authentication.
88pub async fn require_api_key(request: Request, next: Next) -> Response {
89    // Auth state is stored as an extension by the layer setup in lib.rs.
90    let auth_state = request.extensions().get::<AuthState>().cloned();
91
92    let key_set = match auth_state.as_ref().and_then(|s| s.as_ref()) {
93        Some(ks) => ks,
94        None => return next.run(request).await, // Auth disabled
95    };
96
97    // Exempt certain paths (health checks, docs).
98    let path = request.uri().path();
99    if EXEMPT_PATHS.contains(&path) {
100        return next.run(request).await;
101    }
102
103    // Check the X-Api-Key header.
104    match request.headers().get("x-api-key") {
105        Some(value) => {
106            let candidate = value.to_str().unwrap_or("");
107            if key_set.contains(candidate) {
108                next.run(request).await
109            } else {
110                warn!("rejected request with invalid API key");
111                unauthorized("invalid API key")
112            }
113        }
114        None => {
115            warn!(path = %path, "rejected request without API key");
116            unauthorized("missing X-Api-Key header")
117        }
118    }
119}
120
121fn unauthorized(msg: &str) -> Response {
122    let body = AuthError {
123        error: msg.to_string(),
124        code: "UNAUTHORIZED".to_string(),
125    };
126    (StatusCode::UNAUTHORIZED, Json(body)).into_response()
127}
128
129/// Injects the `AuthState` as a request extension so the middleware can access it.
130pub async fn inject_auth_state(
131    axum::extract::State(auth): axum::extract::State<AuthState>,
132    mut request: Request,
133    next: Next,
134) -> Response {
135    request.extensions_mut().insert(auth);
136    next.run(request).await
137}
138
139// ── CORS header exposure ────────────────────────────────────────────────────
140
141/// Header name for API key authentication (needed for CORS `Access-Control-Allow-Headers`).
142pub fn api_key_header_name() -> HeaderValue {
143    HeaderValue::from_static("x-api-key")
144}
145
146#[cfg(test)]
147mod tests {
148    use super::*;
149
150    /// Serialize env var mutations across parallel test threads.
151    fn env_lock() -> &'static std::sync::Mutex<()> {
152        static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
153        &LOCK
154    }
155
156    #[test]
157    fn parse_single_key() {
158        let _lock = env_lock().lock().unwrap();
159        unsafe { std::env::set_var("MOLD_API_KEY", "secret123") };
160        let state = load_api_keys().unwrap();
161        unsafe { std::env::remove_var("MOLD_API_KEY") };
162        let ks = state.as_ref().unwrap();
163        assert!(ks.contains("secret123"));
164        assert!(!ks.contains("wrong"));
165    }
166
167    #[test]
168    fn parse_comma_separated_keys() {
169        let _lock = env_lock().lock().unwrap();
170        unsafe { std::env::set_var("MOLD_API_KEY", "key1,key2, key3 ") };
171        let state = load_api_keys().unwrap();
172        unsafe { std::env::remove_var("MOLD_API_KEY") };
173        let ks = state.as_ref().unwrap();
174        assert!(ks.contains("key1"));
175        assert!(ks.contains("key2"));
176        assert!(ks.contains("key3"));
177        assert!(!ks.contains("key4"));
178    }
179
180    #[test]
181    fn parse_file_keys() {
182        let _lock = env_lock().lock().unwrap();
183        let dir = std::env::temp_dir().join("mold_test_keys");
184        let _ = std::fs::create_dir_all(&dir);
185        let path = dir.join("keys.txt");
186        std::fs::write(&path, "alpha\n# comment\n\nbeta\n").unwrap();
187        let env_val = format!("@{}", path.display());
188        unsafe { std::env::set_var("MOLD_API_KEY", &env_val) };
189        let state = load_api_keys().unwrap();
190        unsafe { std::env::remove_var("MOLD_API_KEY") };
191        let _ = std::fs::remove_file(&path);
192        let ks = state.as_ref().unwrap();
193        assert!(ks.contains("alpha"));
194        assert!(ks.contains("beta"));
195        assert!(!ks.contains("# comment"));
196    }
197
198    #[test]
199    fn empty_env_returns_none() {
200        let _lock = env_lock().lock().unwrap();
201        unsafe { std::env::set_var("MOLD_API_KEY", "") };
202        let state = load_api_keys().unwrap();
203        unsafe { std::env::remove_var("MOLD_API_KEY") };
204        assert!(state.is_none());
205    }
206
207    #[test]
208    fn unset_env_returns_none() {
209        let _lock = env_lock().lock().unwrap();
210        unsafe { std::env::remove_var("MOLD_API_KEY") };
211        let state = load_api_keys().unwrap();
212        assert!(state.is_none());
213    }
214
215    #[test]
216    fn constant_time_comparison_rejects_wrong_key() {
217        let ks = ApiKeySet::new(HashSet::from(["correct-key".to_string()]));
218        assert!(ks.contains("correct-key"));
219        assert!(!ks.contains("wrong-key"));
220        assert!(!ks.contains(""));
221    }
222}