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";
26
27pub struct PostgresTokenStore {
28 pool: Pool<Postgres>,
29 encryption_key: [u8; 32],
30 cache: Arc<RwLock<HashMap<i64, String>>>,
31}
32
33impl PostgresTokenStore {
34 pub async fn new(
35 connection_string: &str, vault_password: Option<String>,
36 ) -> DomainResult<Self> {
37 use sqlx::postgres::PgPoolOptions;
38
39 let pool = PgPoolOptions::new()
40 .max_connections(20)
41 .min_connections(5)
42 .acquire_timeout(std::time::Duration::from_secs(30))
43 .after_connect(|conn, _meta| {
44 Box::pin(async move {
45 sqlx::query("SET search_path TO public")
46 .execute(conn)
47 .await?;
48 Ok(())
49 })
50 })
51 .connect(connection_string)
52 .await
53 .map_err(|e| {
54 DomainError::DatabaseError(format!("Failed to connect to PostgreSQL: {}", e))
55 })?;
56
57 let password = vault_password
58 .or_else(|| std::env::var("PIPEDASH_VAULT_PASSWORD").ok())
59 .unwrap_or_else(Self::generate_random_key);
60
61 let encryption_key = Self::derive_encryption_key(&password);
62
63 let store = Self {
64 pool,
65 encryption_key,
66 cache: Arc::new(RwLock::new(HashMap::new())),
67 };
68
69 store.load_to_cache().await?;
70
71 tracing::info!("PostgreSQL token store initialized");
72
73 Ok(store)
74 }
75
76 fn generate_random_key() -> String {
77 use rand::Rng;
78 const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
79 let mut rng = rand::rng();
80 (0..64)
81 .map(|_| {
82 let idx = rng.random_range(0..CHARSET.len());
83 CHARSET[idx] as char
84 })
85 .collect()
86 }
87
88 fn derive_encryption_key(password: &str) -> [u8; 32] {
89 use argon2::{
90 Argon2,
91 ParamsBuilder,
92 };
93
94 let salt = b"pipedash-postgres-salt-v1";
95 let mut output = [0u8; 32];
96
97 let params = ParamsBuilder::new()
98 .m_cost(4096)
99 .t_cost(1)
100 .p_cost(1)
101 .output_len(32)
102 .build()
103 .expect("Invalid Argon2 parameters");
104
105 let argon2 = Argon2::new(argon2::Algorithm::Argon2id, argon2::Version::V0x13, params);
106
107 argon2
108 .hash_password_into(password.as_bytes(), salt, &mut output)
109 .expect("Failed to derive encryption key");
110
111 output
112 }
113
114 async fn encrypt_token(&self, plaintext: &str) -> DomainResult<(Vec<u8>, Vec<u8>)> {
115 let cipher = Aes256Gcm::new_from_slice(&self.encryption_key)
116 .map_err(|e| DomainError::InternalError(format!("Failed to create cipher: {}", e)))?;
117
118 let nonce_bytes: [u8; 12] = rand::random();
119 let nonce = Nonce::from(nonce_bytes);
120
121 let ciphertext = cipher
122 .encrypt(&nonce, plaintext.as_bytes())
123 .map_err(|e| DomainError::InternalError(format!("Encryption failed: {}", e)))?;
124
125 Ok((nonce_bytes.to_vec(), ciphertext))
126 }
127
128 async fn decrypt_token(&self, nonce: &[u8], ciphertext: &[u8]) -> DomainResult<String> {
129 if nonce.len() != 12 {
130 return Err(DomainError::InvalidConfig("Invalid nonce length".into()));
131 }
132
133 let cipher = Aes256Gcm::new_from_slice(&self.encryption_key)
134 .map_err(|e| DomainError::InternalError(format!("Failed to create cipher: {}", e)))?;
135
136 let nonce_array: [u8; 12] = nonce
137 .try_into()
138 .map_err(|_| DomainError::InvalidConfig("Invalid nonce".into()))?;
139 let nonce = Nonce::from(nonce_array);
140
141 let plaintext = cipher
142 .decrypt(&nonce, ciphertext)
143 .map_err(|_| DomainError::AuthenticationFailed("Decryption failed".into()))?;
144
145 String::from_utf8(plaintext)
146 .map_err(|e| DomainError::DatabaseError(format!("Invalid UTF-8: {}", e)))
147 }
148
149 async fn load_to_cache(&self) -> DomainResult<()> {
150 let rows = sqlx::query_as::<_, (i64, Vec<u8>, Vec<u8>)>(
151 "SELECT provider_id, nonce, ciphertext FROM encrypted_tokens",
152 )
153 .fetch_all(&self.pool)
154 .await
155 .map_err(|e| DomainError::DatabaseError(format!("Failed to load tokens: {}", e)))?;
156
157 let mut cache = self.cache.write().await;
158 cache.clear();
159
160 let mut decryption_errors = Vec::new();
161
162 for (provider_id, nonce, ciphertext) in rows {
163 match self.decrypt_token(&nonce, &ciphertext).await {
164 Ok(token) => {
165 cache.insert(provider_id, token);
166 }
167 Err(e) => {
168 tracing::warn!(
169 provider_id = provider_id,
170 error = %e,
171 "Failed to decrypt token for provider"
172 );
173 decryption_errors.push((provider_id, e));
174 }
175 }
176 }
177
178 if !decryption_errors.is_empty() {
179 let failed_count = decryption_errors.len();
180 return Err(DomainError::AuthenticationFailed(format!(
181 "Failed to decrypt {} token(s) - wrong vault password",
182 failed_count
183 )));
184 }
185
186 tracing::debug!(count = cache.len(), "Loaded tokens to cache");
187 Ok(())
188 }
189}
190
191#[async_trait]
192impl TokenStore for PostgresTokenStore {
193 async fn store_token(&self, provider_id: i64, token: &str) -> DomainResult<()> {
194 let (nonce, ciphertext) = self.encrypt_token(token).await?;
195
196 sqlx::query(
197 "INSERT INTO encrypted_tokens (provider_id, nonce, ciphertext, created_at, updated_at)
198 VALUES ($1, $2, $3, NOW(), NOW())
199 ON CONFLICT (provider_id)
200 DO UPDATE SET nonce = $2, ciphertext = $3, updated_at = NOW()",
201 )
202 .bind(provider_id)
203 .bind(&nonce)
204 .bind(&ciphertext)
205 .execute(&self.pool)
206 .await
207 .map_err(|e| DomainError::DatabaseError(format!("Failed to store token: {}", e)))?;
208
209 let mut cache = self.cache.write().await;
210 cache.insert(provider_id, token.to_string());
211
212 Ok(())
213 }
214
215 async fn get_token(&self, provider_id: i64) -> DomainResult<String> {
216 {
217 let cache = self.cache.read().await;
218 if let Some(token) = cache.get(&provider_id) {
219 return Ok(token.clone());
220 }
221 }
222
223 let row = sqlx::query_as::<_, (Vec<u8>, Vec<u8>)>(
224 "SELECT nonce, ciphertext FROM encrypted_tokens WHERE provider_id = $1",
225 )
226 .bind(provider_id)
227 .fetch_optional(&self.pool)
228 .await
229 .map_err(|e| DomainError::DatabaseError(format!("Failed to get token: {}", e)))?
230 .ok_or_else(|| {
231 DomainError::DatabaseError(format!("Token not found for provider {}", provider_id))
232 })?;
233
234 let token = self.decrypt_token(&row.0, &row.1).await?;
235
236 {
237 let mut cache = self.cache.write().await;
238 cache.insert(provider_id, token.clone());
239 }
240
241 Ok(token)
242 }
243
244 async fn delete_token(&self, provider_id: i64) -> DomainResult<()> {
245 sqlx::query("DELETE FROM encrypted_tokens WHERE provider_id = $1")
246 .bind(provider_id)
247 .execute(&self.pool)
248 .await
249 .map_err(|e| DomainError::DatabaseError(format!("Failed to delete token: {}", e)))?;
250
251 let mut cache = self.cache.write().await;
252 cache.remove(&provider_id);
253
254 Ok(())
255 }
256
257 async fn get_all_tokens(&self) -> DomainResult<HashMap<i64, String>> {
258 let cache = self.cache.read().await;
259 Ok(cache.clone())
260 }
261
262 async fn export_encrypted(&self, password: &str) -> DomainResult<Vec<u8>> {
263 let tokens = self.get_all_tokens().await?;
264 let json = serde_json::to_vec(&tokens).map_err(|e| {
265 DomainError::DatabaseError(format!("Failed to serialize tokens: {}", e))
266 })?;
267
268 let backup_key = Self::derive_backup_key(password);
269 let cipher = Aes256Gcm::new_from_slice(&backup_key)
270 .map_err(|e| DomainError::InternalError(format!("Failed to create cipher: {}", e)))?;
271
272 let nonce_bytes: [u8; 12] = rand::random();
273 let nonce = Nonce::from(nonce_bytes);
274
275 let ciphertext = cipher
276 .encrypt(&nonce, json.as_ref())
277 .map_err(|e| DomainError::InternalError(format!("Encryption failed: {}", e)))?;
278
279 let mut result = nonce_bytes.to_vec();
280 result.extend(ciphertext);
281
282 Ok(result)
283 }
284
285 async fn import_encrypted(&self, data: &[u8], password: &str) -> DomainResult<()> {
286 if data.len() < 12 {
287 return Err(DomainError::InvalidConfig("Invalid backup data".into()));
288 }
289
290 let nonce_bytes: [u8; 12] = data[..12]
291 .try_into()
292 .map_err(|_| DomainError::InvalidConfig("Invalid nonce".into()))?;
293 let ciphertext = &data[12..];
294
295 let backup_key = Self::derive_backup_key(password);
296 let cipher = Aes256Gcm::new_from_slice(&backup_key)
297 .map_err(|e| DomainError::InternalError(format!("Failed to create cipher: {}", e)))?;
298
299 let nonce = Nonce::from(nonce_bytes);
300
301 let plaintext = cipher.decrypt(&nonce, ciphertext).map_err(|_| {
302 DomainError::AuthenticationFailed("Decryption failed (wrong password?)".into())
303 })?;
304
305 let tokens: HashMap<i64, String> = serde_json::from_slice(&plaintext)
306 .map_err(|e| DomainError::DatabaseError(format!("Failed to parse tokens: {}", e)))?;
307
308 for (provider_id, token) in tokens {
309 self.store_token(provider_id, &token).await?;
310 }
311
312 Ok(())
313 }
314}
315
316impl PostgresTokenStore {
317 fn derive_backup_key(password: &str) -> [u8; 32] {
318 use argon2::{
319 Argon2,
320 ParamsBuilder,
321 };
322
323 let mut output = [0u8; 32];
324
325 let params = ParamsBuilder::new()
326 .m_cost(4096)
327 .t_cost(1)
328 .p_cost(1)
329 .output_len(32)
330 .build()
331 .expect("Invalid Argon2 parameters");
332
333 let argon2 = Argon2::new(argon2::Algorithm::Argon2id, argon2::Version::V0x13, params);
334
335 argon2
336 .hash_password_into(password.as_bytes(), BACKUP_KDF_SALT, &mut output)
337 .expect("Failed to derive backup key");
338
339 output
340 }
341}
342
343#[cfg(test)]
344mod tests {
345 use super::*;
346
347 async fn create_test_pool() -> Pool<Postgres> {
348 let database_url = std::env::var("TEST_DATABASE_URL").unwrap_or_else(|_| {
349 "postgresql://postgres:postgres@localhost:5432/pipedash_test".to_string()
350 });
351
352 let pool = Pool::<Postgres>::connect(&database_url)
353 .await
354 .expect("Failed to connect to test database");
355
356 sqlx::query(
357 "CREATE TABLE IF NOT EXISTS encrypted_tokens (
358 provider_id BIGINT PRIMARY KEY,
359 nonce BYTEA NOT NULL,
360 ciphertext BYTEA NOT NULL,
361 created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
362 updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
363 )",
364 )
365 .execute(&pool)
366 .await
367 .expect("Failed to create test table");
368
369 sqlx::query("TRUNCATE TABLE encrypted_tokens")
370 .execute(&pool)
371 .await
372 .expect("Failed to truncate test table");
373
374 pool
375 }
376
377 #[tokio::test]
378 #[ignore]
379 async fn test_postgres_basic() {
380 let pool = create_test_pool().await;
381 let connection_string = std::env::var("TEST_DATABASE_URL").unwrap_or_else(|_| {
382 "postgresql://postgres:postgres@localhost:5432/pipedash_test".to_string()
383 });
384
385 let store = PostgresTokenStore::new(&connection_string, Some("test-password".to_string()))
386 .await
387 .unwrap();
388
389 store.store_token(1, "test-token").await.unwrap();
390 let token = store.get_token(1).await.unwrap();
391 assert_eq!(token, "test-token");
392
393 pool.close().await;
394 }
395
396 #[tokio::test]
397 #[ignore]
398 async fn test_postgres_persistence() {
399 let pool = create_test_pool().await;
400 let connection_string = std::env::var("TEST_DATABASE_URL").unwrap_or_else(|_| {
401 "postgresql://postgres:postgres@localhost:5432/pipedash_test".to_string()
402 });
403
404 {
405 let store =
406 PostgresTokenStore::new(&connection_string, Some("test-password".to_string()))
407 .await
408 .unwrap();
409 store.store_token(1, "persistent-token").await.unwrap();
410 }
411
412 {
413 let store2 =
414 PostgresTokenStore::new(&connection_string, Some("test-password".to_string()))
415 .await
416 .unwrap();
417 let token = store2.get_token(1).await.unwrap();
418 assert_eq!(token, "persistent-token");
419 }
420
421 pool.close().await;
422 }
423
424 #[tokio::test]
425 #[ignore]
426 async fn test_postgres_export_import() {
427 let pool = create_test_pool().await;
428 let connection_string = std::env::var("TEST_DATABASE_URL").unwrap_or_else(|_| {
429 "postgresql://postgres:postgres@localhost:5432/pipedash_test".to_string()
430 });
431
432 let store = PostgresTokenStore::new(&connection_string, Some("test-password".to_string()))
433 .await
434 .unwrap();
435 store.store_token(1, "export-token").await.unwrap();
436 store.store_token(2, "export-token-2").await.unwrap();
437
438 let encrypted = store.export_encrypted("backup-password").await.unwrap();
439
440 sqlx::query("TRUNCATE TABLE encrypted_tokens")
441 .execute(&pool)
442 .await
443 .unwrap();
444
445 let store2 = PostgresTokenStore::new(&connection_string, Some("test-password".to_string()))
446 .await
447 .unwrap();
448
449 store2
450 .import_encrypted(&encrypted, "backup-password")
451 .await
452 .unwrap();
453
454 assert_eq!(store2.get_token(1).await.unwrap(), "export-token");
455 assert_eq!(store2.get_token(2).await.unwrap(), "export-token-2");
456
457 pool.close().await;
458 }
459}