1use async_trait::async_trait;
52use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
53use chrono::{DateTime, Duration, Utc};
54use rand::Rng;
55use serde::{Deserialize, Serialize};
56use sha2::{Digest, Sha256};
57use std::sync::Arc;
58use thiserror::Error;
59
60#[derive(Debug, Error)]
62pub enum ApiKeyError {
63 #[error("Invalid API key")]
64 Invalid,
65
66 #[error("API key expired")]
67 Expired,
68
69 #[error("API key revoked")]
70 Revoked,
71
72 #[error("Insufficient permissions: {0}")]
73 InsufficientPermissions(String),
74
75 #[error("Storage error: {0}")]
76 Storage(String),
77
78 #[error("Rate limit exceeded")]
79 RateLimitExceeded,
80}
81
82#[derive(Debug, Clone, Serialize, Deserialize)]
84pub struct ApiKey {
85 pub id: String,
87
88 pub key: String,
90
91 pub user_id: String,
93
94 pub name: Option<String>,
96
97 pub scopes: Vec<String>,
99
100 pub created_at: DateTime<Utc>,
102
103 pub expires_at: Option<DateTime<Utc>>,
105
106 pub last_used_at: Option<DateTime<Utc>>,
108
109 pub revoked: bool,
111
112 pub rate_limit: Option<u32>,
114}
115
116#[async_trait]
120pub trait ApiKeyStore: Send + Sync {
121 async fn save(&self, key: &ApiKey) -> Result<(), ApiKeyError>;
123
124 async fn find_by_key(&self, key: &str) -> Result<Option<ApiKey>, ApiKeyError>;
126
127 async fn find_by_id(&self, id: &str) -> Result<Option<ApiKey>, ApiKeyError>;
129
130 async fn list_by_user(&self, user_id: &str) -> Result<Vec<ApiKey>, ApiKeyError>;
132
133 async fn revoke(&self, key_id: &str) -> Result<(), ApiKeyError>;
135
136 async fn update_last_used(
138 &self,
139 key_id: &str,
140 timestamp: DateTime<Utc>,
141 ) -> Result<(), ApiKeyError>;
142}
143
144pub struct ApiKeyManager {
148 store: Arc<dyn ApiKeyStore>,
149 key_prefix: String,
150 default_expiration: Option<Duration>,
151}
152
153impl ApiKeyManager {
154 pub fn new(store: Arc<dyn ApiKeyStore>) -> Self {
177 Self {
178 store,
179 key_prefix: "ak".to_string(),
180 default_expiration: Some(Duration::days(365)),
181 }
182 }
183
184 pub fn with_prefix(mut self, prefix: impl Into<String>) -> Self {
186 self.key_prefix = prefix.into();
187 self
188 }
189
190 pub fn with_expiration(mut self, duration: Option<Duration>) -> Self {
192 self.default_expiration = duration;
193 self
194 }
195
196 pub async fn generate(
215 &self,
216 user_id: impl Into<String>,
217 scopes: Vec<String>,
218 ) -> Result<ApiKey, ApiKeyError> {
219 let key_id = uuid::Uuid::new_v4().to_string();
220 let raw_key = self.generate_random_key();
221 let key_string = format!("{}_{}", self.key_prefix, raw_key);
222
223 let expires_at = self.default_expiration.map(|d| Utc::now() + d);
224
225 let api_key = ApiKey {
226 id: key_id,
227 key: key_string.clone(),
228 user_id: user_id.into(),
229 name: None,
230 scopes,
231 created_at: Utc::now(),
232 expires_at,
233 last_used_at: None,
234 revoked: false,
235 rate_limit: None,
236 };
237
238 self.store.save(&api_key).await?;
240
241 Ok(api_key)
242 }
243
244 pub async fn validate(&self, key: &str) -> Result<Option<ApiKey>, ApiKeyError> {
248 let api_key = match self.store.find_by_key(key).await? {
249 Some(k) => k,
250 None => return Ok(None),
251 };
252
253 if api_key.revoked {
255 return Err(ApiKeyError::Revoked);
256 }
257
258 if let Some(expires_at) = api_key.expires_at {
260 if Utc::now() > expires_at {
261 return Err(ApiKeyError::Expired);
262 }
263 }
264
265 self.store.update_last_used(&api_key.id, Utc::now()).await?;
267
268 Ok(Some(api_key))
269 }
270
271 pub fn has_scope(&self, api_key: &ApiKey, required_scope: &str) -> bool {
273 api_key
274 .scopes
275 .iter()
276 .any(|s| s == required_scope || s == "*")
277 }
278
279 pub async fn revoke(&self, key_id: &str) -> Result<(), ApiKeyError> {
281 self.store.revoke(key_id).await
282 }
283
284 pub async fn list_user_keys(&self, user_id: &str) -> Result<Vec<ApiKey>, ApiKeyError> {
286 self.store.list_by_user(user_id).await
287 }
288
289 pub async fn rotate(&self, old_key_id: &str) -> Result<ApiKey, ApiKeyError> {
291 let old_key = self
293 .store
294 .find_by_id(old_key_id)
295 .await?
296 .ok_or(ApiKeyError::Invalid)?;
297
298 self.revoke(old_key_id).await?;
300
301 self.generate(&old_key.user_id, old_key.scopes).await
303 }
304
305 fn generate_random_key(&self) -> String {
307 let mut rng = rand::rng();
308 let bytes: Vec<u8> = (0..32).map(|_| rng.random()).collect();
309 URL_SAFE_NO_PAD.encode(bytes)
310 }
311
312 pub fn hash_key(key: &str) -> String {
314 let mut hasher = Sha256::new();
315 hasher.update(key.as_bytes());
316 hex::encode(hasher.finalize())
317 }
318}
319
320#[cfg(test)]
321mod tests {
322 use super::*;
323
324 struct InMemoryStore {
325 keys: std::sync::Mutex<Vec<ApiKey>>,
326 }
327
328 impl InMemoryStore {
329 fn new() -> Self {
330 Self {
331 keys: std::sync::Mutex::new(Vec::new()),
332 }
333 }
334 }
335
336 #[async_trait]
337 impl ApiKeyStore for InMemoryStore {
338 async fn save(&self, key: &ApiKey) -> Result<(), ApiKeyError> {
339 let mut keys = self.keys.lock().unwrap();
340 keys.push(key.clone());
341 Ok(())
342 }
343
344 async fn find_by_key(&self, key: &str) -> Result<Option<ApiKey>, ApiKeyError> {
345 let keys = self.keys.lock().unwrap();
346 Ok(keys.iter().find(|k| k.key == key).cloned())
347 }
348
349 async fn find_by_id(&self, id: &str) -> Result<Option<ApiKey>, ApiKeyError> {
350 let keys = self.keys.lock().unwrap();
351 Ok(keys.iter().find(|k| k.id == id).cloned())
352 }
353
354 async fn list_by_user(&self, user_id: &str) -> Result<Vec<ApiKey>, ApiKeyError> {
355 let keys = self.keys.lock().unwrap();
356 Ok(keys
357 .iter()
358 .filter(|k| k.user_id == user_id)
359 .cloned()
360 .collect())
361 }
362
363 async fn revoke(&self, key_id: &str) -> Result<(), ApiKeyError> {
364 let mut keys = self.keys.lock().unwrap();
365 if let Some(key) = keys.iter_mut().find(|k| k.id == key_id) {
366 key.revoked = true;
367 }
368 Ok(())
369 }
370
371 async fn update_last_used(
372 &self,
373 key_id: &str,
374 timestamp: DateTime<Utc>,
375 ) -> Result<(), ApiKeyError> {
376 let mut keys = self.keys.lock().unwrap();
377 if let Some(key) = keys.iter_mut().find(|k| k.id == key_id) {
378 key.last_used_at = Some(timestamp);
379 }
380 Ok(())
381 }
382 }
383
384 #[tokio::test]
385 async fn test_generate_api_key() {
386 let store = Arc::new(InMemoryStore::new());
387 let manager = ApiKeyManager::new(store);
388
389 let key = manager
390 .generate("user_123", vec!["read".to_string()])
391 .await
392 .unwrap();
393
394 assert!(key.key.starts_with("ak_"));
395 assert_eq!(key.user_id, "user_123");
396 assert_eq!(key.scopes, vec!["read"]);
397 }
398
399 #[tokio::test]
400 async fn test_validate_api_key() {
401 let store = Arc::new(InMemoryStore::new());
402 let manager = ApiKeyManager::new(store);
403
404 let key = manager
405 .generate("user_123", vec!["read".to_string()])
406 .await
407 .unwrap();
408 let validated = manager.validate(&key.key).await.unwrap();
409
410 assert!(validated.is_some());
411 assert_eq!(validated.unwrap().user_id, "user_123");
412 }
413
414 #[tokio::test]
415 async fn test_revoke_api_key() {
416 let store = Arc::new(InMemoryStore::new());
417 let manager = ApiKeyManager::new(store);
418
419 let key = manager
420 .generate("user_123", vec!["read".to_string()])
421 .await
422 .unwrap();
423 manager.revoke(&key.id).await.unwrap();
424
425 let result = manager.validate(&key.key).await;
426 assert!(matches!(result, Err(ApiKeyError::Revoked)));
427 }
428
429 #[tokio::test]
430 async fn test_has_scope() {
431 let store = Arc::new(InMemoryStore::new());
432 let manager = ApiKeyManager::new(store);
433
434 let key = manager
435 .generate("user_123", vec!["read".to_string(), "write".to_string()])
436 .await
437 .unwrap();
438
439 assert!(manager.has_scope(&key, "read"));
440 assert!(manager.has_scope(&key, "write"));
441 assert!(!manager.has_scope(&key, "admin"));
442 }
443}