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#[derive(Debug, Serialize, Deserialize)]
17pub struct CachedToken {
18 pub access_token: String,
19 pub refresh_token: String,
20 pub expiry_timestamp: u64, }
22
23#[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 pub fn from_env() -> Result<Self, ConfigError> {
34 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 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 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 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#[derive(Clone)]
79pub struct TokenCache {
80 config: TokenCacheConfig,
81 cipher: Aes256Gcm,
82}
83
84impl 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 pub fn new(config: TokenCacheConfig) -> Result<Self, GmailApiError> {
97 if !config.enabled {
98 debug!("Creating TokenCache (disabled)");
99 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 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 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 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 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 let token = CachedToken {
174 access_token: access_token.to_string(),
175 refresh_token: refresh_token.to_string(),
176 expiry_timestamp,
177 };
178
179 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 let encrypted_data = self.encrypt_data(token_json.as_bytes())?;
193
194 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 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 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 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 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 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 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 match serde_json::from_slice::<CachedToken>(&decrypted_data) {
282 Ok(token) => {
283 debug!("Successfully loaded token from cache");
284 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)); }
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 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 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 token.expiry_timestamp > now + 300
322 }
323
324 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 fn encrypt_data(&self, data: &[u8]) -> GmailResult<Vec<u8>> {
353 let nonce_value = rand::random::<[u8; 12]>();
355 let nonce = Nonce::from_slice(&nonce_value);
356
357 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 let mut result = nonce_value.to_vec();
368 result.extend_from_slice(&ciphertext);
369 Ok(result)
370 }
371
372 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 let nonce = Nonce::from_slice(&data[0..12]);
382 let ciphertext = &data[12..];
383
384 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
395fn generate_encryption_key(secret: &str) -> Vec<u8> {
397 let mut key = Vec::with_capacity(32); let source = secret.as_bytes().to_vec();
400
401 if source.len() < 32 {
403 while key.len() < 32 {
405 key.extend_from_slice(&source);
406 }
407 key.truncate(32);
408 } else if source.len() > 32 {
409 key.extend_from_slice(&source[0..32]);
411 } else {
412 key = source;
414 }
415
416 key
417}
418
419fn fallback_encryption_key() -> Vec<u8> {
422 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 let combined = format!("gmail-mcp-rs-{}-{}", hostname, username);
435 generate_encryption_key(&combined)
436}
437
438fn default_cache_path() -> PathBuf {
440 let mut path = match dirs::cache_dir() {
441 Some(cache_dir) => cache_dir,
442 None => {
443 if let Some(temp_dir) = std::env::temp_dir().to_str() {
445 PathBuf::from(temp_dir)
446 } else {
447 PathBuf::from("/tmp") }
449 }
450 };
451
452 path.push("gmail-mcp-rs");
453 path.push("token-cache.dat");
454 path
455}