axon/
api_keys.rs

1//! API Key Management — multi-tenant key registry for AxonServer.
2//!
3//! Manages named API keys with metadata, permissions, and per-key rate limits.
4//! Keys are stored in memory with optional file-backed persistence.
5//!
6//! Features:
7//!   - Create, revoke, list, and validate API keys
8//!   - Per-key metadata: name, role, created_at, last_used
9//!   - Per-key rate limit override
10//!   - Key rotation (revoke old, create new)
11//!
12//! Integration: replaces single auth_token check with multi-key validation.
13
14use std::collections::HashMap;
15use std::time::{SystemTime, UNIX_EPOCH};
16
17use serde::{Deserialize, Serialize};
18
19// ── Types ────────────────────────────────────────────────────────────────
20
21/// Role/permission level for an API key.
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
23#[serde(rename_all = "lowercase")]
24pub enum KeyRole {
25    /// Full access to all endpoints.
26    Admin,
27    /// Can deploy, estimate, read metrics — but not manage keys.
28    Operator,
29    /// Read-only: health, metrics, versions, logs.
30    ReadOnly,
31}
32
33impl KeyRole {
34    pub fn as_str(&self) -> &'static str {
35        match self {
36            KeyRole::Admin => "admin",
37            KeyRole::Operator => "operator",
38            KeyRole::ReadOnly => "readonly",
39        }
40    }
41
42    /// Whether this role can perform write operations (deploy, session writes, etc.).
43    pub fn can_write(&self) -> bool {
44        matches!(self, KeyRole::Admin | KeyRole::Operator)
45    }
46
47    /// Whether this role can manage API keys.
48    pub fn can_manage_keys(&self) -> bool {
49        matches!(self, KeyRole::Admin)
50    }
51}
52
53/// A registered API key with metadata.
54#[derive(Debug, Clone, Serialize)]
55pub struct ApiKey {
56    /// The key name (human-readable identifier).
57    pub name: String,
58    /// The secret token value.
59    #[serde(skip_serializing)]
60    pub token: String,
61    /// Permission role.
62    pub role: KeyRole,
63    /// Creation timestamp (Unix seconds).
64    pub created_at: u64,
65    /// Last used timestamp (Unix seconds), None if never used.
66    pub last_used: Option<u64>,
67    /// Optional per-key rate limit (requests per window).
68    pub rate_limit: Option<u32>,
69    /// Whether the key is active (not revoked).
70    pub active: bool,
71    /// Total requests made with this key.
72    pub request_count: u64,
73}
74
75/// Result of validating a token.
76#[derive(Debug, Clone)]
77pub struct ValidationResult {
78    pub valid: bool,
79    pub key_name: Option<String>,
80    pub role: Option<KeyRole>,
81    pub rate_limit: Option<u32>,
82}
83
84// ── Manager ──────────────────────────────────────────────────────────────
85
86/// API key registry and validator.
87pub struct ApiKeyManager {
88    /// Map of token → ApiKey.
89    keys: HashMap<String, ApiKey>,
90    /// Whether key management is enabled. If false, all auth is bypassed.
91    enabled: bool,
92}
93
94impl ApiKeyManager {
95    /// Create a new manager. If a master token is provided, it becomes the
96    /// initial admin key.
97    pub fn new(master_token: Option<&str>) -> Self {
98        let mut mgr = ApiKeyManager {
99            keys: HashMap::new(),
100            enabled: false,
101        };
102
103        if let Some(token) = master_token {
104            if !token.is_empty() {
105                mgr.create_key("master", token, KeyRole::Admin, None);
106                mgr.enabled = true;
107            }
108        }
109
110        mgr
111    }
112
113    /// Whether key management is enabled.
114    pub fn is_enabled(&self) -> bool {
115        self.enabled
116    }
117
118    /// Create a new API key. Returns true if created (false if token already exists).
119    pub fn create_key(&mut self, name: &str, token: &str, role: KeyRole, rate_limit: Option<u32>) -> bool {
120        if self.keys.contains_key(token) {
121            return false;
122        }
123
124        let key = ApiKey {
125            name: name.to_string(),
126            token: token.to_string(),
127            role,
128            created_at: now_secs(),
129            last_used: None,
130            rate_limit,
131            active: true,
132            request_count: 0,
133        };
134
135        self.keys.insert(token.to_string(), key);
136        if !self.enabled {
137            self.enabled = true;
138        }
139        true
140    }
141
142    /// Validate a token and record usage. Returns validation result.
143    pub fn validate(&mut self, token: &str) -> ValidationResult {
144        if !self.enabled {
145            return ValidationResult {
146                valid: true,
147                key_name: None,
148                role: None,
149                rate_limit: None,
150            };
151        }
152
153        match self.keys.get_mut(token) {
154            Some(key) if key.active => {
155                key.last_used = Some(now_secs());
156                key.request_count += 1;
157                ValidationResult {
158                    valid: true,
159                    key_name: Some(key.name.clone()),
160                    role: Some(key.role),
161                    rate_limit: key.rate_limit,
162                }
163            }
164            _ => ValidationResult {
165                valid: false,
166                key_name: None,
167                role: None,
168                rate_limit: None,
169            },
170        }
171    }
172
173    /// Validate without recording usage (peek).
174    pub fn peek(&self, token: &str) -> ValidationResult {
175        if !self.enabled {
176            return ValidationResult {
177                valid: true,
178                key_name: None,
179                role: None,
180                rate_limit: None,
181            };
182        }
183
184        match self.keys.get(token) {
185            Some(key) if key.active => ValidationResult {
186                valid: true,
187                key_name: Some(key.name.clone()),
188                role: Some(key.role),
189                rate_limit: key.rate_limit,
190            },
191            _ => ValidationResult {
192                valid: false,
193                key_name: None,
194                role: None,
195                rate_limit: None,
196            },
197        }
198    }
199
200    /// Revoke a key by token. Returns true if found and revoked.
201    pub fn revoke(&mut self, token: &str) -> bool {
202        match self.keys.get_mut(token) {
203            Some(key) if key.active => {
204                key.active = false;
205                true
206            }
207            _ => false,
208        }
209    }
210
211    /// Revoke a key by name. Returns true if found and revoked.
212    pub fn revoke_by_name(&mut self, name: &str) -> bool {
213        let token = self.keys.iter()
214            .find(|(_, k)| k.name == name && k.active)
215            .map(|(t, _)| t.clone());
216
217        match token {
218            Some(t) => self.revoke(&t),
219            None => false,
220        }
221    }
222
223    /// Rotate a key: revoke old token, create new one with same name/role.
224    /// Returns the new token if successful.
225    pub fn rotate(&mut self, old_token: &str, new_token: &str) -> Option<String> {
226        let (name, role, rate_limit) = match self.keys.get(old_token) {
227            Some(key) if key.active => (key.name.clone(), key.role, key.rate_limit),
228            _ => return None,
229        };
230
231        self.revoke(old_token);
232        if self.create_key(&name, new_token, role, rate_limit) {
233            Some(name)
234        } else {
235            None
236        }
237    }
238
239    /// List all keys (active and revoked). Tokens are masked.
240    pub fn list(&self) -> Vec<ApiKeySummary> {
241        let mut result: Vec<ApiKeySummary> = self.keys.values().map(|k| {
242            ApiKeySummary {
243                name: k.name.clone(),
244                role: k.role,
245                active: k.active,
246                created_at: k.created_at,
247                last_used: k.last_used,
248                rate_limit: k.rate_limit,
249                request_count: k.request_count,
250                token_prefix: mask_token(&k.token),
251            }
252        }).collect();
253        result.sort_by(|a, b| a.name.cmp(&b.name));
254        result
255    }
256
257    /// Number of active keys.
258    pub fn active_count(&self) -> usize {
259        self.keys.values().filter(|k| k.active).count()
260    }
261
262    /// Total keys (including revoked).
263    pub fn total_count(&self) -> usize {
264        self.keys.len()
265    }
266}
267
268/// Summary of an API key (token masked).
269#[derive(Debug, Clone, Serialize)]
270pub struct ApiKeySummary {
271    pub name: String,
272    pub role: KeyRole,
273    pub active: bool,
274    pub created_at: u64,
275    pub last_used: Option<u64>,
276    pub rate_limit: Option<u32>,
277    pub request_count: u64,
278    pub token_prefix: String,
279}
280
281// ── Helpers ──────────────────────────────────────────────────────────────
282
283fn now_secs() -> u64 {
284    SystemTime::now()
285        .duration_since(UNIX_EPOCH)
286        .unwrap_or_default()
287        .as_secs()
288}
289
290fn mask_token(token: &str) -> String {
291    if token.len() <= 4 {
292        "****".to_string()
293    } else {
294        format!("{}****", &token[..4])
295    }
296}
297
298// ── Tests ────────────────────────────────────────────────────────────────
299
300#[cfg(test)]
301mod tests {
302    use super::*;
303
304    #[test]
305    fn create_and_validate() {
306        let mut mgr = ApiKeyManager::new(None);
307        assert!(!mgr.is_enabled());
308
309        mgr.create_key("test-key", "tok_abc123", KeyRole::Operator, None);
310        assert!(mgr.is_enabled());
311
312        let result = mgr.validate("tok_abc123");
313        assert!(result.valid);
314        assert_eq!(result.key_name.unwrap(), "test-key");
315        assert_eq!(result.role.unwrap(), KeyRole::Operator);
316    }
317
318    #[test]
319    fn master_token_creates_admin() {
320        let mgr = ApiKeyManager::new(Some("master_secret"));
321        assert!(mgr.is_enabled());
322        assert_eq!(mgr.active_count(), 1);
323
324        let result = mgr.peek("master_secret");
325        assert!(result.valid);
326        assert_eq!(result.role.unwrap(), KeyRole::Admin);
327    }
328
329    #[test]
330    fn invalid_token_rejected() {
331        let mut mgr = ApiKeyManager::new(Some("valid"));
332        let result = mgr.validate("invalid");
333        assert!(!result.valid);
334        assert!(result.key_name.is_none());
335    }
336
337    #[test]
338    fn disabled_manager_allows_all() {
339        let mut mgr = ApiKeyManager::new(None);
340        let result = mgr.validate("anything");
341        assert!(result.valid);
342        assert!(result.role.is_none());
343    }
344
345    #[test]
346    fn revoke_key() {
347        let mut mgr = ApiKeyManager::new(None);
348        mgr.create_key("temp", "tok_temp", KeyRole::ReadOnly, None);
349        assert!(mgr.validate("tok_temp").valid);
350
351        assert!(mgr.revoke("tok_temp"));
352        assert!(!mgr.validate("tok_temp").valid);
353
354        // Revoke again returns false
355        assert!(!mgr.revoke("tok_temp"));
356    }
357
358    #[test]
359    fn revoke_by_name() {
360        let mut mgr = ApiKeyManager::new(None);
361        mgr.create_key("mykey", "tok_mykey", KeyRole::Operator, None);
362        assert!(mgr.revoke_by_name("mykey"));
363        assert!(!mgr.validate("tok_mykey").valid);
364        assert!(!mgr.revoke_by_name("mykey")); // already revoked
365    }
366
367    #[test]
368    fn rotate_key() {
369        let mut mgr = ApiKeyManager::new(None);
370        mgr.create_key("service", "old_token", KeyRole::Operator, Some(50));
371
372        let name = mgr.rotate("old_token", "new_token").unwrap();
373        assert_eq!(name, "service");
374
375        // Old token revoked
376        assert!(!mgr.validate("old_token").valid);
377
378        // New token active with same role
379        let result = mgr.validate("new_token");
380        assert!(result.valid);
381        assert_eq!(result.role.unwrap(), KeyRole::Operator);
382        assert_eq!(result.rate_limit, Some(50));
383    }
384
385    #[test]
386    fn rotate_invalid_token_fails() {
387        let mut mgr = ApiKeyManager::new(None);
388        assert!(mgr.rotate("nonexistent", "new").is_none());
389    }
390
391    #[test]
392    fn duplicate_token_rejected() {
393        let mut mgr = ApiKeyManager::new(None);
394        assert!(mgr.create_key("a", "same_token", KeyRole::Admin, None));
395        assert!(!mgr.create_key("b", "same_token", KeyRole::ReadOnly, None));
396    }
397
398    #[test]
399    fn list_keys_masked() {
400        let mut mgr = ApiKeyManager::new(None);
401        mgr.create_key("alpha", "tok_alpha_long", KeyRole::Admin, None);
402        mgr.create_key("beta", "tok_beta_long", KeyRole::ReadOnly, Some(10));
403
404        let list = mgr.list();
405        assert_eq!(list.len(), 2);
406        assert_eq!(list[0].name, "alpha");
407        assert_eq!(list[1].name, "beta");
408        assert!(list[0].token_prefix.contains("****"));
409        assert!(!list[0].token_prefix.contains("alpha_long"));
410    }
411
412    #[test]
413    fn request_count_increments() {
414        let mut mgr = ApiKeyManager::new(None);
415        mgr.create_key("counter", "tok_count", KeyRole::Operator, None);
416
417        for _ in 0..5 {
418            mgr.validate("tok_count");
419        }
420
421        let list = mgr.list();
422        let key = list.iter().find(|k| k.name == "counter").unwrap();
423        assert_eq!(key.request_count, 5);
424    }
425
426    #[test]
427    fn last_used_updated() {
428        let mut mgr = ApiKeyManager::new(None);
429        mgr.create_key("timed", "tok_timed", KeyRole::ReadOnly, None);
430
431        let list = mgr.list();
432        assert!(list[0].last_used.is_none());
433
434        mgr.validate("tok_timed");
435
436        let list = mgr.list();
437        assert!(list[0].last_used.is_some());
438        assert!(list[0].last_used.unwrap() > 1700000000);
439    }
440
441    #[test]
442    fn key_role_permissions() {
443        assert!(KeyRole::Admin.can_write());
444        assert!(KeyRole::Admin.can_manage_keys());
445
446        assert!(KeyRole::Operator.can_write());
447        assert!(!KeyRole::Operator.can_manage_keys());
448
449        assert!(!KeyRole::ReadOnly.can_write());
450        assert!(!KeyRole::ReadOnly.can_manage_keys());
451    }
452
453    #[test]
454    fn per_key_rate_limit() {
455        let mut mgr = ApiKeyManager::new(None);
456        mgr.create_key("limited", "tok_limited", KeyRole::Operator, Some(25));
457
458        let result = mgr.validate("tok_limited");
459        assert_eq!(result.rate_limit, Some(25));
460
461        mgr.create_key("unlimited", "tok_unlimited", KeyRole::Admin, None);
462        let result = mgr.validate("tok_unlimited");
463        assert_eq!(result.rate_limit, None);
464    }
465
466    #[test]
467    fn active_and_total_counts() {
468        let mut mgr = ApiKeyManager::new(None);
469        mgr.create_key("a", "t1", KeyRole::Admin, None);
470        mgr.create_key("b", "t2", KeyRole::Operator, None);
471        mgr.create_key("c", "t3", KeyRole::ReadOnly, None);
472
473        assert_eq!(mgr.active_count(), 3);
474        assert_eq!(mgr.total_count(), 3);
475
476        mgr.revoke("t2");
477        assert_eq!(mgr.active_count(), 2);
478        assert_eq!(mgr.total_count(), 3);
479    }
480
481    #[test]
482    fn mask_token_works() {
483        assert_eq!(mask_token("abcdef"), "abcd****");
484        assert_eq!(mask_token("ab"), "****");
485        assert_eq!(mask_token(""), "****");
486    }
487
488    #[test]
489    fn summary_serializes() {
490        let summary = ApiKeySummary {
491            name: "test".into(),
492            role: KeyRole::Admin,
493            active: true,
494            created_at: 1700000000,
495            last_used: None,
496            rate_limit: None,
497            request_count: 0,
498            token_prefix: "tok_****".into(),
499        };
500        let json = serde_json::to_string(&summary).unwrap();
501        assert!(json.contains("\"role\":\"admin\""));
502        assert!(json.contains("\"active\":true"));
503        assert!(json.contains("\"tok_****\""));
504    }
505
506    #[test]
507    fn peek_does_not_update_usage() {
508        let mut mgr = ApiKeyManager::new(None);
509        mgr.create_key("peek_test", "tok_peek", KeyRole::ReadOnly, None);
510
511        mgr.peek("tok_peek");
512        mgr.peek("tok_peek");
513
514        let list = mgr.list();
515        let key = list.iter().find(|k| k.name == "peek_test").unwrap();
516        assert_eq!(key.request_count, 0);
517        assert!(key.last_used.is_none());
518    }
519}
axon/api_keys.rs

axon/
api_keys.rs
