litellm_rs/core/keys/
repository.rs

1//! API Key Repository
2//!
3//! This module provides the storage abstraction for API keys.
4
5use super::types::{KeyStatus, ManagedApiKey, UpdateKeyConfig, UsageRecord};
6use crate::utils::error::gateway_error::{GatewayError, Result};
7use async_trait::async_trait;
8use chrono::Utc;
9use dashmap::DashMap;
10use std::sync::Arc;
11use uuid::Uuid;
12
13/// Repository trait for API key storage operations
14#[async_trait]
15pub trait KeyRepository: Send + Sync {
16    /// Store a new API key
17    async fn create(&self, key: ManagedApiKey) -> Result<ManagedApiKey>;
18
19    /// Find a key by ID
20    async fn find_by_id(&self, id: Uuid) -> Result<Option<ManagedApiKey>>;
21
22    /// Find a key by its hash
23    async fn find_by_hash(&self, key_hash: &str) -> Result<Option<ManagedApiKey>>;
24
25    /// Update a key
26    async fn update(&self, id: Uuid, config: UpdateKeyConfig) -> Result<ManagedApiKey>;
27
28    /// Update key status
29    async fn update_status(&self, id: Uuid, status: KeyStatus) -> Result<()>;
30
31    /// Update last used timestamp
32    async fn update_last_used(&self, id: Uuid) -> Result<()>;
33
34    /// Update usage statistics
35    async fn update_usage(&self, id: Uuid, record: UsageRecord) -> Result<()>;
36
37    /// List keys by user ID
38    async fn list_by_user(&self, user_id: Uuid) -> Result<Vec<ManagedApiKey>>;
39
40    /// List keys by team ID
41    async fn list_by_team(&self, team_id: Uuid) -> Result<Vec<ManagedApiKey>>;
42
43    /// List all keys with optional filters
44    async fn list_all(
45        &self,
46        status: Option<KeyStatus>,
47        limit: Option<usize>,
48        offset: Option<usize>,
49    ) -> Result<Vec<ManagedApiKey>>;
50
51    /// Delete a key permanently
52    async fn delete(&self, id: Uuid) -> Result<()>;
53
54    /// Delete expired keys
55    async fn delete_expired(&self) -> Result<u64>;
56
57    /// Count total keys
58    async fn count(&self, status: Option<KeyStatus>) -> Result<u64>;
59}
60
61/// In-memory implementation of KeyRepository for testing and development
62#[derive(Debug, Default)]
63pub struct InMemoryKeyRepository {
64    /// Storage for keys by ID
65    keys_by_id: DashMap<Uuid, ManagedApiKey>,
66    /// Index for looking up keys by hash
67    hash_index: DashMap<String, Uuid>,
68}
69
70impl InMemoryKeyRepository {
71    /// Create a new in-memory repository
72    pub fn new() -> Self {
73        Self {
74            keys_by_id: DashMap::new(),
75            hash_index: DashMap::new(),
76        }
77    }
78
79    /// Create a new in-memory repository wrapped in Arc
80    pub fn new_arc() -> Arc<Self> {
81        Arc::new(Self::new())
82    }
83}
84
85#[async_trait]
86impl KeyRepository for InMemoryKeyRepository {
87    async fn create(&self, key: ManagedApiKey) -> Result<ManagedApiKey> {
88        // Check if hash already exists
89        if self.hash_index.contains_key(&key.key_hash) {
90            return Err(GatewayError::conflict("API key already exists"));
91        }
92
93        // Store the key
94        self.hash_index.insert(key.key_hash.clone(), key.id);
95        self.keys_by_id.insert(key.id, key.clone());
96
97        Ok(key)
98    }
99
100    async fn find_by_id(&self, id: Uuid) -> Result<Option<ManagedApiKey>> {
101        Ok(self.keys_by_id.get(&id).map(|r| r.value().clone()))
102    }
103
104    async fn find_by_hash(&self, key_hash: &str) -> Result<Option<ManagedApiKey>> {
105        if let Some(id) = self.hash_index.get(key_hash) {
106            return self.find_by_id(*id).await;
107        }
108        Ok(None)
109    }
110
111    async fn update(&self, id: Uuid, config: UpdateKeyConfig) -> Result<ManagedApiKey> {
112        let mut entry = self
113            .keys_by_id
114            .get_mut(&id)
115            .ok_or_else(|| GatewayError::not_found("API key not found"))?;
116
117        let key = entry.value_mut();
118
119        if let Some(name) = config.name {
120            key.name = name;
121        }
122        if let Some(description) = config.description {
123            key.description = description;
124        }
125        if let Some(permissions) = config.permissions {
126            key.permissions = permissions;
127        }
128        if let Some(rate_limits) = config.rate_limits {
129            key.rate_limits = rate_limits;
130        }
131        if let Some(budget_id) = config.budget_id {
132            key.budget_id = budget_id;
133        }
134        if let Some(expires_at) = config.expires_at {
135            key.expires_at = expires_at;
136        }
137        if let Some(metadata) = config.metadata {
138            key.metadata = metadata;
139        }
140
141        key.updated_at = Utc::now();
142
143        Ok(key.clone())
144    }
145
146    async fn update_status(&self, id: Uuid, status: KeyStatus) -> Result<()> {
147        let mut entry = self
148            .keys_by_id
149            .get_mut(&id)
150            .ok_or_else(|| GatewayError::not_found("API key not found"))?;
151
152        entry.value_mut().status = status;
153        entry.value_mut().updated_at = Utc::now();
154
155        Ok(())
156    }
157
158    async fn update_last_used(&self, id: Uuid) -> Result<()> {
159        if let Some(mut entry) = self.keys_by_id.get_mut(&id) {
160            entry.value_mut().last_used_at = Some(Utc::now());
161        }
162        Ok(())
163    }
164
165    async fn update_usage(&self, id: Uuid, record: UsageRecord) -> Result<()> {
166        if let Some(mut entry) = self.keys_by_id.get_mut(&id) {
167            let key = entry.value_mut();
168            key.usage_stats.reset_daily_if_needed();
169            key.usage_stats.record_usage_record(&record);
170            key.last_used_at = Some(Utc::now());
171        }
172        Ok(())
173    }
174
175    async fn list_by_user(&self, user_id: Uuid) -> Result<Vec<ManagedApiKey>> {
176        let keys: Vec<ManagedApiKey> = self
177            .keys_by_id
178            .iter()
179            .filter(|r| r.value().user_id == Some(user_id))
180            .map(|r| r.value().clone())
181            .collect();
182        Ok(keys)
183    }
184
185    async fn list_by_team(&self, team_id: Uuid) -> Result<Vec<ManagedApiKey>> {
186        let keys: Vec<ManagedApiKey> = self
187            .keys_by_id
188            .iter()
189            .filter(|r| r.value().team_id == Some(team_id))
190            .map(|r| r.value().clone())
191            .collect();
192        Ok(keys)
193    }
194
195    async fn list_all(
196        &self,
197        status: Option<KeyStatus>,
198        limit: Option<usize>,
199        offset: Option<usize>,
200    ) -> Result<Vec<ManagedApiKey>> {
201        let mut keys: Vec<ManagedApiKey> = self
202            .keys_by_id
203            .iter()
204            .filter(|r| {
205                if let Some(s) = status {
206                    r.value().effective_status() == s
207                } else {
208                    true
209                }
210            })
211            .map(|r| r.value().clone())
212            .collect();
213
214        // Sort by created_at descending
215        keys.sort_by_key(|b| std::cmp::Reverse(b.created_at));
216
217        // Apply pagination
218        let offset = offset.unwrap_or(0);
219        let limit = limit.unwrap_or(100);
220
221        Ok(keys.into_iter().skip(offset).take(limit).collect())
222    }
223
224    async fn delete(&self, id: Uuid) -> Result<()> {
225        if let Some((_, key)) = self.keys_by_id.remove(&id) {
226            self.hash_index.remove(&key.key_hash);
227        }
228        Ok(())
229    }
230
231    async fn delete_expired(&self) -> Result<u64> {
232        let now = Utc::now();
233        let mut deleted = 0u64;
234
235        let expired_ids: Vec<Uuid> = self
236            .keys_by_id
237            .iter()
238            .filter(|r| {
239                if let Some(expires_at) = r.value().expires_at {
240                    expires_at < now
241                } else {
242                    false
243                }
244            })
245            .map(|r| *r.key())
246            .collect();
247
248        for id in expired_ids {
249            if let Some((_, key)) = self.keys_by_id.remove(&id) {
250                self.hash_index.remove(&key.key_hash);
251                deleted += 1;
252            }
253        }
254
255        Ok(deleted)
256    }
257
258    async fn count(&self, status: Option<KeyStatus>) -> Result<u64> {
259        let count = self
260            .keys_by_id
261            .iter()
262            .filter(|r| {
263                if let Some(s) = status {
264                    r.value().effective_status() == s
265                } else {
266                    true
267                }
268            })
269            .count();
270        Ok(count as u64)
271    }
272}
273
274#[cfg(test)]
275mod repository_tests {
276    use super::*;
277    use crate::core::keys::types::{KeyPermissions, KeyRateLimits, KeyUsageStats};
278
279    fn create_test_key() -> ManagedApiKey {
280        ManagedApiKey {
281            id: Uuid::new_v4(),
282            key_hash: format!("hash_{}", Uuid::new_v4()),
283            key_prefix: "gw-test...1234".to_string(),
284            name: "Test Key".to_string(),
285            description: None,
286            user_id: None,
287            team_id: None,
288            budget_id: None,
289            permissions: KeyPermissions::default(),
290            rate_limits: KeyRateLimits::default(),
291            status: KeyStatus::Active,
292            expires_at: None,
293            created_at: Utc::now(),
294            updated_at: Utc::now(),
295            last_used_at: None,
296            usage_stats: KeyUsageStats::new(),
297            metadata: serde_json::Value::Null,
298        }
299    }
300
301    #[tokio::test]
302    async fn test_create_key() {
303        let repo = InMemoryKeyRepository::new();
304        let key = create_test_key();
305        let id = key.id;
306
307        let created = repo.create(key).await.unwrap();
308        assert_eq!(created.id, id);
309
310        let found = repo.find_by_id(id).await.unwrap();
311        assert!(found.is_some());
312    }
313
314    #[tokio::test]
315    async fn test_find_by_hash() {
316        let repo = InMemoryKeyRepository::new();
317        let key = create_test_key();
318        let hash = key.key_hash.clone();
319
320        repo.create(key).await.unwrap();
321
322        let found = repo.find_by_hash(&hash).await.unwrap();
323        assert!(found.is_some());
324        assert_eq!(found.unwrap().key_hash, hash);
325    }
326
327    #[tokio::test]
328    async fn test_update_key() {
329        let repo = InMemoryKeyRepository::new();
330        let key = create_test_key();
331        let id = key.id;
332
333        repo.create(key).await.unwrap();
334
335        let config = UpdateKeyConfig {
336            name: Some("Updated Name".to_string()),
337            ..Default::default()
338        };
339
340        let updated = repo.update(id, config).await.unwrap();
341        assert_eq!(updated.name, "Updated Name");
342    }
343
344    #[tokio::test]
345    async fn test_update_status() {
346        let repo = InMemoryKeyRepository::new();
347        let key = create_test_key();
348        let id = key.id;
349
350        repo.create(key).await.unwrap();
351        repo.update_status(id, KeyStatus::Revoked).await.unwrap();
352
353        let found = repo.find_by_id(id).await.unwrap().unwrap();
354        assert_eq!(found.status, KeyStatus::Revoked);
355    }
356
357    #[tokio::test]
358    async fn test_list_by_user() {
359        let repo = InMemoryKeyRepository::new();
360        let user_id = Uuid::new_v4();
361
362        let mut key1 = create_test_key();
363        key1.user_id = Some(user_id);
364        repo.create(key1).await.unwrap();
365
366        let mut key2 = create_test_key();
367        key2.user_id = Some(user_id);
368        repo.create(key2).await.unwrap();
369
370        let mut key3 = create_test_key();
371        key3.user_id = Some(Uuid::new_v4());
372        repo.create(key3).await.unwrap();
373
374        let keys = repo.list_by_user(user_id).await.unwrap();
375        assert_eq!(keys.len(), 2);
376    }
377
378    #[tokio::test]
379    async fn test_delete_key() {
380        let repo = InMemoryKeyRepository::new();
381        let key = create_test_key();
382        let id = key.id;
383        let hash = key.key_hash.clone();
384
385        repo.create(key).await.unwrap();
386        repo.delete(id).await.unwrap();
387
388        assert!(repo.find_by_id(id).await.unwrap().is_none());
389        assert!(repo.find_by_hash(&hash).await.unwrap().is_none());
390    }
391
392    #[tokio::test]
393    async fn test_delete_expired() {
394        let repo = InMemoryKeyRepository::new();
395
396        // Create an expired key
397        let mut key1 = create_test_key();
398        key1.expires_at = Some(Utc::now() - chrono::Duration::hours(1));
399        repo.create(key1).await.unwrap();
400
401        // Create a valid key
402        let key2 = create_test_key();
403        let valid_id = key2.id;
404        repo.create(key2).await.unwrap();
405
406        let deleted = repo.delete_expired().await.unwrap();
407        assert_eq!(deleted, 1);
408
409        // Valid key should still exist
410        assert!(repo.find_by_id(valid_id).await.unwrap().is_some());
411    }
412
413    #[tokio::test]
414    async fn test_count() {
415        let repo = InMemoryKeyRepository::new();
416
417        repo.create(create_test_key()).await.unwrap();
418        repo.create(create_test_key()).await.unwrap();
419
420        let count = repo.count(None).await.unwrap();
421        assert_eq!(count, 2);
422
423        let active_count = repo.count(Some(KeyStatus::Active)).await.unwrap();
424        assert_eq!(active_count, 2);
425    }
426
427    #[tokio::test]
428    async fn test_update_usage() {
429        let repo = InMemoryKeyRepository::new();
430        let key = create_test_key();
431        let id = key.id;
432
433        repo.create(key).await.unwrap();
434        let result = repo.update_usage(id, UsageRecord::priced(100, 0.01)).await;
435        assert!(result.is_ok(), "usage update should succeed: {result:?}");
436
437        let found = repo.find_by_id(id).await.unwrap().unwrap();
438        assert_eq!(found.usage_stats.total_requests, 1);
439        assert_eq!(found.usage_stats.total_tokens, 100);
440    }
441
442    #[tokio::test]
443    async fn test_duplicate_hash_rejected() {
444        let repo = InMemoryKeyRepository::new();
445        let mut key1 = create_test_key();
446        let hash = "same_hash".to_string();
447        key1.key_hash = hash.clone();
448
449        repo.create(key1).await.unwrap();
450
451        let mut key2 = create_test_key();
452        key2.key_hash = hash;
453
454        let result = repo.create(key2).await;
455        assert!(result.is_err());
456    }
457}
litellm_rs/core/keys/repository.rs

litellm_rs/core/keys/
repository.rs
