nixfleet_control_plane/db/
tokens.rs

1//! Bootstrap-token nonces (soft state); loss bounded by one TTL.
2
3use anyhow::{Context, Result};
4use rusqlite::{Connection, params};
5use std::sync::Mutex;
6
7pub struct Tokens<'a> {
8    pub(super) conn: &'a Mutex<Connection>,
9}
10
11/// Distinguishes concurrent-replay race (409) from transient DB failure (500).
12#[derive(Debug, PartialEq, Eq)]
13pub enum RecordTokenOutcome {
14    Recorded,
15    AlreadyRecorded,
16}
17
18impl Tokens<'_> {
19    pub fn token_seen(&self, nonce: &str) -> Result<bool> {
20        super::read(self.conn, |c| {
21            c.query_row(
22                "SELECT 1 FROM token_replay WHERE nonce = ?1",
23                params![nonce],
24                |_| Ok(true),
25            )
26            .or_else(|err| match err {
27                rusqlite::Error::QueryReturnedNoRows => Ok(false),
28                e => Err(e),
29            })
30            .context("query token_replay")
31        })
32    }
33
34    /// Plain INSERT (not OR IGNORE): PK conflict surfaces as `AlreadyRecorded` for atomic check-and-set.
35    pub fn record_token_nonce(&self, nonce: &str, hostname: &str) -> Result<RecordTokenOutcome> {
36        super::read(self.conn, |c| {
37            match c.execute(
38                "INSERT INTO token_replay(nonce, hostname) VALUES (?1, ?2)",
39                params![nonce, hostname],
40            ) {
41                Ok(_) => Ok(RecordTokenOutcome::Recorded),
42                Err(rusqlite::Error::SqliteFailure(err, _))
43                    if err.code == rusqlite::ErrorCode::ConstraintViolation =>
44                {
45                    Ok(RecordTokenOutcome::AlreadyRecorded)
46                }
47                Err(e) => Err(anyhow::Error::from(e).context("insert token_replay")),
48            }
49        })
50    }
51
52    pub fn prune_token_replay(&self, max_age_hours: i64) -> Result<usize> {
53        super::read(self.conn, |c| {
54            c.execute(
55                "DELETE FROM token_replay
56                 WHERE first_seen < datetime('now', ?1)",
57                params![format!("-{max_age_hours} hours")],
58            )
59            .context("prune token_replay")
60        })
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use super::super::test_helpers::fresh_db;
67
68    #[test]
69    fn token_replay_round_trip() {
70        let db = fresh_db();
71        assert!(!db.tokens().token_seen("nonce-1").unwrap());
72        let outcome = db
73            .tokens()
74            .record_token_nonce("nonce-1", "test-host")
75            .unwrap();
76        assert_eq!(outcome, super::RecordTokenOutcome::Recorded);
77        assert!(db.tokens().token_seen("nonce-1").unwrap());
78    }
79
80    #[test]
81    fn record_token_nonce_returns_already_recorded_on_repeat() {
82        let db = fresh_db();
83        let first = db
84            .tokens()
85            .record_token_nonce("nonce-1", "test-host")
86            .unwrap();
87        assert_eq!(first, super::RecordTokenOutcome::Recorded);
88
89        let second = db
90            .tokens()
91            .record_token_nonce("nonce-1", "test-host")
92            .unwrap();
93        assert_eq!(second, super::RecordTokenOutcome::AlreadyRecorded);
94    }
95}
nixfleet_control_plane/db/tokens.rs

nixfleet_control_plane/db/
tokens.rs
