pipedash_core/infrastructure/secrets/
postgres.rs1use 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 Postgres,
16};
17use tokio::sync::RwLock;
18
19use crate::domain::{
20 DomainError,
21 DomainResult,
22};
23use crate::infrastructure::TokenStore;
24
25const BACKUP_KDF_SALT: &[u8] = b"pipedash-backup-salt-v1";
26const VAULT_VERIFIER: &str = "pipedash-vault-verifier-v1";
27
28pub struct PostgresTokenStore {
29 pool: Pool<Postgres>,
30 encryption_key: [u8; 32],
31 cache: Arc<RwLock<HashMap<i64, String>>>,
32}
33
34impl PostgresTokenStore {
35 pub async fn new(
36 connection_string: &str, vault_password: Option<String>,
37 ) -> DomainResult<Self> {
38 use sqlx::postgres::PgPoolOptions;
39
40 let pool = PgPoolOptions::new()
41 .max_connections(20)
42 .min_connections(5)
43 .acquire_timeout(std::time::Duration::from_secs(30))
44 .after_connect(|conn, _meta| {
45 Box::pin(async move {
46 sqlx::query("SET search_path TO public")
47 .execute(conn)
48 .await?;
49 Ok(())
50 })
51 })
52 .connect(connection_string)
53 .await
54 .map_err(|e| {
55 DomainError::DatabaseError(format!("Failed to connect to PostgreSQL: {}", e))
56 })?;
57
58 let password = vault_password.or_else(|| std::env::var("PIPEDASH_VAULT_PASSWORD").ok());
59 let generated_key = password.is_none();
60 let password = password.unwrap_or_else(Self::generate_random_key);
61
62 let encryption_key = Self::derive_encryption_key(&password);
63
64 let store = Self {
65 pool,
66 encryption_key,
67 cache: Arc::new(RwLock::new(HashMap::new())),
68 };
69
70 let has_verifier = store.check_verifier().await?;
71 store.load_to_cache().await?;
72 if !has_verifier && !generated_key {
73 store.store_verifier().await?;
74 }
75
76 tracing::info!("PostgreSQL token store initialized");
77
78 Ok(store)
79 }
80
81 fn generate_random_key() -> String {
82 use rand::RngExt;
83 const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
84 let mut rng = rand::rng();
85 (0..64)
86 .map(|_| {
87 let idx = rng.random_range(0..CHARSET.len());
88 CHARSET[idx] as char
89 })
90 .collect()
91 }
92
93 fn derive_encryption_key(password: &str) -> [u8; 32] {
94 use argon2::{
95 Argon2,
96 ParamsBuilder,
97 };
98
99 let salt = b"pipedash-postgres-salt-v1";
100 let mut output = [0u8; 32];
101
102 let params = ParamsBuilder::new()
103 .m_cost(4096)
104 .t_cost(1)
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(), salt, &mut output)
114 .expect("Failed to derive encryption 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
147 .decrypt(&nonce, ciphertext)
148 .map_err(|_| DomainError::AuthenticationFailed("Decryption failed".into()))?;
149
150 String::from_utf8(plaintext)
151 .map_err(|e| DomainError::DatabaseError(format!("Invalid UTF-8: {}", 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 provider_id, nonce, ciphertext FROM encrypted_tokens",
194 )
195 .fetch_all(&self.pool)
196 .await
197 .map_err(|e| DomainError::DatabaseError(format!("Failed to load tokens: {}", e)))?;
198
199 let mut cache = self.cache.write().await;
200 cache.clear();
201
202 let mut decryption_errors = Vec::new();
203
204 for (provider_id, nonce, ciphertext) in rows {
205 match self.decrypt_token(&nonce, &ciphertext).await {
206 Ok(token) => {
207 cache.insert(provider_id, token);
208 }
209 Err(e) => {
210 tracing::warn!(
211 provider_id = provider_id,
212 error = %e,
213 "Failed to decrypt token for provider"
214 );
215 decryption_errors.push((provider_id, e));
216 }
217 }
218 }
219
220 if !decryption_errors.is_empty() {
221 let failed_count = decryption_errors.len();
222 return Err(DomainError::AuthenticationFailed(format!(
223 "Failed to decrypt {} token(s) - wrong vault password",
224 failed_count
225 )));
226 }
227
228 tracing::debug!(count = cache.len(), "Loaded tokens to cache");
229 Ok(())
230 }
231}
232
233#[async_trait]
234impl TokenStore for PostgresTokenStore {
235 async fn store_token(&self, provider_id: i64, token: &str) -> DomainResult<()> {
236 let (nonce, ciphertext) = self.encrypt_token(token).await?;
237
238 sqlx::query(
239 "INSERT INTO encrypted_tokens (provider_id, nonce, ciphertext, created_at, updated_at)
240 VALUES ($1, $2, $3, NOW(), NOW())
241 ON CONFLICT (provider_id)
242 DO UPDATE SET nonce = $2, ciphertext = $3, updated_at = NOW()",
243 )
244 .bind(provider_id)
245 .bind(&nonce)
246 .bind(&ciphertext)
247 .execute(&self.pool)
248 .await
249 .map_err(|e| DomainError::DatabaseError(format!("Failed to store token: {}", e)))?;
250
251 let mut cache = self.cache.write().await;
252 cache.insert(provider_id, token.to_string());
253
254 Ok(())
255 }
256
257 async fn get_token(&self, provider_id: i64) -> DomainResult<String> {
258 {
259 let cache = self.cache.read().await;
260 if let Some(token) = cache.get(&provider_id) {
261 return Ok(token.clone());
262 }
263 }
264
265 let row = sqlx::query_as::<_, (Vec<u8>, Vec<u8>)>(
266 "SELECT nonce, ciphertext FROM encrypted_tokens WHERE provider_id = $1",
267 )
268 .bind(provider_id)
269 .fetch_optional(&self.pool)
270 .await
271 .map_err(|e| DomainError::DatabaseError(format!("Failed to get token: {}", e)))?
272 .ok_or_else(|| {
273 DomainError::NotFound(format!("Token not found for provider {}", provider_id))
274 })?;
275
276 let token = self.decrypt_token(&row.0, &row.1).await?;
277
278 {
279 let mut cache = self.cache.write().await;
280 cache.insert(provider_id, token.clone());
281 }
282
283 Ok(token)
284 }
285
286 async fn delete_token(&self, provider_id: i64) -> DomainResult<()> {
287 sqlx::query("DELETE FROM encrypted_tokens WHERE provider_id = $1")
288 .bind(provider_id)
289 .execute(&self.pool)
290 .await
291 .map_err(|e| DomainError::DatabaseError(format!("Failed to delete token: {}", e)))?;
292
293 let mut cache = self.cache.write().await;
294 cache.remove(&provider_id);
295
296 Ok(())
297 }
298
299 async fn get_all_tokens(&self) -> DomainResult<HashMap<i64, String>> {
300 let cache = self.cache.read().await;
301 Ok(cache.clone())
302 }
303
304 async fn export_encrypted(&self, password: &str) -> DomainResult<Vec<u8>> {
305 let tokens = self.get_all_tokens().await?;
306 let json = serde_json::to_vec(&tokens).map_err(|e| {
307 DomainError::DatabaseError(format!("Failed to serialize tokens: {}", e))
308 })?;
309
310 let backup_key = Self::derive_backup_key(password);
311 let cipher = Aes256Gcm::new_from_slice(&backup_key)
312 .map_err(|e| DomainError::InternalError(format!("Failed to create cipher: {}", e)))?;
313
314 let nonce_bytes: [u8; 12] = rand::random();
315 let nonce = Nonce::from(nonce_bytes);
316
317 let ciphertext = cipher
318 .encrypt(&nonce, json.as_ref())
319 .map_err(|e| DomainError::InternalError(format!("Encryption failed: {}", e)))?;
320
321 let mut result = nonce_bytes.to_vec();
322 result.extend(ciphertext);
323
324 Ok(result)
325 }
326
327 async fn import_encrypted(&self, data: &[u8], password: &str) -> DomainResult<()> {
328 if data.len() < 12 {
329 return Err(DomainError::InvalidConfig("Invalid backup data".into()));
330 }
331
332 let nonce_bytes: [u8; 12] = data[..12]
333 .try_into()
334 .map_err(|_| DomainError::InvalidConfig("Invalid nonce".into()))?;
335 let ciphertext = &data[12..];
336
337 let backup_key = Self::derive_backup_key(password);
338 let cipher = Aes256Gcm::new_from_slice(&backup_key)
339 .map_err(|e| DomainError::InternalError(format!("Failed to create cipher: {}", e)))?;
340
341 let nonce = Nonce::from(nonce_bytes);
342
343 let plaintext = cipher.decrypt(&nonce, ciphertext).map_err(|_| {
344 DomainError::AuthenticationFailed("Decryption failed (wrong password?)".into())
345 })?;
346
347 let tokens: HashMap<i64, String> = serde_json::from_slice(&plaintext)
348 .map_err(|e| DomainError::DatabaseError(format!("Failed to parse tokens: {}", e)))?;
349
350 for (provider_id, token) in tokens {
351 self.store_token(provider_id, &token).await?;
352 }
353
354 Ok(())
355 }
356}
357
358impl PostgresTokenStore {
359 fn derive_backup_key(password: &str) -> [u8; 32] {
360 use argon2::{
361 Argon2,
362 ParamsBuilder,
363 };
364
365 let mut output = [0u8; 32];
366
367 let params = ParamsBuilder::new()
368 .m_cost(4096)
369 .t_cost(1)
370 .p_cost(1)
371 .output_len(32)
372 .build()
373 .expect("Invalid Argon2 parameters");
374
375 let argon2 = Argon2::new(argon2::Algorithm::Argon2id, argon2::Version::V0x13, params);
376
377 argon2
378 .hash_password_into(password.as_bytes(), BACKUP_KDF_SALT, &mut output)
379 .expect("Failed to derive backup key");
380
381 output
382 }
383}
384
385#[cfg(test)]
386mod tests {
387 use super::*;
388
389 const TEST_LOCK_ID: i64 = 0x7069_7065_6461_7368;
390
391 fn test_database_url() -> String {
392 std::env::var("TEST_DATABASE_URL")
393 .expect("TEST_DATABASE_URL must be set to run the PostgreSQL tests")
394 }
395
396 async fn create_test_pool() -> (Pool<Postgres>, [i64; 2]) {
397 crate::infrastructure::database::init_postgres_database(&test_database_url())
398 .await
399 .expect("Failed to migrate test database")
400 .close()
401 .await;
402
403 let pool = sqlx::postgres::PgPoolOptions::new()
404 .max_connections(1)
405 .connect(&test_database_url())
406 .await
407 .expect("Failed to connect to test database");
408
409 sqlx::query("SELECT pg_advisory_lock($1)")
410 .bind(TEST_LOCK_ID)
411 .execute(&pool)
412 .await
413 .expect("Failed to take test lock");
414
415 let mut providers = [0; 2];
416 for (provider, name) in providers
417 .iter_mut()
418 .zip(["secrets-test-1", "secrets-test-2"])
419 {
420 *provider = sqlx::query_scalar(
421 "INSERT INTO providers (name, provider_type, token_encrypted)
422 VALUES ($1, 'test', '')
423 ON CONFLICT (name) DO UPDATE SET name = EXCLUDED.name
424 RETURNING id",
425 )
426 .bind(name)
427 .fetch_one(&pool)
428 .await
429 .expect("Failed to create test provider");
430 }
431
432 sqlx::query("TRUNCATE TABLE encrypted_tokens, vault_verifier")
433 .execute(&pool)
434 .await
435 .expect("Failed to truncate test tables");
436
437 (pool, providers)
438 }
439
440 #[tokio::test]
441 #[ignore = "requires TEST_DATABASE_URL"]
442 async fn test_postgres_basic() {
443 let (pool, [first, _]) = create_test_pool().await;
444 let connection_string = test_database_url();
445
446 let store = PostgresTokenStore::new(&connection_string, Some("test-password".to_string()))
447 .await
448 .unwrap();
449
450 store.store_token(first, "test-token").await.unwrap();
451 let token = store.get_token(first).await.unwrap();
452 assert_eq!(token, "test-token");
453
454 pool.close().await;
455 }
456
457 #[tokio::test]
458 #[ignore = "requires TEST_DATABASE_URL"]
459 async fn test_postgres_persistence() {
460 let (pool, [first, _]) = create_test_pool().await;
461 let connection_string = test_database_url();
462
463 {
464 let store =
465 PostgresTokenStore::new(&connection_string, Some("test-password".to_string()))
466 .await
467 .unwrap();
468 store.store_token(first, "persistent-token").await.unwrap();
469 }
470
471 {
472 let store2 =
473 PostgresTokenStore::new(&connection_string, Some("test-password".to_string()))
474 .await
475 .unwrap();
476 let token = store2.get_token(first).await.unwrap();
477 assert_eq!(token, "persistent-token");
478 }
479
480 pool.close().await;
481 }
482
483 #[tokio::test]
484 #[ignore = "requires TEST_DATABASE_URL"]
485 async fn test_postgres_export_import() {
486 let (pool, [first, second]) = create_test_pool().await;
487 let connection_string = test_database_url();
488
489 let store = PostgresTokenStore::new(&connection_string, Some("test-password".to_string()))
490 .await
491 .unwrap();
492 store.store_token(first, "export-token").await.unwrap();
493 store.store_token(second, "export-token-2").await.unwrap();
494
495 let encrypted = store.export_encrypted("backup-password").await.unwrap();
496
497 sqlx::query("TRUNCATE TABLE encrypted_tokens, vault_verifier")
498 .execute(&pool)
499 .await
500 .unwrap();
501
502 let store2 = PostgresTokenStore::new(&connection_string, Some("test-password".to_string()))
503 .await
504 .unwrap();
505
506 store2
507 .import_encrypted(&encrypted, "backup-password")
508 .await
509 .unwrap();
510
511 assert_eq!(store2.get_token(first).await.unwrap(), "export-token");
512 assert_eq!(store2.get_token(second).await.unwrap(), "export-token-2");
513
514 pool.close().await;
515 }
516
517 #[tokio::test]
518 #[ignore = "requires TEST_DATABASE_URL"]
519 async fn test_postgres_empty_vault_rejects_another_password() {
520 let (pool, _) = create_test_pool().await;
521 let connection_string = test_database_url();
522
523 PostgresTokenStore::new(&connection_string, Some("operator-password".to_string()))
524 .await
525 .unwrap();
526
527 let result =
528 PostgresTokenStore::new(&connection_string, Some("attacker-password".to_string()))
529 .await;
530
531 assert!(matches!(result, Err(DomainError::AuthenticationFailed(_))));
532 PostgresTokenStore::new(&connection_string, Some("operator-password".to_string()))
533 .await
534 .unwrap();
535
536 pool.close().await;
537 }
538}