Skip to main content

mcp_gmailcal/
token_cache.rs

1use aes_gcm::{
2    aead::{Aead, KeyInit},
3    Aes256Gcm, Nonce,
4};
5use base64::{decode, encode};
6use log::{debug, error, warn};
7use serde::{Deserialize, Serialize};
8use std::fs::{self, File};
9use std::io::{Read, Write};
10use std::path::PathBuf;
11use std::time::SystemTime;
12
13use crate::errors::{ConfigError, GmailApiError, GmailResult};
14
15/// Cached token information
16#[derive(Debug, Serialize, Deserialize)]
17pub struct CachedToken {
18    pub access_token: String,
19    pub refresh_token: String,
20    pub expiry_timestamp: u64, // Unix timestamp when token expires
21}
22
23/// Configuration for the token cache
24#[derive(Debug, Clone)]
25pub struct TokenCacheConfig {
26    pub enabled: bool,
27    pub cache_file_path: PathBuf,
28    pub encryption_key: Vec<u8>,
29}
30
31impl TokenCacheConfig {
32    /// Create a new TokenCacheConfig from environment variables
33    pub fn from_env() -> Result<Self, ConfigError> {
34        // Check if token caching is enabled
35        let enabled = std::env::var("TOKEN_CACHE_ENABLED")
36            .map(|s| s.to_lowercase() == "true" || s == "1")
37            .unwrap_or(false);
38
39        if !enabled {
40            debug!("Token caching is disabled");
41            // Return with default values even though it's disabled
42            return Ok(Self {
43                enabled: false,
44                cache_file_path: default_cache_path(),
45                encryption_key: generate_encryption_key("default_unused_key"),
46            });
47        }
48
49        // Get cache file path
50        let cache_file_path = match std::env::var("TOKEN_CACHE_FILE") {
51            Ok(path) => PathBuf::from(path),
52            Err(_) => default_cache_path(),
53        };
54
55        // Get encryption key - NEVER log this
56        let encryption_key = match std::env::var("TOKEN_CACHE_ENCRYPTION_KEY") {
57            Ok(key) => generate_encryption_key(&key),
58            Err(_) => {
59                warn!("TOKEN_CACHE_ENCRYPTION_KEY not found, using less secure device-derived key");
60                fallback_encryption_key()
61            }
62        };
63
64        debug!(
65            "Token cache configured to use file: {}",
66            cache_file_path.display()
67        );
68
69        Ok(Self {
70            enabled,
71            cache_file_path,
72            encryption_key,
73        })
74    }
75}
76
77/// TokenCache provides secure persistence of OAuth tokens between application runs
78#[derive(Clone)]
79pub struct TokenCache {
80    config: TokenCacheConfig,
81    cipher: Aes256Gcm,
82}
83
84// Manual Debug implementation since Aes256Gcm doesn't implement Debug
85impl std::fmt::Debug for TokenCache {
86    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
87        f.debug_struct("TokenCache")
88            .field("config", &self.config)
89            .field("cipher", &"<AES-GCM cipher>")
90            .finish()
91    }
92}
93
94impl TokenCache {
95    /// Create a new TokenCache with the provided configuration
96    pub fn new(config: TokenCacheConfig) -> Result<Self, GmailApiError> {
97        if !config.enabled {
98            debug!("Creating TokenCache (disabled)");
99            // Even though it's disabled, create a valid instance
100            let cipher = match Aes256Gcm::new_from_slice(&config.encryption_key) {
101                Ok(cipher) => cipher,
102                Err(e) => {
103                    error!("Failed to initialize encryption: {}", e);
104                    return Err(GmailApiError::CacheError(format!(
105                        "Failed to initialize encryption: {}",
106                        e
107                    )));
108                }
109            };
110            return Ok(Self { config, cipher });
111        }
112
113        debug!("Creating TokenCache (enabled)");
114
115        // Ensure parent directory exists
116        if let Some(parent) = config.cache_file_path.parent() {
117            if !parent.exists() {
118                debug!("Creating parent directory for token cache");
119                match fs::create_dir_all(parent) {
120                    Ok(_) => {}
121                    Err(e) => {
122                        error!("Failed to create cache directory: {}", e);
123                        return Err(GmailApiError::CacheError(format!(
124                            "Failed to create cache directory: {}",
125                            e
126                        )));
127                    }
128                }
129            }
130        }
131
132        // Initialize AES-GCM cipher
133        let cipher = match Aes256Gcm::new_from_slice(&config.encryption_key) {
134            Ok(cipher) => cipher,
135            Err(e) => {
136                error!("Failed to initialize encryption: {}", e);
137                return Err(GmailApiError::CacheError(format!(
138                    "Failed to initialize encryption: {}",
139                    e
140                )));
141            }
142        };
143
144        Ok(Self { config, cipher })
145    }
146
147    /// Save token to the cache file
148    pub fn save_token(
149        &self,
150        access_token: &str,
151        refresh_token: &str,
152        expiry: SystemTime,
153    ) -> GmailResult<()> {
154        if !self.config.enabled {
155            debug!("Token caching disabled, not saving token");
156            return Ok(());
157        }
158
159        debug!("Saving token to cache");
160
161        // Convert expiry to Unix timestamp
162        let expiry_timestamp = match expiry.duration_since(SystemTime::UNIX_EPOCH) {
163            Ok(duration) => duration.as_secs(),
164            Err(e) => {
165                error!("Invalid expiry time: {}", e);
166                return Err(GmailApiError::CacheError(
167                    "Invalid expiry timestamp".to_string(),
168                ));
169            }
170        };
171
172        // Create token data
173        let token = CachedToken {
174            access_token: access_token.to_string(),
175            refresh_token: refresh_token.to_string(),
176            expiry_timestamp,
177        };
178
179        // Serialize token
180        let token_json = match serde_json::to_string(&token) {
181            Ok(json) => json,
182            Err(e) => {
183                error!("Failed to serialize token: {}", e);
184                return Err(GmailApiError::CacheError(format!(
185                    "Failed to serialize token: {}",
186                    e
187                )));
188            }
189        };
190
191        // Encrypt token
192        let encrypted_data = self.encrypt_data(token_json.as_bytes())?;
193
194        // Write to file
195        let encrypted_b64 = encode(&encrypted_data);
196        match File::create(&self.config.cache_file_path) {
197            Ok(mut file) => {
198                if let Err(e) = file.write_all(encrypted_b64.as_bytes()) {
199                    error!("Failed to write token cache: {}", e);
200                    return Err(GmailApiError::CacheError(format!(
201                        "Failed to write token cache: {}",
202                        e
203                    )));
204                }
205            }
206            Err(e) => {
207                error!("Failed to create token cache file: {}", e);
208                return Err(GmailApiError::CacheError(format!(
209                    "Failed to create token cache file: {}",
210                    e
211                )));
212            }
213        }
214
215        debug!("Token successfully cached to {}", self.config.cache_file_path.display());
216        Ok(())
217    }
218
219    /// Load token from the cache file
220    pub fn load_token(&self) -> GmailResult<Option<CachedToken>> {
221        if !self.config.enabled {
222            debug!("Token caching disabled, not loading token");
223            return Ok(None);
224        }
225
226        // Check if cache file exists
227        if !self.config.cache_file_path.exists() {
228            debug!("Token cache file not found");
229            return Ok(None);
230        }
231
232        debug!("Loading token from cache");
233
234        // Read encrypted data
235        let mut file = match File::open(&self.config.cache_file_path) {
236            Ok(file) => file,
237            Err(e) => {
238                error!("Failed to open token cache file: {}", e);
239                return Err(GmailApiError::CacheError(format!(
240                    "Failed to open token cache file: {}",
241                    e
242                )));
243            }
244        };
245
246        let mut encrypted_b64 = String::new();
247        if let Err(e) = file.read_to_string(&mut encrypted_b64) {
248            error!("Failed to read token cache: {}", e);
249            return Err(GmailApiError::CacheError(format!(
250                "Failed to read token cache: {}",
251                e
252            )));
253        }
254
255        // Decode base64
256        let encrypted_data = match decode(&encrypted_b64) {
257            Ok(data) => data,
258            Err(e) => {
259                error!("Failed to decode cached token data: {}", e);
260                return Err(GmailApiError::CacheError(format!(
261                    "Invalid token cache format: {}",
262                    e
263                )));
264            }
265        };
266
267        // Decrypt data
268        let decrypted_data = match self.decrypt_data(&encrypted_data) {
269            Ok(data) => data,
270            Err(e) => {
271                warn!("Failed to decrypt token cache: {}", e);
272                // If decryption fails, the cache is corrupt or key changed - delete it
273                if let Err(e) = fs::remove_file(&self.config.cache_file_path) {
274                    debug!("Failed to delete corrupt token cache: {}", e);
275                }
276                return Ok(None);
277            }
278        };
279
280        // Deserialize token
281        match serde_json::from_slice::<CachedToken>(&decrypted_data) {
282            Ok(token) => {
283                debug!("Successfully loaded token from cache");
284                // Check if token is expired
285                let now = match SystemTime::now().duration_since(SystemTime::UNIX_EPOCH) {
286                    Ok(duration) => duration.as_secs(),
287                    Err(_) => {
288                        error!("System time error when checking token expiry");
289                        return Ok(Some(token)); // Return token anyway, expiry will be checked elsewhere
290                    }
291                };
292
293                if token.expiry_timestamp <= now {
294                    debug!("Cached token has expired, will need refreshing");
295                }
296
297                Ok(Some(token))
298            }
299            Err(e) => {
300                error!("Failed to deserialize cached token: {}", e);
301                // Delete corrupt cache
302                if let Err(e) = fs::remove_file(&self.config.cache_file_path) {
303                    debug!("Failed to delete corrupt token cache: {}", e);
304                }
305                Ok(None)
306            }
307        }
308    }
309
310    /// Check if the token is still valid
311    pub fn is_token_valid(&self, token: &CachedToken) -> bool {
312        let now = match SystemTime::now().duration_since(SystemTime::UNIX_EPOCH) {
313            Ok(duration) => duration.as_secs(),
314            Err(_) => {
315                error!("System time error when checking token validity");
316                return false;
317            }
318        };
319
320        // Consider token valid if it has at least 5 minutes left
321        token.expiry_timestamp > now + 300
322    }
323
324    /// Delete the token cache file
325    pub fn clear_cache(&self) -> GmailResult<()> {
326        if !self.config.enabled {
327            return Ok(());
328        }
329
330        if self.config.cache_file_path.exists() {
331            debug!("Clearing token cache");
332            match fs::remove_file(&self.config.cache_file_path) {
333                Ok(_) => {
334                    debug!("Token cache cleared successfully");
335                    Ok(())
336                }
337                Err(e) => {
338                    error!("Failed to clear token cache: {}", e);
339                    Err(GmailApiError::CacheError(format!(
340                        "Failed to clear token cache: {}",
341                        e
342                    )))
343                }
344            }
345        } else {
346            debug!("No token cache to clear");
347            Ok(())
348        }
349    }
350
351    // Encrypt data using AES-GCM
352    fn encrypt_data(&self, data: &[u8]) -> GmailResult<Vec<u8>> {
353        // Use random 12-byte nonce (IV)
354        let nonce_value = rand::random::<[u8; 12]>();
355        let nonce = Nonce::from_slice(&nonce_value);
356
357        // Encrypt the data
358        let ciphertext = match self.cipher.encrypt(nonce, data) {
359            Ok(ciphertext) => ciphertext,
360            Err(e) => {
361                error!("Encryption failed: {}", e);
362                return Err(GmailApiError::CacheError(format!("Encryption failed: {}", e)));
363            }
364        };
365
366        // Combine nonce and ciphertext (nonce needs to be stored for decryption)
367        let mut result = nonce_value.to_vec();
368        result.extend_from_slice(&ciphertext);
369        Ok(result)
370    }
371
372    // Decrypt data using AES-GCM
373    fn decrypt_data(&self, data: &[u8]) -> GmailResult<Vec<u8>> {
374        if data.len() < 12 {
375            return Err(GmailApiError::CacheError(
376                "Invalid encrypted data: too short".to_string(),
377            ));
378        }
379
380        // Split data into nonce and ciphertext
381        let nonce = Nonce::from_slice(&data[0..12]);
382        let ciphertext = &data[12..];
383
384        // Decrypt the data
385        match self.cipher.decrypt(nonce, ciphertext) {
386            Ok(plaintext) => Ok(plaintext),
387            Err(e) => {
388                error!("Decryption failed: {}", e);
389                Err(GmailApiError::CacheError(format!("Decryption failed: {}", e)))
390            }
391        }
392    }
393}
394
395// Generate a consistent encryption key from a password/secret
396fn generate_encryption_key(secret: &str) -> Vec<u8> {
397    // Simple key derivation - in a production system, use a proper KDF like PBKDF2
398    let mut key = Vec::with_capacity(32); // 256 bits for AES-256
399    let source = secret.as_bytes().to_vec();
400    
401    // Pad or truncate the key to exactly 32 bytes
402    if source.len() < 32 {
403        // If key is too short, repeat it
404        while key.len() < 32 {
405            key.extend_from_slice(&source);
406        }
407        key.truncate(32);
408    } else if source.len() > 32 {
409        // If key is too long, truncate it
410        key.extend_from_slice(&source[0..32]);
411    } else {
412        // Key is exactly right size
413        key = source;
414    }
415    
416    key
417}
418
419// Fallback encryption key derived from machine-specific information
420// This is less secure than using a provided key but better than nothing
421fn fallback_encryption_key() -> Vec<u8> {
422    // Combine hostname and username to create a device-specific key
423    let hostname = match std::process::Command::new("hostname").output() {
424        Ok(output) => String::from_utf8_lossy(&output.stdout).to_string(),
425        Err(_) => "unknown-host".to_string(),
426    };
427    
428    let username = match std::env::var("USER") {
429        Ok(user) => user,
430        Err(_) => "unknown-user".to_string(),
431    };
432    
433    // Combine and hash to get a unique key
434    let combined = format!("gmail-mcp-rs-{}-{}", hostname, username);
435    generate_encryption_key(&combined)
436}
437
438// Get default cache file location (platform-specific)
439fn default_cache_path() -> PathBuf {
440    let mut path = match dirs::cache_dir() {
441        Some(cache_dir) => cache_dir,
442        None => {
443            // Fallback to temp directory if cache dir not available
444            if let Some(temp_dir) = std::env::temp_dir().to_str() {
445                PathBuf::from(temp_dir)
446            } else {
447                PathBuf::from("/tmp") // Unix-like systems fallback
448            }
449        }
450    };
451    
452    path.push("gmail-mcp-rs");
453    path.push("token-cache.dat");
454    path
455}