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> {
397 let pool = sqlx::postgres::PgPoolOptions::new()
398 .max_connections(1)
399 .connect(&test_database_url())
400 .await
401 .expect("Failed to connect to test database");
402
403 sqlx::query("SELECT pg_advisory_lock($1)")
404 .bind(TEST_LOCK_ID)
405 .execute(&pool)
406 .await
407 .expect("Failed to take test lock");
408
409 sqlx::query(
410 "CREATE TABLE IF NOT EXISTS encrypted_tokens (
411 provider_id BIGINT PRIMARY KEY,
412 nonce BYTEA NOT NULL,
413 ciphertext BYTEA NOT NULL,
414 created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
415 updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
416 )",
417 )
418 .execute(&pool)
419 .await
420 .expect("Failed to create test table");
421
422 sqlx::query(include_str!(
423 "../../../migrations/postgres/20261020000002_vault_verifier.sql"
424 ))
425 .execute(&pool)
426 .await
427 .expect("Failed to create vault verifier table");
428
429 sqlx::query("TRUNCATE TABLE encrypted_tokens, vault_verifier")
430 .execute(&pool)
431 .await
432 .expect("Failed to truncate test tables");
433
434 pool
435 }
436
437 #[tokio::test]
438 #[ignore = "requires TEST_DATABASE_URL"]
439 async fn test_postgres_basic() {
440 let pool = create_test_pool().await;
441 let connection_string = test_database_url();
442
443 let store = PostgresTokenStore::new(&connection_string, Some("test-password".to_string()))
444 .await
445 .unwrap();
446
447 store.store_token(1, "test-token").await.unwrap();
448 let token = store.get_token(1).await.unwrap();
449 assert_eq!(token, "test-token");
450
451 pool.close().await;
452 }
453
454 #[tokio::test]
455 #[ignore = "requires TEST_DATABASE_URL"]
456 async fn test_postgres_persistence() {
457 let pool = create_test_pool().await;
458 let connection_string = test_database_url();
459
460 {
461 let store =
462 PostgresTokenStore::new(&connection_string, Some("test-password".to_string()))
463 .await
464 .unwrap();
465 store.store_token(1, "persistent-token").await.unwrap();
466 }
467
468 {
469 let store2 =
470 PostgresTokenStore::new(&connection_string, Some("test-password".to_string()))
471 .await
472 .unwrap();
473 let token = store2.get_token(1).await.unwrap();
474 assert_eq!(token, "persistent-token");
475 }
476
477 pool.close().await;
478 }
479
480 #[tokio::test]
481 #[ignore = "requires TEST_DATABASE_URL"]
482 async fn test_postgres_export_import() {
483 let pool = create_test_pool().await;
484 let connection_string = test_database_url();
485
486 let store = PostgresTokenStore::new(&connection_string, Some("test-password".to_string()))
487 .await
488 .unwrap();
489 store.store_token(1, "export-token").await.unwrap();
490 store.store_token(2, "export-token-2").await.unwrap();
491
492 let encrypted = store.export_encrypted("backup-password").await.unwrap();
493
494 sqlx::query("TRUNCATE TABLE encrypted_tokens, vault_verifier")
495 .execute(&pool)
496 .await
497 .unwrap();
498
499 let store2 = PostgresTokenStore::new(&connection_string, Some("test-password".to_string()))
500 .await
501 .unwrap();
502
503 store2
504 .import_encrypted(&encrypted, "backup-password")
505 .await
506 .unwrap();
507
508 assert_eq!(store2.get_token(1).await.unwrap(), "export-token");
509 assert_eq!(store2.get_token(2).await.unwrap(), "export-token-2");
510
511 pool.close().await;
512 }
513
514 #[tokio::test]
515 #[ignore = "requires TEST_DATABASE_URL"]
516 async fn test_postgres_empty_vault_rejects_another_password() {
517 let pool = create_test_pool().await;
518 let connection_string = test_database_url();
519
520 PostgresTokenStore::new(&connection_string, Some("operator-password".to_string()))
521 .await
522 .unwrap();
523
524 let result =
525 PostgresTokenStore::new(&connection_string, Some("attacker-password".to_string()))
526 .await;
527
528 assert!(matches!(result, Err(DomainError::AuthenticationFailed(_))));
529 PostgresTokenStore::new(&connection_string, Some("operator-password".to_string()))
530 .await
531 .unwrap();
532
533 pool.close().await;
534 }
535}