oxify_connect_vision/
access_control.rs

1//! Access control module for API key management, rate limiting, and quotas.
2//!
3//! This module provides:
4//! - API key generation and validation
5//! - Rate limiting (requests per minute/hour/day)
6//! - Usage quotas (total requests, bytes processed)
7//! - Permission-based access control
8
9use serde::{Deserialize, Serialize};
10use std::collections::HashMap;
11use std::sync::atomic::{AtomicU64, Ordering};
12use std::sync::{Arc, RwLock};
13use std::time::{SystemTime, UNIX_EPOCH};
14
15/// Permission levels for access control
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
17pub enum Permission {
18    /// Read-only access (process images)
19    Read,
20    /// Write access (manage API keys)
21    Write,
22    /// Admin access (all operations)
23    Admin,
24}
25
26impl Permission {
27    /// Check if this permission includes another permission
28    pub fn includes(&self, other: Permission) -> bool {
29        match self {
30            Permission::Admin => true,
31            Permission::Write => matches!(other, Permission::Read | Permission::Write),
32            Permission::Read => matches!(other, Permission::Read),
33        }
34    }
35}
36
37/// API key with associated permissions and quotas
38#[derive(Debug, Clone, Serialize, Deserialize)]
39pub struct ApiKey {
40    /// Unique key identifier
41    pub key: String,
42
43    /// Human-readable name
44    pub name: String,
45
46    /// Permissions granted to this key
47    pub permissions: Vec<Permission>,
48
49    /// Rate limit (requests per minute)
50    pub rate_limit_per_minute: Option<u64>,
51
52    /// Rate limit (requests per hour)
53    pub rate_limit_per_hour: Option<u64>,
54
55    /// Rate limit (requests per day)
56    pub rate_limit_per_day: Option<u64>,
57
58    /// Total quota (maximum total requests)
59    pub total_quota: Option<u64>,
60
61    /// Bytes quota (maximum bytes processed)
62    pub bytes_quota: Option<u64>,
63
64    /// Whether the key is active
65    pub active: bool,
66
67    /// Creation timestamp (Unix timestamp in seconds)
68    pub created_at: u64,
69
70    /// Expiration timestamp (Unix timestamp in seconds, None = never expires)
71    pub expires_at: Option<u64>,
72
73    /// Optional metadata
74    pub metadata: HashMap<String, String>,
75}
76
77impl ApiKey {
78    /// Create a new API key
79    pub fn new(name: impl Into<String>) -> Self {
80        Self {
81            key: generate_api_key(),
82            name: name.into(),
83            permissions: vec![Permission::Read],
84            rate_limit_per_minute: None,
85            rate_limit_per_hour: None,
86            rate_limit_per_day: None,
87            total_quota: None,
88            bytes_quota: None,
89            active: true,
90            created_at: current_timestamp(),
91            expires_at: None,
92            metadata: HashMap::new(),
93        }
94    }
95
96    /// Set permissions
97    pub fn with_permissions(mut self, permissions: Vec<Permission>) -> Self {
98        self.permissions = permissions;
99        self
100    }
101
102    /// Set rate limit per minute
103    pub fn with_rate_limit_per_minute(mut self, limit: u64) -> Self {
104        self.rate_limit_per_minute = Some(limit);
105        self
106    }
107
108    /// Set rate limit per hour
109    pub fn with_rate_limit_per_hour(mut self, limit: u64) -> Self {
110        self.rate_limit_per_hour = Some(limit);
111        self
112    }
113
114    /// Set rate limit per day
115    pub fn with_rate_limit_per_day(mut self, limit: u64) -> Self {
116        self.rate_limit_per_day = Some(limit);
117        self
118    }
119
120    /// Set total quota
121    pub fn with_total_quota(mut self, quota: u64) -> Self {
122        self.total_quota = Some(quota);
123        self
124    }
125
126    /// Set bytes quota
127    pub fn with_bytes_quota(mut self, quota: u64) -> Self {
128        self.bytes_quota = Some(quota);
129        self
130    }
131
132    /// Set expiration time
133    pub fn with_expiration(mut self, expires_at: u64) -> Self {
134        self.expires_at = Some(expires_at);
135        self
136    }
137
138    /// Add metadata
139    pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
140        self.metadata.insert(key.into(), value.into());
141        self
142    }
143
144    /// Check if the key has a specific permission
145    pub fn has_permission(&self, permission: Permission) -> bool {
146        self.permissions.iter().any(|p| p.includes(permission))
147    }
148
149    /// Check if the key is expired
150    pub fn is_expired(&self) -> bool {
151        if let Some(expires_at) = self.expires_at {
152            current_timestamp() >= expires_at
153        } else {
154            false
155        }
156    }
157
158    /// Check if the key is valid (active and not expired)
159    pub fn is_valid(&self) -> bool {
160        self.active && !self.is_expired()
161    }
162}
163
164/// Generate a random API key
165fn generate_api_key() -> String {
166    use std::collections::hash_map::RandomState;
167    use std::hash::{BuildHasher, Hasher};
168
169    let mut hasher = RandomState::new().build_hasher();
170    let timestamp = current_timestamp();
171    hasher.write_u64(timestamp);
172
173    // Generate a pseudo-random key
174    let hash = hasher.finish();
175    format!("oxify_{:016x}{:016x}", hash, timestamp)
176}
177
178/// Get current timestamp in seconds
179fn current_timestamp() -> u64 {
180    SystemTime::now()
181        .duration_since(UNIX_EPOCH)
182        .unwrap()
183        .as_secs()
184}
185
186/// Usage statistics for an API key
187#[derive(Debug, Clone, Default)]
188pub struct UsageStats {
189    /// Total requests made
190    pub total_requests: Arc<AtomicU64>,
191
192    /// Total bytes processed
193    pub total_bytes: Arc<AtomicU64>,
194
195    /// Requests in the current minute
196    pub requests_this_minute: Arc<AtomicU64>,
197
198    /// Requests in the current hour
199    pub requests_this_hour: Arc<AtomicU64>,
200
201    /// Requests in the current day
202    pub requests_this_day: Arc<AtomicU64>,
203
204    /// Timestamp of last request
205    pub last_request_at: Arc<AtomicU64>,
206
207    /// Timestamp for minute window
208    pub minute_window_start: Arc<AtomicU64>,
209
210    /// Timestamp for hour window
211    pub hour_window_start: Arc<AtomicU64>,
212
213    /// Timestamp for day window
214    pub day_window_start: Arc<AtomicU64>,
215}
216
217impl UsageStats {
218    /// Create new usage stats
219    pub fn new() -> Self {
220        let now = current_timestamp();
221        Self {
222            total_requests: Arc::new(AtomicU64::new(0)),
223            total_bytes: Arc::new(AtomicU64::new(0)),
224            requests_this_minute: Arc::new(AtomicU64::new(0)),
225            requests_this_hour: Arc::new(AtomicU64::new(0)),
226            requests_this_day: Arc::new(AtomicU64::new(0)),
227            last_request_at: Arc::new(AtomicU64::new(0)),
228            minute_window_start: Arc::new(AtomicU64::new(now)),
229            hour_window_start: Arc::new(AtomicU64::new(now)),
230            day_window_start: Arc::new(AtomicU64::new(now)),
231        }
232    }
233
234    /// Record a request
235    pub fn record_request(&self, bytes: u64) {
236        let now = current_timestamp();
237
238        // Update total counters
239        self.total_requests.fetch_add(1, Ordering::Relaxed);
240        self.total_bytes.fetch_add(bytes, Ordering::Relaxed);
241        self.last_request_at.store(now, Ordering::Relaxed);
242
243        // Reset windows if needed
244        self.reset_windows_if_needed(now);
245
246        // Increment window counters
247        self.requests_this_minute.fetch_add(1, Ordering::Relaxed);
248        self.requests_this_hour.fetch_add(1, Ordering::Relaxed);
249        self.requests_this_day.fetch_add(1, Ordering::Relaxed);
250    }
251
252    /// Reset windows if they've expired
253    fn reset_windows_if_needed(&self, now: u64) {
254        // Reset minute window
255        let minute_start = self.minute_window_start.load(Ordering::Relaxed);
256        if now - minute_start >= 60 {
257            self.requests_this_minute.store(0, Ordering::Relaxed);
258            self.minute_window_start.store(now, Ordering::Relaxed);
259        }
260
261        // Reset hour window
262        let hour_start = self.hour_window_start.load(Ordering::Relaxed);
263        if now - hour_start >= 3600 {
264            self.requests_this_hour.store(0, Ordering::Relaxed);
265            self.hour_window_start.store(now, Ordering::Relaxed);
266        }
267
268        // Reset day window
269        let day_start = self.day_window_start.load(Ordering::Relaxed);
270        if now - day_start >= 86400 {
271            self.requests_this_day.store(0, Ordering::Relaxed);
272            self.day_window_start.store(now, Ordering::Relaxed);
273        }
274    }
275
276    /// Get current statistics
277    pub fn get(&self) -> UsageSnapshot {
278        UsageSnapshot {
279            total_requests: self.total_requests.load(Ordering::Relaxed),
280            total_bytes: self.total_bytes.load(Ordering::Relaxed),
281            requests_this_minute: self.requests_this_minute.load(Ordering::Relaxed),
282            requests_this_hour: self.requests_this_hour.load(Ordering::Relaxed),
283            requests_this_day: self.requests_this_day.load(Ordering::Relaxed),
284            last_request_at: self.last_request_at.load(Ordering::Relaxed),
285        }
286    }
287}
288
289/// Snapshot of usage statistics
290#[derive(Debug, Clone, Serialize, Deserialize)]
291pub struct UsageSnapshot {
292    pub total_requests: u64,
293    pub total_bytes: u64,
294    pub requests_this_minute: u64,
295    pub requests_this_hour: u64,
296    pub requests_this_day: u64,
297    pub last_request_at: u64,
298}
299
300/// Access control error
301#[derive(Debug, Clone, thiserror::Error)]
302pub enum AccessError {
303    /// Invalid API key
304    #[error("Invalid API key")]
305    InvalidKey,
306
307    /// API key expired
308    #[error("API key expired")]
309    Expired,
310
311    /// API key inactive
312    #[error("API key inactive")]
313    Inactive,
314
315    /// Insufficient permissions
316    #[error("Insufficient permissions: {0}")]
317    InsufficientPermissions(String),
318
319    /// Rate limit exceeded
320    #[error("Rate limit exceeded: {0}")]
321    RateLimitExceeded(String),
322
323    /// Quota exceeded
324    #[error("Quota exceeded: {0}")]
325    QuotaExceeded(String),
326}
327
328/// Access controller for managing API keys and enforcing limits
329pub struct AccessController {
330    keys: Arc<RwLock<HashMap<String, ApiKey>>>,
331    usage: Arc<RwLock<HashMap<String, UsageStats>>>,
332}
333
334impl AccessController {
335    /// Create a new access controller
336    pub fn new() -> Self {
337        Self {
338            keys: Arc::new(RwLock::new(HashMap::new())),
339            usage: Arc::new(RwLock::new(HashMap::new())),
340        }
341    }
342
343    /// Create a new API key
344    pub fn create_key(&self, name: impl Into<String>) -> ApiKey {
345        let key = ApiKey::new(name);
346        let mut keys = self.keys.write().unwrap();
347        let mut usage = self.usage.write().unwrap();
348
349        keys.insert(key.key.clone(), key.clone());
350        usage.insert(key.key.clone(), UsageStats::new());
351
352        key
353    }
354
355    /// Create a custom API key
356    pub fn create_custom_key(&self, key: ApiKey) -> String {
357        let key_str = key.key.clone();
358        let mut keys = self.keys.write().unwrap();
359        let mut usage = self.usage.write().unwrap();
360
361        keys.insert(key_str.clone(), key);
362        usage.insert(key_str.clone(), UsageStats::new());
363
364        key_str
365    }
366
367    /// Get an API key
368    pub fn get_key(&self, key: &str) -> Option<ApiKey> {
369        self.keys.read().unwrap().get(key).cloned()
370    }
371
372    /// Revoke an API key
373    pub fn revoke_key(&self, key: &str) -> bool {
374        if let Some(api_key) = self.keys.write().unwrap().get_mut(key) {
375            api_key.active = false;
376            true
377        } else {
378            false
379        }
380    }
381
382    /// Delete an API key
383    pub fn delete_key(&self, key: &str) -> bool {
384        let mut keys = self.keys.write().unwrap();
385        let mut usage = self.usage.write().unwrap();
386
387        keys.remove(key).is_some() || usage.remove(key).is_some()
388    }
389
390    /// Validate an API key and check permissions
391    pub fn validate(
392        &self,
393        key: &str,
394        permission: Permission,
395        bytes: u64,
396    ) -> Result<(), AccessError> {
397        // Get the API key
398        let api_key = self.get_key(key).ok_or(AccessError::InvalidKey)?;
399
400        // Check if key is valid
401        if !api_key.active {
402            return Err(AccessError::Inactive);
403        }
404
405        if api_key.is_expired() {
406            return Err(AccessError::Expired);
407        }
408
409        // Check permissions
410        if !api_key.has_permission(permission) {
411            return Err(AccessError::InsufficientPermissions(format!(
412                "Required: {:?}",
413                permission
414            )));
415        }
416
417        // Get usage stats
418        let usage = self.usage.read().unwrap();
419        let stats = usage.get(key).ok_or(AccessError::InvalidKey)?;
420
421        // Reset windows if needed
422        stats.reset_windows_if_needed(current_timestamp());
423
424        // Check rate limits
425        if let Some(limit) = api_key.rate_limit_per_minute {
426            let current = stats.requests_this_minute.load(Ordering::Relaxed);
427            if current >= limit {
428                return Err(AccessError::RateLimitExceeded(format!(
429                    "{} requests per minute",
430                    limit
431                )));
432            }
433        }
434
435        if let Some(limit) = api_key.rate_limit_per_hour {
436            let current = stats.requests_this_hour.load(Ordering::Relaxed);
437            if current >= limit {
438                return Err(AccessError::RateLimitExceeded(format!(
439                    "{} requests per hour",
440                    limit
441                )));
442            }
443        }
444
445        if let Some(limit) = api_key.rate_limit_per_day {
446            let current = stats.requests_this_day.load(Ordering::Relaxed);
447            if current >= limit {
448                return Err(AccessError::RateLimitExceeded(format!(
449                    "{} requests per day",
450                    limit
451                )));
452            }
453        }
454
455        // Check total quota
456        if let Some(quota) = api_key.total_quota {
457            let current = stats.total_requests.load(Ordering::Relaxed);
458            if current >= quota {
459                return Err(AccessError::QuotaExceeded(format!(
460                    "Total quota of {} requests",
461                    quota
462                )));
463            }
464        }
465
466        // Check bytes quota
467        if let Some(quota) = api_key.bytes_quota {
468            let current = stats.total_bytes.load(Ordering::Relaxed);
469            if current + bytes > quota {
470                return Err(AccessError::QuotaExceeded(format!(
471                    "Bytes quota of {} bytes",
472                    quota
473                )));
474            }
475        }
476
477        // Record the request
478        stats.record_request(bytes);
479
480        Ok(())
481    }
482
483    /// Get usage statistics for a key
484    pub fn get_usage(&self, key: &str) -> Option<UsageSnapshot> {
485        self.usage.read().unwrap().get(key).map(|stats| stats.get())
486    }
487
488    /// List all API keys
489    pub fn list_keys(&self) -> Vec<ApiKey> {
490        self.keys.read().unwrap().values().cloned().collect()
491    }
492
493    /// Get statistics for all keys
494    pub fn get_all_usage(&self) -> HashMap<String, UsageSnapshot> {
495        self.usage
496            .read()
497            .unwrap()
498            .iter()
499            .map(|(key, stats)| (key.clone(), stats.get()))
500            .collect()
501    }
502}
503
504impl Default for AccessController {
505    fn default() -> Self {
506        Self::new()
507    }
508}
509
510#[cfg(test)]
511mod tests {
512    use super::*;
513
514    #[test]
515    fn test_permission_includes() {
516        assert!(Permission::Admin.includes(Permission::Read));
517        assert!(Permission::Admin.includes(Permission::Write));
518        assert!(Permission::Admin.includes(Permission::Admin));
519
520        assert!(Permission::Write.includes(Permission::Read));
521        assert!(Permission::Write.includes(Permission::Write));
522        assert!(!Permission::Write.includes(Permission::Admin));
523
524        assert!(Permission::Read.includes(Permission::Read));
525        assert!(!Permission::Read.includes(Permission::Write));
526        assert!(!Permission::Read.includes(Permission::Admin));
527    }
528
529    #[test]
530    fn test_api_key_creation() {
531        let key = ApiKey::new("test-key");
532        assert_eq!(key.name, "test-key");
533        assert!(key.key.starts_with("oxify_"));
534        assert!(key.active);
535        assert!(!key.is_expired());
536    }
537
538    #[test]
539    fn test_api_key_with_permissions() {
540        let key = ApiKey::new("test").with_permissions(vec![Permission::Read, Permission::Write]);
541
542        assert!(key.has_permission(Permission::Read));
543        assert!(key.has_permission(Permission::Write));
544        assert!(!key.has_permission(Permission::Admin));
545    }
546
547    #[test]
548    fn test_api_key_with_rate_limits() {
549        let key = ApiKey::new("test")
550            .with_rate_limit_per_minute(100)
551            .with_rate_limit_per_hour(1000);
552
553        assert_eq!(key.rate_limit_per_minute, Some(100));
554        assert_eq!(key.rate_limit_per_hour, Some(1000));
555    }
556
557    #[test]
558    fn test_api_key_with_quotas() {
559        let key = ApiKey::new("test")
560            .with_total_quota(5000)
561            .with_bytes_quota(1_000_000);
562
563        assert_eq!(key.total_quota, Some(5000));
564        assert_eq!(key.bytes_quota, Some(1_000_000));
565    }
566
567    #[test]
568    fn test_api_key_expiration() {
569        let past = current_timestamp() - 3600;
570        let future = current_timestamp() + 3600;
571
572        let expired_key = ApiKey::new("test").with_expiration(past);
573        assert!(expired_key.is_expired());
574
575        let valid_key = ApiKey::new("test").with_expiration(future);
576        assert!(!valid_key.is_expired());
577    }
578
579    #[test]
580    fn test_usage_stats_record_request() {
581        let stats = UsageStats::new();
582        stats.record_request(1024);
583
584        let snapshot = stats.get();
585        assert_eq!(snapshot.total_requests, 1);
586        assert_eq!(snapshot.total_bytes, 1024);
587        assert_eq!(snapshot.requests_this_minute, 1);
588    }
589
590    #[test]
591    fn test_access_controller_create_key() {
592        let controller = AccessController::new();
593        let key = controller.create_key("test-key");
594
595        assert_eq!(key.name, "test-key");
596        assert!(controller.get_key(&key.key).is_some());
597    }
598
599    #[test]
600    fn test_access_controller_revoke_key() {
601        let controller = AccessController::new();
602        let key = controller.create_key("test");
603
604        assert!(controller.revoke_key(&key.key));
605
606        let revoked = controller.get_key(&key.key).unwrap();
607        assert!(!revoked.active);
608    }
609
610    #[test]
611    fn test_access_controller_delete_key() {
612        let controller = AccessController::new();
613        let key = controller.create_key("test");
614
615        assert!(controller.delete_key(&key.key));
616        assert!(controller.get_key(&key.key).is_none());
617    }
618
619    #[test]
620    fn test_access_controller_validate_invalid_key() {
621        let controller = AccessController::new();
622        let result = controller.validate("invalid-key", Permission::Read, 0);
623
624        assert!(result.is_err());
625        assert!(matches!(result.unwrap_err(), AccessError::InvalidKey));
626    }
627
628    #[test]
629    fn test_access_controller_validate_inactive_key() {
630        let controller = AccessController::new();
631        let key = controller.create_key("test");
632        controller.revoke_key(&key.key);
633
634        let result = controller.validate(&key.key, Permission::Read, 0);
635        assert!(result.is_err());
636        assert!(matches!(result.unwrap_err(), AccessError::Inactive));
637    }
638
639    #[test]
640    fn test_access_controller_validate_permissions() {
641        let controller = AccessController::new();
642        let key = ApiKey::new("test").with_permissions(vec![Permission::Read]);
643        let key_str = controller.create_custom_key(key);
644
645        let result = controller.validate(&key_str, Permission::Write, 0);
646        assert!(result.is_err());
647        assert!(matches!(
648            result.unwrap_err(),
649            AccessError::InsufficientPermissions(_)
650        ));
651    }
652
653    #[test]
654    fn test_access_controller_validate_rate_limit() {
655        let controller = AccessController::new();
656        let key = ApiKey::new("test").with_rate_limit_per_minute(1);
657        let key_str = controller.create_custom_key(key);
658
659        // First request should succeed
660        assert!(controller.validate(&key_str, Permission::Read, 0).is_ok());
661
662        // Second request should fail
663        let result = controller.validate(&key_str, Permission::Read, 0);
664        assert!(result.is_err());
665        assert!(matches!(
666            result.unwrap_err(),
667            AccessError::RateLimitExceeded(_)
668        ));
669    }
670
671    #[test]
672    fn test_access_controller_validate_quota() {
673        let controller = AccessController::new();
674        let key = ApiKey::new("test").with_total_quota(1);
675        let key_str = controller.create_custom_key(key);
676
677        // First request should succeed
678        assert!(controller.validate(&key_str, Permission::Read, 0).is_ok());
679
680        // Second request should fail
681        let result = controller.validate(&key_str, Permission::Read, 0);
682        assert!(result.is_err());
683        assert!(matches!(result.unwrap_err(), AccessError::QuotaExceeded(_)));
684    }
685
686    #[test]
687    fn test_access_controller_get_usage() {
688        let controller = AccessController::new();
689        let key = controller.create_key("test");
690
691        controller
692            .validate(&key.key, Permission::Read, 1024)
693            .unwrap();
694
695        let usage = controller.get_usage(&key.key).unwrap();
696        assert_eq!(usage.total_requests, 1);
697        assert_eq!(usage.total_bytes, 1024);
698    }
699
700    #[test]
701    fn test_access_controller_list_keys() {
702        let controller = AccessController::new();
703        controller.create_key("key1");
704        controller.create_key("key2");
705
706        let keys = controller.list_keys();
707        assert_eq!(keys.len(), 2);
708    }
709
710    #[test]
711    fn test_generate_api_key_format() {
712        let key = generate_api_key();
713        assert!(key.starts_with("oxify_"));
714        assert!(key.len() > 10);
715    }
716}
oxify_connect_vision/access_control.rs

oxify_connect_vision/
access_control.rs
