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
14pub type AuthState = Option<Arc<ApiKeySet>>;
16
17pub 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 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
45pub 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
84const EXEMPT_PATHS: &[&str] = &["/health", "/api/docs", "/api/openapi.json"];
86
87pub async fn require_api_key(request: Request, next: Next) -> Response {
89 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, };
96
97 let path = request.uri().path();
99 if EXEMPT_PATHS.contains(&path) {
100 return next.run(request).await;
101 }
102
103 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
129pub 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
139pub 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 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}