liter_llm_proxy/auth/
key_store.rs

1use dashmap::DashMap;
2
3use crate::config::VirtualKeyConfig;
4
5/// Context injected into request extensions after successful auth.
6#[derive(Debug, Clone)]
7pub struct KeyContext {
8    pub key_id: String,
9    pub allowed_models: Option<Vec<String>>,
10    pub is_master: bool,
11}
12
13impl KeyContext {
14    /// Create a context representing the master key (unrestricted access).
15    pub fn master() -> Self {
16        Self {
17            key_id: "master".into(),
18            allowed_models: None,
19            is_master: true,
20        }
21    }
22
23    /// Create a context from a virtual key configuration.
24    pub fn from_config(config: &VirtualKeyConfig) -> Self {
25        let allowed_models = if config.models.is_empty() {
26            None
27        } else {
28            Some(config.models.clone())
29        };
30        Self {
31            key_id: config.key.clone(),
32            allowed_models,
33            is_master: false,
34        }
35    }
36
37    /// Returns true if this key is allowed to access the given model.
38    pub fn can_access_model(&self, model: &str) -> bool {
39        match &self.allowed_models {
40            None => true,
41            Some(models) => models.iter().any(|m| m == model),
42        }
43    }
44}
45
46/// In-memory virtual key store backed by `DashMap` for concurrent access.
47pub struct KeyStore {
48    keys: DashMap<String, VirtualKeyConfig>,
49    master_key: Option<String>,
50}
51
52impl KeyStore {
53    /// Build a key store from the proxy configuration values.
54    pub fn from_config(master_key: Option<String>, keys: &[VirtualKeyConfig]) -> Self {
55        let map = DashMap::new();
56        for k in keys {
57            map.insert(k.key.clone(), k.clone());
58        }
59        Self { keys: map, master_key }
60    }
61
62    /// Check whether `token` matches the configured master key.
63    pub fn is_master_key(&self, token: &str) -> bool {
64        self.master_key.as_deref() == Some(token)
65    }
66
67    /// Look up a virtual key configuration by its token string.
68    pub fn get(&self, token: &str) -> Option<VirtualKeyConfig> {
69        self.keys.get(token).map(|r| r.value().clone())
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76
77    fn sample_key_config(key: &str, models: Vec<String>) -> VirtualKeyConfig {
78        VirtualKeyConfig {
79            key: key.to_string(),
80            description: None,
81            models,
82            rpm: Some(60),
83            tpm: Some(100_000),
84            budget_limit: Some(50.0),
85        }
86    }
87
88    // ── KeyStore tests ──────────────────────────────────────────────────
89
90    #[test]
91    fn master_key_match_returns_true() {
92        let store = KeyStore::from_config(Some("sk-master".into()), &[]);
93        assert!(store.is_master_key("sk-master"));
94    }
95
96    #[test]
97    fn master_key_mismatch_returns_false() {
98        let store = KeyStore::from_config(Some("sk-master".into()), &[]);
99        assert!(!store.is_master_key("sk-wrong"));
100    }
101
102    #[test]
103    fn no_master_key_always_returns_false() {
104        let store = KeyStore::from_config(None, &[]);
105        assert!(!store.is_master_key("sk-anything"));
106    }
107
108    #[test]
109    fn get_existing_key_returns_config() {
110        let cfg = sample_key_config("vk-team-a", vec!["gpt-4o".into()]);
111        let store = KeyStore::from_config(None, std::slice::from_ref(&cfg));
112
113        let result = store.get("vk-team-a");
114        assert!(result.is_some());
115        let found = result.unwrap();
116        assert_eq!(found.key, "vk-team-a");
117        assert_eq!(found.models, vec!["gpt-4o"]);
118    }
119
120    #[test]
121    fn get_nonexistent_key_returns_none() {
122        let store = KeyStore::from_config(None, &[]);
123        assert!(store.get("vk-missing").is_none());
124    }
125
126    // ── KeyContext tests ────────────────────────────────────────────────
127
128    #[test]
129    fn master_context_has_no_restrictions() {
130        let ctx = KeyContext::master();
131        assert!(ctx.is_master);
132        assert!(ctx.allowed_models.is_none());
133        assert!(ctx.can_access_model("any-model"));
134    }
135
136    #[test]
137    fn context_with_allowed_models_permits_listed_model() {
138        let cfg = sample_key_config("vk-1", vec!["gpt-4o".into(), "claude-sonnet".into()]);
139        let ctx = KeyContext::from_config(&cfg);
140
141        assert!(!ctx.is_master);
142        assert!(ctx.can_access_model("gpt-4o"));
143        assert!(ctx.can_access_model("claude-sonnet"));
144    }
145
146    #[test]
147    fn context_with_allowed_models_denies_unlisted_model() {
148        let cfg = sample_key_config("vk-1", vec!["gpt-4o".into()]);
149        let ctx = KeyContext::from_config(&cfg);
150
151        assert!(!ctx.can_access_model("claude-sonnet"));
152    }
153
154    #[test]
155    fn context_with_empty_models_allows_all() {
156        let cfg = sample_key_config("vk-1", vec![]);
157        let ctx = KeyContext::from_config(&cfg);
158
159        assert!(ctx.allowed_models.is_none());
160        assert!(ctx.can_access_model("any-model"));
161    }
162}
liter_llm_proxy/auth/key_store.rs

liter_llm_proxy/auth/
key_store.rs
