1use std::collections::HashMap;
2use std::sync::Arc;
3
4use aes_gcm::{
5 aead::{
6 Aead,
7 KeyInit,
8 },
9 Aes256Gcm,
10 Nonce,
11};
12use async_trait::async_trait;
13use sqlx::{
14 Pool,
15 Sqlite,
16};
17use tokio::sync::RwLock;
18
19use crate::domain::{
20 DomainError,
21 DomainResult,
22};
23use crate::infrastructure::TokenStore;
24
25const SQLITE_KDF_SALT: &[u8] = b"pipedash-sqlite-vault-v1";
26const BACKUP_KDF_SALT: &[u8] = b"pipedash-backup-salt-v1";
27const VAULT_VERIFIER: &str = "pipedash-vault-verifier-v1";
28
29pub struct SqliteTokenStore {
30 pool: Pool<Sqlite>,
31 encryption_key: [u8; 32],
32 cache: Arc<RwLock<HashMap<i64, String>>>,
33}
34
35impl SqliteTokenStore {
36 pub async fn new(pool: Pool<Sqlite>, vault_password: Option<String>) -> DomainResult<Self> {
37 let password = vault_password
38 .or_else(|| std::env::var("PIPEDASH_VAULT_PASSWORD").ok())
39 .ok_or_else(|| {
40 DomainError::InvalidConfig(
41 "PIPEDASH_VAULT_PASSWORD environment variable required for SQLite token encryption".into(),
42 )
43 })?;
44
45 let encryption_key = Self::derive_encryption_key(&password);
46
47 let store = Self {
48 pool,
49 encryption_key,
50 cache: Arc::new(RwLock::new(HashMap::new())),
51 };
52
53 let has_verifier = store.check_verifier().await?;
54 store.load_to_cache().await?;
55 if !has_verifier {
56 store.store_verifier().await?;
57 }
58
59 tracing::info!("SQLite token store initialized");
60 Ok(store)
61 }
62
63 pub async fn with_pool(
64 pool: Pool<Sqlite>, vault_password: Option<String>,
65 ) -> DomainResult<Self> {
66 Self::new(pool, vault_password).await
67 }
68
69 fn derive_encryption_key(password: &str) -> [u8; 32] {
70 use argon2::{
71 Argon2,
72 ParamsBuilder,
73 };
74
75 let mut output = [0u8; 32];
76
77 let params = ParamsBuilder::new()
78 .m_cost(65536)
79 .t_cost(3)
80 .p_cost(1)
81 .output_len(32)
82 .build()
83 .expect("Invalid Argon2 parameters");
84
85 let argon2 = Argon2::new(argon2::Algorithm::Argon2id, argon2::Version::V0x13, params);
86
87 argon2
88 .hash_password_into(password.as_bytes(), SQLITE_KDF_SALT, &mut output)
89 .expect("Failed to derive encryption key");
90
91 output
92 }
93
94 fn derive_backup_key(password: &str) -> [u8; 32] {
95 use argon2::{
96 Argon2,
97 ParamsBuilder,
98 };
99
100 let mut output = [0u8; 32];
101
102 let params = ParamsBuilder::new()
103 .m_cost(65536)
104 .t_cost(3)
105 .p_cost(1)
106 .output_len(32)
107 .build()
108 .expect("Invalid Argon2 parameters");
109
110 let argon2 = Argon2::new(argon2::Algorithm::Argon2id, argon2::Version::V0x13, params);
111
112 argon2
113 .hash_password_into(password.as_bytes(), BACKUP_KDF_SALT, &mut output)
114 .expect("Failed to derive backup key");
115
116 output
117 }
118
119 async fn encrypt_token(&self, plaintext: &str) -> DomainResult<(Vec<u8>, Vec<u8>)> {
120 let cipher = Aes256Gcm::new_from_slice(&self.encryption_key)
121 .map_err(|e| DomainError::InternalError(format!("Failed to create cipher: {}", e)))?;
122
123 let nonce_bytes: [u8; 12] = rand::random();
124 let nonce = Nonce::from(nonce_bytes);
125
126 let ciphertext = cipher
127 .encrypt(&nonce, plaintext.as_bytes())
128 .map_err(|e| DomainError::InternalError(format!("Encryption failed: {}", e)))?;
129
130 Ok((nonce_bytes.to_vec(), ciphertext))
131 }
132
133 async fn decrypt_token(&self, nonce: &[u8], ciphertext: &[u8]) -> DomainResult<String> {
134 if nonce.len() != 12 {
135 return Err(DomainError::InvalidConfig("Invalid nonce length".into()));
136 }
137
138 let cipher = Aes256Gcm::new_from_slice(&self.encryption_key)
139 .map_err(|e| DomainError::InternalError(format!("Failed to create cipher: {}", e)))?;
140
141 let nonce_array: [u8; 12] = nonce
142 .try_into()
143 .map_err(|_| DomainError::InvalidConfig("Invalid nonce".into()))?;
144 let nonce = Nonce::from(nonce_array);
145
146 let plaintext = cipher.decrypt(&nonce, ciphertext).map_err(|_| {
147 DomainError::AuthenticationFailed("Token decryption failed - wrong password?".into())
148 })?;
149
150 String::from_utf8(plaintext)
151 .map_err(|e| DomainError::DatabaseError(format!("Invalid UTF-8 in token: {}", e)))
152 }
153
154 async fn check_verifier(&self) -> DomainResult<bool> {
155 let row = sqlx::query_as::<_, (Vec<u8>, Vec<u8>)>(
156 "SELECT nonce, ciphertext FROM vault_verifier WHERE id = 1",
157 )
158 .fetch_optional(&self.pool)
159 .await
160 .map_err(|e| DomainError::DatabaseError(format!("Failed to read vault verifier: {}", e)))?;
161
162 let Some((nonce, ciphertext)) = row else {
163 return Ok(false);
164 };
165
166 self.decrypt_token(&nonce, &ciphertext)
167 .await
168 .map_err(|_| DomainError::AuthenticationFailed("Wrong vault password".into()))?;
169 Ok(true)
170 }
171
172 async fn store_verifier(&self) -> DomainResult<()> {
173 let (nonce, ciphertext) = self.encrypt_token(VAULT_VERIFIER).await?;
174
175 sqlx::query(
176 "INSERT INTO vault_verifier (id, nonce, ciphertext) VALUES (1, ?1, ?2)
177 ON CONFLICT (id) DO NOTHING",
178 )
179 .bind(&nonce)
180 .bind(&ciphertext)
181 .execute(&self.pool)
182 .await
183 .map_err(|e| {
184 DomainError::DatabaseError(format!("Failed to store vault verifier: {}", e))
185 })?;
186
187 self.check_verifier().await?;
188 Ok(())
189 }
190
191 async fn load_to_cache(&self) -> DomainResult<()> {
192 let rows = sqlx::query_as::<_, (i64, Vec<u8>, Vec<u8>)>(
193 "SELECT id, encrypted_token, token_nonce FROM providers
194 WHERE encrypted_token IS NOT NULL AND token_nonce IS NOT NULL",
195 )
196 .fetch_all(&self.pool)
197 .await
198 .map_err(|e| DomainError::DatabaseError(format!("Failed to load tokens: {}", e)))?;
199
200 let mut cache = self.cache.write().await;
201 cache.clear();
202
203 let mut decryption_errors = Vec::new();
204
205 for (provider_id, ciphertext, nonce) in rows {
206 match self.decrypt_token(&nonce, &ciphertext).await {
207 Ok(token) => {
208 cache.insert(provider_id, token);
209 }
210 Err(e) => {
211 tracing::warn!(
212 provider_id = provider_id,
213 error = %e,
214 "Failed to decrypt token for provider"
215 );
216 decryption_errors.push((provider_id, e));
217 }
218 }
219 }
220
221 if !decryption_errors.is_empty() {
222 let failed_count = decryption_errors.len();
223 return Err(DomainError::AuthenticationFailed(format!(
224 "Failed to decrypt {} token(s) - wrong vault password",
225 failed_count
226 )));
227 }
228
229 tracing::debug!(count = cache.len(), "Loaded tokens to cache");
230 Ok(())
231 }
232}
233
234#[async_trait]
235impl TokenStore for SqliteTokenStore {
236 async fn store_token(&self, provider_id: i64, token: &str) -> DomainResult<()> {
237 tracing::debug!(
238 "[SqliteTokenStore] Storing token for provider {} (token length: {})",
239 provider_id,
240 token.len()
241 );
242
243 let (nonce, ciphertext) = self.encrypt_token(token).await?;
244 tracing::debug!(
245 "[SqliteTokenStore] Token encrypted: nonce={} bytes, ciphertext={} bytes",
246 nonce.len(),
247 ciphertext.len()
248 );
249
250 let result = sqlx::query(
251 "UPDATE providers SET encrypted_token = ?1, token_nonce = ?2, updated_at = datetime('now')
252 WHERE id = ?3",
253 )
254 .bind(&ciphertext)
255 .bind(&nonce)
256 .bind(provider_id)
257 .execute(&self.pool)
258 .await
259 .map_err(|e| DomainError::DatabaseError(format!("Failed to store token: {}", e)))?;
260
261 tracing::info!(
262 "[SqliteTokenStore] Token stored for provider {} (rows affected: {})",
263 provider_id,
264 result.rows_affected()
265 );
266
267 if result.rows_affected() == 0 {
268 tracing::warn!(
269 "[SqliteTokenStore] No rows affected when storing token for provider {} - provider may not exist!",
270 provider_id
271 );
272 }
273
274 let mut cache = self.cache.write().await;
275 cache.insert(provider_id, token.to_string());
276
277 tracing::debug!(provider_id = provider_id, "Token stored successfully");
278 Ok(())
279 }
280
281 async fn get_token(&self, provider_id: i64) -> DomainResult<String> {
282 {
283 let cache = self.cache.read().await;
284 if let Some(token) = cache.get(&provider_id) {
285 return Ok(token.clone());
286 }
287 }
288
289 let row = sqlx::query_as::<_, (Vec<u8>, Vec<u8>)>(
290 "SELECT encrypted_token, token_nonce FROM providers
291 WHERE id = ?1 AND encrypted_token IS NOT NULL",
292 )
293 .bind(provider_id)
294 .fetch_optional(&self.pool)
295 .await
296 .map_err(|e| DomainError::DatabaseError(format!("Failed to get token: {}", e)))?
297 .ok_or_else(|| {
298 DomainError::NotFound(format!("Token not found for provider {}", provider_id))
299 })?;
300
301 let token = self.decrypt_token(&row.1, &row.0).await?;
302
303 {
304 let mut cache = self.cache.write().await;
305 cache.insert(provider_id, token.clone());
306 }
307
308 Ok(token)
309 }
310
311 async fn delete_token(&self, provider_id: i64) -> DomainResult<()> {
312 sqlx::query(
313 "UPDATE providers SET encrypted_token = NULL, token_nonce = NULL, updated_at = datetime('now')
314 WHERE id = ?1",
315 )
316 .bind(provider_id)
317 .execute(&self.pool)
318 .await
319 .map_err(|e| DomainError::DatabaseError(format!("Failed to delete token: {}", e)))?;
320
321 let mut cache = self.cache.write().await;
322 cache.remove(&provider_id);
323
324 tracing::debug!(provider_id = provider_id, "Token deleted successfully");
325 Ok(())
326 }
327
328 async fn get_all_tokens(&self) -> DomainResult<HashMap<i64, String>> {
329 let cache = self.cache.read().await;
330 Ok(cache.clone())
331 }
332
333 async fn export_encrypted(&self, password: &str) -> DomainResult<Vec<u8>> {
334 let tokens = self.get_all_tokens().await?;
335 let json = serde_json::to_vec(&tokens).map_err(|e| {
336 DomainError::DatabaseError(format!("Failed to serialize tokens: {}", e))
337 })?;
338
339 let backup_key = Self::derive_backup_key(password);
340 let cipher = Aes256Gcm::new_from_slice(&backup_key)
341 .map_err(|e| DomainError::InternalError(format!("Failed to create cipher: {}", e)))?;
342
343 let nonce_bytes: [u8; 12] = rand::random();
344 let nonce = Nonce::from(nonce_bytes);
345
346 let ciphertext = cipher
347 .encrypt(&nonce, json.as_ref())
348 .map_err(|e| DomainError::InternalError(format!("Encryption failed: {}", e)))?;
349
350 let mut result = nonce_bytes.to_vec();
351 result.extend(ciphertext);
352 Ok(result)
353 }
354
355 async fn import_encrypted(&self, data: &[u8], password: &str) -> DomainResult<()> {
356 tracing::info!(
357 "[SqliteTokenStore] Starting import_encrypted (data size: {} bytes, password length: {})",
358 data.len(),
359 password.len()
360 );
361
362 if data.len() < 12 {
363 return Err(DomainError::InvalidConfig("Invalid backup data".into()));
364 }
365
366 let nonce_bytes: [u8; 12] = data[..12]
367 .try_into()
368 .map_err(|_| DomainError::InvalidConfig("Invalid nonce".into()))?;
369 let ciphertext = &data[12..];
370
371 tracing::debug!(
372 "[SqliteTokenStore] Extracted nonce (12 bytes) and ciphertext ({} bytes)",
373 ciphertext.len()
374 );
375
376 let backup_key = Self::derive_backup_key(password);
377 let cipher = Aes256Gcm::new_from_slice(&backup_key)
378 .map_err(|e| DomainError::InternalError(format!("Failed to create cipher: {}", e)))?;
379
380 let nonce = Nonce::from(nonce_bytes);
381
382 tracing::debug!("[SqliteTokenStore] Attempting to decrypt token blob...");
383 let plaintext = cipher.decrypt(&nonce, ciphertext).map_err(|_| {
384 DomainError::AuthenticationFailed("Decryption failed (wrong password?)".into())
385 })?;
386 tracing::debug!(
387 "[SqliteTokenStore] Decryption successful ({} bytes)",
388 plaintext.len()
389 );
390
391 tracing::debug!("[SqliteTokenStore] Parsing JSON tokens...");
392 let tokens: HashMap<i64, String> = serde_json::from_slice(&plaintext)
393 .map_err(|e| DomainError::DatabaseError(format!("Failed to parse tokens: {}", e)))?;
394
395 tracing::info!(
396 "[SqliteTokenStore] Parsed {} tokens from encrypted blob",
397 tokens.len()
398 );
399
400 let mut stored_count = 0;
401 for (provider_id, token) in tokens {
402 tracing::debug!(
403 "[SqliteTokenStore] Importing token for provider {} (length: {})",
404 provider_id,
405 token.len()
406 );
407 self.store_token(provider_id, &token).await?;
408 stored_count += 1;
409 }
410
411 tracing::info!(
412 "[SqliteTokenStore] Import completed: {} tokens stored",
413 stored_count
414 );
415
416 Ok(())
417 }
418
419 async fn warmup(&self) -> DomainResult<()> {
420 tracing::info!("SQLite token store warmup complete (instant)");
421 Ok(())
422 }
423}
424
425#[cfg(test)]
426mod tests {
427 use super::*;
428
429 async fn create_test_pool() -> Pool<Sqlite> {
430 use sqlx::sqlite::SqlitePoolOptions;
431
432 let pool = SqlitePoolOptions::new()
433 .max_connections(1)
434 .connect(":memory:")
435 .await
436 .expect("Failed to create in-memory pool");
437
438 sqlx::query(
439 "CREATE TABLE providers (
440 id INTEGER PRIMARY KEY AUTOINCREMENT,
441 name TEXT NOT NULL UNIQUE,
442 provider_type TEXT NOT NULL,
443 token_encrypted TEXT NOT NULL DEFAULT '',
444 config_json TEXT NOT NULL DEFAULT '{}',
445 encrypted_token BLOB,
446 token_nonce BLOB,
447 created_at TEXT NOT NULL DEFAULT (datetime('now')),
448 updated_at TEXT NOT NULL DEFAULT (datetime('now'))
449 )",
450 )
451 .execute(&pool)
452 .await
453 .expect("Failed to create table");
454
455 sqlx::query(include_str!(
456 "../../../migrations/sqlite/20261020000002_vault_verifier.sql"
457 ))
458 .execute(&pool)
459 .await
460 .expect("Failed to create vault verifier table");
461
462 sqlx::query(
463 "INSERT INTO providers (name, provider_type) VALUES ('test-provider', 'github')",
464 )
465 .execute(&pool)
466 .await
467 .expect("Failed to insert provider");
468
469 pool
470 }
471
472 #[tokio::test]
473 async fn test_sqlite_basic() {
474 let pool = create_test_pool().await;
475 let store = SqliteTokenStore::new(pool, Some("test-password".to_string()))
476 .await
477 .unwrap();
478
479 store.store_token(1, "test-token").await.unwrap();
480 let token = store.get_token(1).await.unwrap();
481 assert_eq!(token, "test-token");
482 }
483
484 #[tokio::test]
485 async fn test_sqlite_cache() {
486 let pool = create_test_pool().await;
487 let store = SqliteTokenStore::new(pool, Some("test-password".to_string()))
488 .await
489 .unwrap();
490
491 store.store_token(1, "cached-token").await.unwrap();
492
493 let token = store.get_token(1).await.unwrap();
494 assert_eq!(token, "cached-token");
495 }
496
497 #[tokio::test]
498 async fn test_sqlite_delete() {
499 let pool = create_test_pool().await;
500 let store = SqliteTokenStore::new(pool, Some("test-password".to_string()))
501 .await
502 .unwrap();
503
504 store.store_token(1, "to-delete").await.unwrap();
505 store.delete_token(1).await.unwrap();
506
507 assert!(store.get_token(1).await.is_err());
508 }
509
510 #[tokio::test]
511 async fn test_sqlite_export_import() {
512 let pool = create_test_pool().await;
513 let store = SqliteTokenStore::new(pool.clone(), Some("test-password".to_string()))
514 .await
515 .unwrap();
516
517 store.store_token(1, "export-token").await.unwrap();
518
519 let encrypted = store.export_encrypted("backup-password").await.unwrap();
520
521 store.delete_token(1).await.unwrap();
522
523 store
524 .import_encrypted(&encrypted, "backup-password")
525 .await
526 .unwrap();
527
528 let token = store.get_token(1).await.unwrap();
529 assert_eq!(token, "export-token");
530 }
531
532 #[tokio::test]
533 async fn test_sqlite_wrong_backup_password() {
534 let pool = create_test_pool().await;
535 let store = SqliteTokenStore::new(pool, Some("test-password".to_string()))
536 .await
537 .unwrap();
538
539 store.store_token(1, "secret").await.unwrap();
540 let encrypted = store.export_encrypted("correct-password").await.unwrap();
541
542 let result = store.import_encrypted(&encrypted, "wrong-password").await;
543 assert!(result.is_err());
544 }
545
546 #[tokio::test]
547 async fn test_sqlite_empty_vault_rejects_another_password() {
548 let pool = create_test_pool().await;
549 SqliteTokenStore::new(pool.clone(), Some("operator-password".to_string()))
550 .await
551 .unwrap();
552
553 let result =
554 SqliteTokenStore::new(pool.clone(), Some("attacker-password".to_string())).await;
555
556 assert!(matches!(result, Err(DomainError::AuthenticationFailed(_))));
557 SqliteTokenStore::new(pool, Some("operator-password".to_string()))
558 .await
559 .unwrap();
560 }
561
562 #[tokio::test]
563 async fn test_sqlite_wrong_password_does_not_claim_vault_without_verifier() {
564 let pool = create_test_pool().await;
565 let store = SqliteTokenStore::new(pool.clone(), Some("operator-password".to_string()))
566 .await
567 .unwrap();
568 store.store_token(1, "secret").await.unwrap();
569 sqlx::query("DELETE FROM vault_verifier")
570 .execute(&pool)
571 .await
572 .unwrap();
573
574 let result =
575 SqliteTokenStore::new(pool.clone(), Some("attacker-password".to_string())).await;
576
577 assert!(result.is_err());
578 let store = SqliteTokenStore::new(pool, Some("operator-password".to_string()))
579 .await
580 .unwrap();
581 assert_eq!(store.get_token(1).await.unwrap(), "secret");
582 }
583}