ricecoder_providers/
api_key.rs

1//! API key management for secure credential handling and storage
2//!
3//! This module provides secure API key management with support for:
4//! - Storing keys in config files
5//! - Environment variable overrides
6//! - Key rotation support
7//! - Secure error messages that don't expose credentials
8
9use std::collections::HashMap;
10
11use crate::error::ProviderError;
12use crate::models::ApiKeyConfig;
13
14/// Manages API keys for providers with support for secure storage and retrieval
15pub struct ApiKeyManager {
16    /// Cached API keys (provider_id -> api_key)
17    keys: HashMap<String, String>,
18    /// API key configurations (provider_id -> config)
19    configs: HashMap<String, ApiKeyConfig>,
20}
21
22impl ApiKeyManager {
23    /// Create a new API key manager
24    pub fn new() -> Self {
25        Self {
26            keys: HashMap::new(),
27            configs: HashMap::new(),
28        }
29    }
30
31    /// Register an API key configuration for a provider
32    pub fn register_config(&mut self, provider_id: String, config: ApiKeyConfig) {
33        self.configs.insert(provider_id, config);
34    }
35
36    /// Store an API key for a provider
37    pub fn store_key(&mut self, provider_id: String, api_key: String) {
38        self.keys.insert(provider_id, api_key);
39    }
40
41    /// Get an API key for a provider
42    ///
43    /// Retrieves API key in the following order:
44    /// 1. From cache (if already loaded)
45    /// 2. From environment variable (if configured)
46    /// 3. From config file (if stored)
47    /// 4. Error if not found
48    pub fn get_key(&self, provider_id: &str) -> Result<String, ProviderError> {
49        // First check cache
50        if let Some(key) = self.keys.get(provider_id) {
51            return Ok(key.clone());
52        }
53
54        // Then check environment variable
55        if let Some(config) = self.configs.get(provider_id) {
56            if let Ok(key) = std::env::var(&config.env_var) {
57                return Ok(key);
58            }
59        }
60
61        // If not found, return error
62        Err(ProviderError::ConfigError(format!(
63            "API key not found for provider '{}'",
64            provider_id
65        )))
66    }
67
68    /// Check if an API key is available for a provider
69    pub fn has_key(&self, provider_id: &str) -> bool {
70        // Check cache
71        if self.keys.contains_key(provider_id) {
72            return true;
73        }
74
75        // Check environment variable
76        if let Some(config) = self.configs.get(provider_id) {
77            if std::env::var(&config.env_var).is_ok() {
78                return true;
79            }
80        }
81
82        false
83    }
84
85    /// Rotate an API key for a provider
86    ///
87    /// This updates the cached key and can optionally persist to config
88    pub fn rotate_key(
89        &mut self,
90        provider_id: String,
91        new_key: String,
92    ) -> Result<(), ProviderError> {
93        // Validate that the new key is not empty
94        if new_key.is_empty() {
95            return Err(ProviderError::ConfigError(
96                "API key cannot be empty".to_string(),
97            ));
98        }
99
100        // Update the cached key
101        self.keys.insert(provider_id, new_key);
102        Ok(())
103    }
104
105    /// Clear a cached API key (but not the environment variable)
106    pub fn clear_cached_key(&mut self, provider_id: &str) {
107        self.keys.remove(provider_id);
108    }
109
110    /// Clear all cached API keys
111    pub fn clear_all_cached_keys(&mut self) {
112        self.keys.clear();
113    }
114
115    /// Load API keys from environment variables
116    pub fn load_from_env(&mut self) -> Result<(), ProviderError> {
117        for (provider_id, config) in &self.configs {
118            if let Ok(key) = std::env::var(&config.env_var) {
119                self.keys.insert(provider_id.clone(), key);
120            }
121        }
122        Ok(())
123    }
124
125    /// Get the number of cached keys
126    pub fn cached_key_count(&self) -> usize {
127        self.keys.len()
128    }
129
130    /// Get all provider IDs with configured API key sources
131    pub fn configured_providers(&self) -> Vec<String> {
132        self.configs.keys().cloned().collect()
133    }
134}
135
136impl Default for ApiKeyManager {
137    fn default() -> Self {
138        Self::new()
139    }
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145
146    #[test]
147    fn test_new_api_key_manager() {
148        let manager = ApiKeyManager::new();
149        assert_eq!(manager.cached_key_count(), 0);
150    }
151
152    #[test]
153    fn test_store_and_get_key() {
154        let mut manager = ApiKeyManager::new();
155        manager.store_key("openai".to_string(), "sk-test-123".to_string());
156
157        let key = manager.get_key("openai");
158        assert!(key.is_ok());
159        assert_eq!(key.unwrap(), "sk-test-123");
160    }
161
162    #[test]
163    fn test_get_nonexistent_key() {
164        let manager = ApiKeyManager::new();
165        let key = manager.get_key("nonexistent");
166        assert!(key.is_err());
167    }
168
169    #[test]
170    fn test_has_key() {
171        let mut manager = ApiKeyManager::new();
172        assert!(!manager.has_key("openai"));
173
174        manager.store_key("openai".to_string(), "sk-test-123".to_string());
175        assert!(manager.has_key("openai"));
176    }
177
178    #[test]
179    fn test_rotate_key() {
180        let mut manager = ApiKeyManager::new();
181        manager.store_key("openai".to_string(), "sk-old-key".to_string());
182
183        let result = manager.rotate_key("openai".to_string(), "sk-new-key".to_string());
184        assert!(result.is_ok());
185
186        let key = manager.get_key("openai");
187        assert_eq!(key.unwrap(), "sk-new-key");
188    }
189
190    #[test]
191    fn test_rotate_key_empty() {
192        let mut manager = ApiKeyManager::new();
193        let result = manager.rotate_key("openai".to_string(), "".to_string());
194        assert!(result.is_err());
195    }
196
197    #[test]
198    fn test_clear_cached_key() {
199        let mut manager = ApiKeyManager::new();
200        manager.store_key("openai".to_string(), "sk-test-123".to_string());
201        assert!(manager.has_key("openai"));
202
203        manager.clear_cached_key("openai");
204        assert!(!manager.has_key("openai"));
205    }
206
207    #[test]
208    fn test_clear_all_cached_keys() {
209        let mut manager = ApiKeyManager::new();
210        manager.store_key("openai".to_string(), "sk-test-123".to_string());
211        manager.store_key("anthropic".to_string(), "sk-test-456".to_string());
212        assert_eq!(manager.cached_key_count(), 2);
213
214        manager.clear_all_cached_keys();
215        assert_eq!(manager.cached_key_count(), 0);
216    }
217
218    #[test]
219    fn test_register_config() {
220        let mut manager = ApiKeyManager::new();
221        let config = ApiKeyConfig {
222            env_var: "OPENAI_API_KEY".to_string(),
223            secure_storage: false,
224        };
225        manager.register_config("openai".to_string(), config);
226
227        assert!(manager.configs.contains_key("openai"));
228    }
229
230    #[test]
231    fn test_get_key_from_env() {
232        let mut manager = ApiKeyManager::new();
233        let config = ApiKeyConfig {
234            env_var: "TEST_API_KEY_ENV".to_string(),
235            secure_storage: false,
236        };
237        manager.register_config("test".to_string(), config);
238
239        // Set environment variable
240        std::env::set_var("TEST_API_KEY_ENV", "env-key-123");
241
242        let key = manager.get_key("test");
243        assert!(key.is_ok());
244        assert_eq!(key.unwrap(), "env-key-123");
245
246        // Cleanup
247        std::env::remove_var("TEST_API_KEY_ENV");
248    }
249
250    #[test]
251    fn test_load_from_env() {
252        let mut manager = ApiKeyManager::new();
253
254        let config1 = ApiKeyConfig {
255            env_var: "TEST_KEY_1".to_string(),
256            secure_storage: false,
257        };
258        let config2 = ApiKeyConfig {
259            env_var: "TEST_KEY_2".to_string(),
260            secure_storage: false,
261        };
262
263        manager.register_config("provider1".to_string(), config1);
264        manager.register_config("provider2".to_string(), config2);
265
266        std::env::set_var("TEST_KEY_1", "key-1");
267        std::env::set_var("TEST_KEY_2", "key-2");
268
269        manager.load_from_env().unwrap();
270
271        assert_eq!(manager.cached_key_count(), 2);
272        assert_eq!(manager.get_key("provider1").unwrap(), "key-1");
273        assert_eq!(manager.get_key("provider2").unwrap(), "key-2");
274
275        std::env::remove_var("TEST_KEY_1");
276        std::env::remove_var("TEST_KEY_2");
277    }
278
279    #[test]
280    fn test_configured_providers() {
281        let mut manager = ApiKeyManager::new();
282
283        let config1 = ApiKeyConfig {
284            env_var: "KEY_1".to_string(),
285            secure_storage: false,
286        };
287        let config2 = ApiKeyConfig {
288            env_var: "KEY_2".to_string(),
289            secure_storage: false,
290        };
291
292        manager.register_config("openai".to_string(), config1);
293        manager.register_config("anthropic".to_string(), config2);
294
295        let providers = manager.configured_providers();
296        assert_eq!(providers.len(), 2);
297        assert!(providers.contains(&"openai".to_string()));
298        assert!(providers.contains(&"anthropic".to_string()));
299    }
300
301    #[test]
302    fn test_cached_key_takes_precedence_over_env() {
303        let mut manager = ApiKeyManager::new();
304
305        let config = ApiKeyConfig {
306            env_var: "TEST_PRECEDENCE_KEY".to_string(),
307            secure_storage: false,
308        };
309        manager.register_config("test".to_string(), config);
310
311        // Set environment variable
312        std::env::set_var("TEST_PRECEDENCE_KEY", "env-key");
313
314        // Store a different key in cache
315        manager.store_key("test".to_string(), "cached-key".to_string());
316
317        // Cached key should take precedence
318        let key = manager.get_key("test");
319        assert_eq!(key.unwrap(), "cached-key");
320
321        std::env::remove_var("TEST_PRECEDENCE_KEY");
322    }
323}
ricecoder_providers/api_key.rs

ricecoder_providers/
api_key.rs
