Skip to main content

ironflow_store/memory/
api_key_store.rs

1//! [`ApiKeyStore`] trait implementation for [`InMemoryStore`].
2
3use chrono::Utc;
4use uuid::Uuid;
5
6use crate::api_key_store::ApiKeyStore;
7use crate::entities::{ApiKey, ApiKeyUpdate, NewApiKey};
8use crate::error::StoreError;
9use crate::store::StoreFuture;
10
11use super::InMemoryStore;
12
13impl ApiKeyStore for InMemoryStore {
14    fn create_api_key(&self, req: NewApiKey) -> StoreFuture<'_, ApiKey> {
15        Box::pin(async move {
16            let mut state = self.state.write().await;
17            let now = Utc::now();
18            let id = Uuid::now_v7();
19            let key = ApiKey {
20                id,
21                user_id: req.user_id,
22                name: req.name,
23                key_hash: req.key_hash,
24                key_prefix: req.key_prefix,
25                scopes: req.scopes,
26                is_active: true,
27                expires_at: req.expires_at,
28                last_used_at: None,
29                created_at: now,
30                updated_at: now,
31                rate_limit_override: req.rate_limit_override,
32            };
33            state.api_keys.insert(id, key.clone());
34            Ok(key)
35        })
36    }
37
38    fn find_api_key_by_prefix(&self, prefix: &str) -> StoreFuture<'_, Option<ApiKey>> {
39        let prefix = prefix.to_string();
40        Box::pin(async move {
41            let state = self.state.read().await;
42            Ok(state
43                .api_keys
44                .values()
45                .find(|k| k.key_prefix == prefix && k.is_active)
46                .cloned())
47        })
48    }
49
50    fn find_api_key_by_id(&self, id: Uuid) -> StoreFuture<'_, Option<ApiKey>> {
51        Box::pin(async move {
52            let state = self.state.read().await;
53            Ok(state.api_keys.get(&id).cloned())
54        })
55    }
56
57    fn list_api_keys_by_user(&self, user_id: Uuid) -> StoreFuture<'_, Vec<ApiKey>> {
58        Box::pin(async move {
59            let state = self.state.read().await;
60            let keys: Vec<ApiKey> = state
61                .api_keys
62                .values()
63                .filter(|k| k.user_id == user_id)
64                .cloned()
65                .collect();
66            Ok(keys)
67        })
68    }
69
70    fn update_api_key(&self, id: Uuid, update: ApiKeyUpdate) -> StoreFuture<'_, ()> {
71        Box::pin(async move {
72            let mut state = self.state.write().await;
73            let key = state
74                .api_keys
75                .get_mut(&id)
76                .ok_or(StoreError::Database(format!("API key {id} not found")))?;
77            if let Some(name) = update.name {
78                key.name = name;
79            }
80            if let Some(scopes) = update.scopes {
81                key.scopes = scopes;
82            }
83            if let Some(is_active) = update.is_active {
84                key.is_active = is_active;
85            }
86            if let Some(expires_at) = update.expires_at {
87                key.expires_at = expires_at;
88            }
89            if let Some(rate_limit_override) = update.rate_limit_override {
90                key.rate_limit_override = rate_limit_override;
91            }
92            key.updated_at = Utc::now();
93            Ok(())
94        })
95    }
96
97    fn touch_api_key(&self, id: Uuid) -> StoreFuture<'_, ()> {
98        Box::pin(async move {
99            let mut state = self.state.write().await;
100            if let Some(key) = state.api_keys.get_mut(&id) {
101                key.last_used_at = Some(Utc::now());
102            }
103            Ok(())
104        })
105    }
106
107    fn delete_api_key(&self, id: Uuid) -> StoreFuture<'_, ()> {
108        Box::pin(async move {
109            let mut state = self.state.write().await;
110            state.api_keys.remove(&id);
111            Ok(())
112        })
113    }
114}
115
116#[cfg(test)]
117mod tests {
118    use crate::entities::ApiKeyScope;
119
120    use super::*;
121
122    #[tokio::test]
123    async fn create_api_key_returns_active_key() {
124        let store = InMemoryStore::new();
125        let user_id = Uuid::now_v7();
126
127        let key = store
128            .create_api_key(NewApiKey {
129                user_id,
130                name: "production".to_string(),
131                key_hash: "bcrypt_hash".to_string(),
132                key_prefix: "sk_prod_abc123".to_string(),
133                scopes: vec![ApiKeyScope::WorkflowsRead],
134                expires_at: None,
135                rate_limit_override: None,
136            })
137            .await
138            .unwrap();
139
140        assert_eq!(key.user_id, user_id);
141        assert_eq!(key.name, "production");
142        assert!(key.is_active);
143        assert!(key.last_used_at.is_none());
144    }
145
146    #[tokio::test]
147    async fn find_api_key_by_prefix_existing() {
148        let store = InMemoryStore::new();
149        let user_id = Uuid::now_v7();
150
151        let created = store
152            .create_api_key(NewApiKey {
153                user_id,
154                name: "test".to_string(),
155                key_hash: "hash".to_string(),
156                key_prefix: "sk_test_xyz".to_string(),
157                scopes: vec![],
158                expires_at: None,
159                rate_limit_override: None,
160            })
161            .await
162            .unwrap();
163
164        let found = store
165            .find_api_key_by_prefix("sk_test_xyz")
166            .await
167            .unwrap()
168            .expect("key should exist");
169
170        assert_eq!(found.id, created.id);
171    }
172
173    #[tokio::test]
174    async fn find_api_key_by_prefix_missing_returns_none() {
175        let store = InMemoryStore::new();
176        let found = store
177            .find_api_key_by_prefix("sk_nonexistent")
178            .await
179            .unwrap();
180
181        assert!(found.is_none());
182    }
183
184    #[tokio::test]
185    async fn find_api_key_by_prefix_inactive_returns_none() {
186        let store = InMemoryStore::new();
187        let user_id = Uuid::now_v7();
188
189        let created = store
190            .create_api_key(NewApiKey {
191                user_id,
192                name: "test".to_string(),
193                key_hash: "hash".to_string(),
194                key_prefix: "sk_inactive".to_string(),
195                scopes: vec![],
196                expires_at: None,
197                rate_limit_override: None,
198            })
199            .await
200            .unwrap();
201
202        // Deactivate the key
203        store
204            .update_api_key(
205                created.id,
206                ApiKeyUpdate {
207                    name: None,
208                    scopes: None,
209                    is_active: Some(false),
210                    expires_at: None,
211                    rate_limit_override: None,
212                },
213            )
214            .await
215            .unwrap();
216
217        let found = store.find_api_key_by_prefix("sk_inactive").await.unwrap();
218
219        assert!(found.is_none());
220    }
221
222    #[tokio::test]
223    async fn find_api_key_by_id_existing() {
224        let store = InMemoryStore::new();
225        let user_id = Uuid::now_v7();
226
227        let created = store
228            .create_api_key(NewApiKey {
229                user_id,
230                name: "test".to_string(),
231                key_hash: "hash".to_string(),
232                key_prefix: "sk_test".to_string(),
233                scopes: vec![],
234                expires_at: None,
235                rate_limit_override: None,
236            })
237            .await
238            .unwrap();
239
240        let found = store
241            .find_api_key_by_id(created.id)
242            .await
243            .unwrap()
244            .expect("key should exist");
245
246        assert_eq!(found.id, created.id);
247    }
248
249    #[tokio::test]
250    async fn find_api_key_by_id_missing_returns_none() {
251        let store = InMemoryStore::new();
252        let found = store.find_api_key_by_id(Uuid::now_v7()).await.unwrap();
253
254        assert!(found.is_none());
255    }
256
257    #[tokio::test]
258    async fn list_api_keys_by_user_returns_only_user_keys() {
259        let store = InMemoryStore::new();
260        let user1 = Uuid::now_v7();
261        let user2 = Uuid::now_v7();
262
263        store
264            .create_api_key(NewApiKey {
265                user_id: user1,
266                name: "key1".to_string(),
267                key_hash: "hash1".to_string(),
268                key_prefix: "sk_1".to_string(),
269                scopes: vec![],
270                expires_at: None,
271                rate_limit_override: None,
272            })
273            .await
274            .unwrap();
275
276        store
277            .create_api_key(NewApiKey {
278                user_id: user1,
279                name: "key2".to_string(),
280                key_hash: "hash2".to_string(),
281                key_prefix: "sk_2".to_string(),
282                scopes: vec![],
283                expires_at: None,
284                rate_limit_override: None,
285            })
286            .await
287            .unwrap();
288
289        store
290            .create_api_key(NewApiKey {
291                user_id: user2,
292                name: "key3".to_string(),
293                key_hash: "hash3".to_string(),
294                key_prefix: "sk_3".to_string(),
295                scopes: vec![],
296                expires_at: None,
297                rate_limit_override: None,
298            })
299            .await
300            .unwrap();
301
302        let user1_keys = store.list_api_keys_by_user(user1).await.unwrap();
303        assert_eq!(user1_keys.len(), 2);
304        assert!(user1_keys.iter().all(|k| k.user_id == user1));
305    }
306
307    #[tokio::test]
308    async fn update_api_key_name() {
309        let store = InMemoryStore::new();
310        let user_id = Uuid::now_v7();
311
312        let created = store
313            .create_api_key(NewApiKey {
314                user_id,
315                name: "old-name".to_string(),
316                key_hash: "hash".to_string(),
317                key_prefix: "sk_test".to_string(),
318                scopes: vec![],
319                expires_at: None,
320                rate_limit_override: None,
321            })
322            .await
323            .unwrap();
324
325        store
326            .update_api_key(
327                created.id,
328                ApiKeyUpdate {
329                    name: Some("new-name".to_string()),
330                    scopes: None,
331                    is_active: None,
332                    expires_at: None,
333                    rate_limit_override: None,
334                },
335            )
336            .await
337            .unwrap();
338
339        let updated = store.find_api_key_by_id(created.id).await.unwrap().unwrap();
340
341        assert_eq!(updated.name, "new-name");
342    }
343
344    #[tokio::test]
345    async fn update_api_key_scopes() {
346        let store = InMemoryStore::new();
347        let user_id = Uuid::now_v7();
348
349        let created = store
350            .create_api_key(NewApiKey {
351                user_id,
352                name: "test".to_string(),
353                key_hash: "hash".to_string(),
354                key_prefix: "sk_test".to_string(),
355                scopes: vec![ApiKeyScope::WorkflowsRead],
356                expires_at: None,
357                rate_limit_override: None,
358            })
359            .await
360            .unwrap();
361
362        let new_scopes = vec![ApiKeyScope::WorkflowsRead, ApiKeyScope::RunsWrite];
363
364        store
365            .update_api_key(
366                created.id,
367                ApiKeyUpdate {
368                    name: None,
369                    scopes: Some(new_scopes.clone()),
370                    is_active: None,
371                    expires_at: None,
372                    rate_limit_override: None,
373                },
374            )
375            .await
376            .unwrap();
377
378        let updated = store.find_api_key_by_id(created.id).await.unwrap().unwrap();
379
380        assert_eq!(updated.scopes, new_scopes);
381    }
382
383    #[tokio::test]
384    async fn touch_api_key_updates_last_used() {
385        let store = InMemoryStore::new();
386        let user_id = Uuid::now_v7();
387
388        let created = store
389            .create_api_key(NewApiKey {
390                user_id,
391                name: "test".to_string(),
392                key_hash: "hash".to_string(),
393                key_prefix: "sk_test".to_string(),
394                scopes: vec![],
395                expires_at: None,
396                rate_limit_override: None,
397            })
398            .await
399            .unwrap();
400
401        assert!(created.last_used_at.is_none());
402
403        store.touch_api_key(created.id).await.unwrap();
404
405        let touched = store.find_api_key_by_id(created.id).await.unwrap().unwrap();
406
407        assert!(touched.last_used_at.is_some());
408    }
409
410    #[tokio::test]
411    async fn delete_api_key() {
412        let store = InMemoryStore::new();
413        let user_id = Uuid::now_v7();
414
415        let created = store
416            .create_api_key(NewApiKey {
417                user_id,
418                name: "test".to_string(),
419                key_hash: "hash".to_string(),
420                key_prefix: "sk_test".to_string(),
421                scopes: vec![],
422                expires_at: None,
423                rate_limit_override: None,
424            })
425            .await
426            .unwrap();
427
428        store.delete_api_key(created.id).await.unwrap();
429
430        let found = store.find_api_key_by_id(created.id).await.unwrap();
431
432        assert!(found.is_none());
433    }
434
435    #[tokio::test]
436    async fn delete_nonexistent_api_key_is_idempotent() {
437        let store = InMemoryStore::new();
438        let result = store.delete_api_key(Uuid::now_v7()).await;
439        assert!(result.is_ok());
440    }
441
442    #[tokio::test]
443    async fn create_api_key_with_rate_limit_override() {
444        let store = InMemoryStore::new();
445        let user_id = Uuid::now_v7();
446
447        let key = store
448            .create_api_key(NewApiKey {
449                user_id,
450                name: "high-volume".to_string(),
451                key_hash: "hash".to_string(),
452                key_prefix: "sk_hv".to_string(),
453                scopes: vec![],
454                expires_at: None,
455                rate_limit_override: Some(500),
456            })
457            .await
458            .unwrap();
459
460        assert_eq!(key.rate_limit_override, Some(500));
461
462        let found = store.find_api_key_by_id(key.id).await.unwrap().unwrap();
463        assert_eq!(found.rate_limit_override, Some(500));
464    }
465
466    #[tokio::test]
467    async fn update_api_key_rate_limit_override() {
468        let store = InMemoryStore::new();
469        let user_id = Uuid::now_v7();
470
471        let created = store
472            .create_api_key(NewApiKey {
473                user_id,
474                name: "test".to_string(),
475                key_hash: "hash".to_string(),
476                key_prefix: "sk_rl".to_string(),
477                scopes: vec![],
478                expires_at: None,
479                rate_limit_override: None,
480            })
481            .await
482            .unwrap();
483
484        assert!(created.rate_limit_override.is_none());
485
486        store
487            .update_api_key(
488                created.id,
489                ApiKeyUpdate {
490                    name: None,
491                    scopes: None,
492                    is_active: None,
493                    expires_at: None,
494                    rate_limit_override: Some(Some(200)),
495                },
496            )
497            .await
498            .unwrap();
499
500        let updated = store.find_api_key_by_id(created.id).await.unwrap().unwrap();
501        assert_eq!(updated.rate_limit_override, Some(200));
502
503        store
504            .update_api_key(
505                created.id,
506                ApiKeyUpdate {
507                    name: None,
508                    scopes: None,
509                    is_active: None,
510                    expires_at: None,
511                    rate_limit_override: Some(None),
512                },
513            )
514            .await
515            .unwrap();
516
517        let cleared = store.find_api_key_by_id(created.id).await.unwrap().unwrap();
518        assert!(cleared.rate_limit_override.is_none());
519    }
520}