secrets_storage_postgres/
lib.rs

1use std::collections::HashMap;
2
3use async_trait::async_trait;
4use secrets_core::storage::{StorageBackend, StorageEntry, StorageError, StorageResult};
5use sha2::{Digest, Sha256};
6use sqlx::pool::PoolConnection;
7use sqlx::{PgPool, Postgres};
8use tokio::sync::Mutex;
9
10pub struct PgStorage {
11    pool: PgPool,
12    /// Connections held open purely to keep `pg_advisory_lock` sessions alive.
13    ///
14    /// A Postgres advisory lock belongs to the *session* that took it, so the
15    /// connection cannot go back to the pool: another query reusing it could
16    /// release the lock, and a pooled `pg_advisory_unlock` might run on a
17    /// connection that never held it. Holding the connection also gives
18    /// failover for free — when the process dies the connection closes and
19    /// Postgres drops the lock, with no lease timeout to tune.
20    locks: Mutex<HashMap<String, PoolConnection<Postgres>>>,
21}
22
23impl PgStorage {
24    pub async fn connect(database_url: &str) -> Result<Self, sqlx::Error> {
25        let pool = PgPool::connect(database_url).await?;
26        sqlx::migrate!("./src/migrations").run(&pool).await?;
27        Ok(Self {
28            pool,
29            locks: Mutex::new(HashMap::new()),
30        })
31    }
32}
33
34/// Advisory locks are keyed by a single 64-bit integer, so a name has to be
35/// folded down to one. It must be stable across processes and releases —
36/// `DefaultHasher` is explicitly not, so hash with SHA-256 instead.
37fn advisory_lock_id(key: &str) -> i64 {
38    let digest = Sha256::digest(key.as_bytes());
39    i64::from_be_bytes(digest[..8].try_into().expect("sha256 digest is 32 bytes"))
40}
41
42#[async_trait]
43impl StorageBackend for PgStorage {
44    async fn get(&self, path: &str) -> StorageResult<Option<StorageEntry>> {
45        let row = sqlx::query_as::<_, (Vec<u8>, Option<chrono::DateTime<chrono::Utc>>)>(
46            "SELECT value, expires_at FROM kv_store WHERE path = $1",
47        )
48        .bind(path)
49        .fetch_optional(&self.pool)
50        .await
51        .map_err(|e| StorageError::Backend(e.to_string()))?;
52
53        Ok(row.map(|(value, expires_at)| StorageEntry { value, expires_at }))
54    }
55
56    async fn put(&self, path: &str, entry: StorageEntry) -> StorageResult<()> {
57        sqlx::query(
58            "INSERT INTO kv_store (path, value, expires_at, updated_at)
59             VALUES ($1, $2, $3, now())
60             ON CONFLICT (path) DO UPDATE SET value = $2, expires_at = $3, updated_at = now()",
61        )
62        .bind(path)
63        .bind(entry.value)
64        .bind(entry.expires_at)
65        .execute(&self.pool)
66        .await
67        .map_err(|e| StorageError::Backend(e.to_string()))?;
68        Ok(())
69    }
70
71    async fn delete(&self, path: &str) -> StorageResult<()> {
72        sqlx::query("DELETE FROM kv_store WHERE path = $1")
73            .bind(path)
74            .execute(&self.pool)
75            .await
76            .map_err(|e| StorageError::Backend(e.to_string()))?;
77        Ok(())
78    }
79
80    async fn ping(&self) -> StorageResult<()> {
81        sqlx::query("SELECT 1")
82            .execute(&self.pool)
83            .await
84            .map_err(|e| StorageError::Backend(e.to_string()))?;
85        Ok(())
86    }
87
88    /// Served by the partial index on `expires_at` from the initial
89    /// migration, so the reaper no longer reads every lease to find the few
90    /// that expired.
91    async fn list_expired(
92        &self,
93        prefix: &str,
94        now: chrono::DateTime<chrono::Utc>,
95    ) -> StorageResult<Vec<String>> {
96        let rows: Vec<(String,)> = sqlx::query_as(
97            "SELECT path FROM kv_store
98             WHERE path LIKE $1 || '%' AND expires_at IS NOT NULL AND expires_at <= $2",
99        )
100        .bind(prefix)
101        .bind(now)
102        .fetch_all(&self.pool)
103        .await
104        .map_err(|e| StorageError::Backend(e.to_string()))?;
105        Ok(rows.into_iter().map(|(p,)| p).collect())
106    }
107
108    /// One statement, so the comparison and the write cannot be interleaved.
109    async fn replace_if_unchanged(
110        &self,
111        path: &str,
112        expected: &[u8],
113        entry: StorageEntry,
114    ) -> StorageResult<bool> {
115        let result = sqlx::query(
116            "UPDATE kv_store SET value = $3, expires_at = $4, updated_at = now()
117             WHERE path = $1 AND value = $2",
118        )
119        .bind(path)
120        .bind(expected)
121        .bind(entry.value)
122        .bind(entry.expires_at)
123        .execute(&self.pool)
124        .await
125        .map_err(|e| StorageError::Backend(e.to_string()))?;
126        Ok(result.rows_affected() == 1)
127    }
128
129    async fn try_acquire_lock(&self, key: &str) -> StorageResult<bool> {
130        let mut locks = self.locks.lock().await;
131        // Already leading. Re-taking the same advisory lock on the same
132        // session would succeed and just bump Postgres' own counter, so skip
133        // the round trip entirely.
134        if locks.contains_key(key) {
135            return Ok(true);
136        }
137
138        let mut conn = self
139            .pool
140            .acquire()
141            .await
142            .map_err(|e| StorageError::Backend(e.to_string()))?;
143        let (acquired,): (bool,) = sqlx::query_as("SELECT pg_try_advisory_lock($1)")
144            .bind(advisory_lock_id(key))
145            .fetch_one(&mut *conn)
146            .await
147            .map_err(|e| StorageError::Backend(e.to_string()))?;
148
149        if acquired {
150            locks.insert(key.to_string(), conn);
151        }
152        Ok(acquired)
153    }
154
155    async fn list(&self, prefix: &str) -> StorageResult<Vec<String>> {
156        let rows: Vec<(String,)> =
157            sqlx::query_as("SELECT path FROM kv_store WHERE path LIKE $1 || '%'")
158                .bind(prefix)
159                .fetch_all(&self.pool)
160                .await
161                .map_err(|e| StorageError::Backend(e.to_string()))?;
162        Ok(rows.into_iter().map(|(p,)| p).collect())
163    }
164}
165
166#[cfg(test)]
167mod tests {
168    use super::advisory_lock_id;
169
170    /// The id has to be identical in every process that competes for the
171    /// lock, so this is a pinned value, not a round-trip check.
172    #[test]
173    fn advisory_lock_ids_are_stable_and_distinct() {
174        assert_eq!(advisory_lock_id("secrets/lease-reaper"), advisory_lock_id("secrets/lease-reaper"));
175        assert_ne!(advisory_lock_id("secrets/lease-reaper"), advisory_lock_id("secrets/other"));
176    }
177}
secrets_storage_postgres/lib.rs

secrets_storage_postgres/
lib.rs
