armature_auth/
api_key.rs

1//! API Key Management
2//!
3//! Provides API key generation, validation, and rotation with database injection.
4//!
5//! # Features
6//!
7//! - API key generation with secure random
8//! - Key validation and verification
9//! - Key rotation and expiration
10//! - Rate limiting per key
11//! - Scopes/permissions per key
12//! - Database-agnostic with DI
13//!
14//! # Usage
15//!
16//! ```no_run
17//! use armature_auth::api_key::*;
18//! use chrono::{DateTime, Utc};
19//! use std::sync::Arc;
20//!
21//! # async fn example() -> Result<(), Box<dyn std::error::Error>> {
22//! // Implement your database storage
23//! struct MyApiKeyStore;
24//!
25//! #[async_trait::async_trait]
26//! impl ApiKeyStore for MyApiKeyStore {
27//!     async fn save(&self, key: &ApiKey) -> Result<(), ApiKeyError> { Ok(()) }
28//!     async fn find_by_key(&self, key: &str) -> Result<Option<ApiKey>, ApiKeyError> { Ok(None) }
29//!     async fn find_by_id(&self, id: &str) -> Result<Option<ApiKey>, ApiKeyError> { Ok(None) }
30//!     async fn list_by_user(&self, user_id: &str) -> Result<Vec<ApiKey>, ApiKeyError> { Ok(vec![]) }
31//!     async fn revoke(&self, key_id: &str) -> Result<(), ApiKeyError> { Ok(()) }
32//!     async fn update_last_used(&self, key_id: &str, timestamp: DateTime<Utc>) -> Result<(), ApiKeyError> { Ok(()) }
33//! }
34//!
35//! // Inject store via DI container
36//! let store: Arc<dyn ApiKeyStore> = Arc::new(MyApiKeyStore);
37//! let manager = ApiKeyManager::new(store);
38//!
39//! // Generate API key
40//! let api_key = manager.generate("user_123", vec!["read".to_string(), "write".to_string()]).await?;
41//! println!("API Key: {}", api_key.key);
42//!
43//! // Validate API key
44//! if let Some(key) = manager.validate(&api_key.key).await? {
45//!     println!("Valid key for user: {}", key.user_id);
46//! }
47//! # Ok(())
48//! # }
49//! ```
50
51use 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/// API Key errors
61#[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/// API Key structure
83#[derive(Debug, Clone, Serialize, Deserialize)]
84pub struct ApiKey {
85    /// Unique key ID
86    pub id: String,
87
88    /// The actual API key (store hash, not plaintext!)
89    pub key: String,
90
91    /// User/account ID this key belongs to
92    pub user_id: String,
93
94    /// Key name/description
95    pub name: Option<String>,
96
97    /// Scopes/permissions
98    pub scopes: Vec<String>,
99
100    /// Creation timestamp
101    pub created_at: DateTime<Utc>,
102
103    /// Expiration timestamp
104    pub expires_at: Option<DateTime<Utc>>,
105
106    /// Last used timestamp
107    pub last_used_at: Option<DateTime<Utc>>,
108
109    /// Whether key is revoked
110    pub revoked: bool,
111
112    /// Rate limit (requests per minute)
113    pub rate_limit: Option<u32>,
114}
115
116/// API Key storage trait (implement with your database)
117///
118/// Users must provide their own implementation using their database of choice.
119#[async_trait]
120pub trait ApiKeyStore: Send + Sync {
121    /// Save or update an API key
122    async fn save(&self, key: &ApiKey) -> Result<(), ApiKeyError>;
123
124    /// Find API key by the key string (should query by hash)
125    async fn find_by_key(&self, key: &str) -> Result<Option<ApiKey>, ApiKeyError>;
126
127    /// Find API key by ID
128    async fn find_by_id(&self, id: &str) -> Result<Option<ApiKey>, ApiKeyError>;
129
130    /// List all keys for a user
131    async fn list_by_user(&self, user_id: &str) -> Result<Vec<ApiKey>, ApiKeyError>;
132
133    /// Revoke an API key
134    async fn revoke(&self, key_id: &str) -> Result<(), ApiKeyError>;
135
136    /// Update last used timestamp
137    async fn update_last_used(
138        &self,
139        key_id: &str,
140        timestamp: DateTime<Utc>,
141    ) -> Result<(), ApiKeyError>;
142}
143
144/// API Key Manager
145///
146/// Manages API key lifecycle with injected storage.
147pub struct ApiKeyManager {
148    store: Arc<dyn ApiKeyStore>,
149    key_prefix: String,
150    default_expiration: Option<Duration>,
151}
152
153impl ApiKeyManager {
154    /// Create new API key manager with injected store
155    ///
156    /// # Examples
157    ///
158    /// ```no_run
159    /// use armature_auth::api_key::*;
160    /// use chrono::{DateTime, Utc};
161    /// use std::sync::Arc;
162    ///
163    /// # struct MyStore;
164    /// # #[async_trait::async_trait]
165    /// # impl ApiKeyStore for MyStore {
166    /// #     async fn save(&self, key: &ApiKey) -> Result<(), ApiKeyError> { Ok(()) }
167    /// #     async fn find_by_key(&self, key: &str) -> Result<Option<ApiKey>, ApiKeyError> { Ok(None) }
168    /// #     async fn find_by_id(&self, id: &str) -> Result<Option<ApiKey>, ApiKeyError> { Ok(None) }
169    /// #     async fn list_by_user(&self, user_id: &str) -> Result<Vec<ApiKey>, ApiKeyError> { Ok(vec![]) }
170    /// #     async fn revoke(&self, key_id: &str) -> Result<(), ApiKeyError> { Ok(()) }
171    /// #     async fn update_last_used(&self, key_id: &str, timestamp: DateTime<Utc>) -> Result<(), ApiKeyError> { Ok(()) }
172    /// # }
173    /// let store: Arc<dyn ApiKeyStore> = Arc::new(MyStore);
174    /// let manager = ApiKeyManager::new(store);
175    /// ```
176    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    /// Set key prefix (default: "ak")
185    pub fn with_prefix(mut self, prefix: impl Into<String>) -> Self {
186        self.key_prefix = prefix.into();
187        self
188    }
189
190    /// Set default expiration duration
191    pub fn with_expiration(mut self, duration: Option<Duration>) -> Self {
192        self.default_expiration = duration;
193        self
194    }
195
196    /// Generate a new API key
197    ///
198    /// # Examples
199    ///
200    /// ```no_run
201    /// # use armature_auth::api_key::*;
202    /// # use std::sync::Arc;
203    /// # async fn example(manager: ApiKeyManager) -> Result<(), ApiKeyError> {
204    /// let key = manager.generate(
205    ///     "user_123",
206    ///     vec!["read".to_string(), "write".to_string()]
207    /// ).await?;
208    ///
209    /// println!("API Key: {}", key.key);
210    /// println!("Key ID: {}", key.id);
211    /// # Ok(())
212    /// # }
213    /// ```
214    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        // Store key (implementation should hash the key!)
239        self.store.save(&api_key).await?;
240
241        Ok(api_key)
242    }
243
244    /// Validate an API key
245    ///
246    /// Returns the API key if valid, None if invalid, or an error.
247    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        // Check if revoked
254        if api_key.revoked {
255            return Err(ApiKeyError::Revoked);
256        }
257
258        // Check if expired
259        if let Some(expires_at) = api_key.expires_at {
260            if Utc::now() > expires_at {
261                return Err(ApiKeyError::Expired);
262            }
263        }
264
265        // Update last used timestamp
266        self.store.update_last_used(&api_key.id, Utc::now()).await?;
267
268        Ok(Some(api_key))
269    }
270
271    /// Check if key has required scope
272    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    /// Revoke an API key
280    pub async fn revoke(&self, key_id: &str) -> Result<(), ApiKeyError> {
281        self.store.revoke(key_id).await
282    }
283
284    /// List all keys for a user
285    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    /// Rotate an API key (revoke old, generate new)
290    pub async fn rotate(&self, old_key_id: &str) -> Result<ApiKey, ApiKeyError> {
291        // Get old key
292        let old_key = self
293            .store
294            .find_by_id(old_key_id)
295            .await?
296            .ok_or(ApiKeyError::Invalid)?;
297
298        // Revoke old key
299        self.revoke(old_key_id).await?;
300
301        // Generate new key with same scopes
302        self.generate(&old_key.user_id, old_key.scopes).await
303    }
304
305    /// Generate random key string
306    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    /// Hash an API key (for storage)
313    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}