nordisk/monetize

crates/monetize-ui-server/src/state.rs

//! The BFF's small persistent state: opaque bytes under string keys — the
//! registered passkeys, the session MAC key, the revocation epoch. redb rather
//! than a file because [`StateStore::get_or_init`] must be atomic: two requests
//! racing to create the session key must not end up with two keys.

use std::fmt;
use std::path::Path;
use std::sync::Arc;

use redb::{Database, ReadableTable, TableDefinition};

const AUTH: TableDefinition<&str, &[u8]> = TableDefinition::new("auth");

#[derive(Clone)]
pub struct StateStore {
    db: Arc<Database>,
}

impl fmt::Debug for StateStore {
    /// Never renders values: they include the session MAC key.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("StateStore").finish_non_exhaustive()
    }
}

fn err(what: &str, e: impl fmt::Display) -> String {
    format!("{what}: {e}")
}

impl StateStore {
    pub fn open(path: impl AsRef<Path>) -> Result<StateStore, String> {
        let db = Database::create(path.as_ref()).map_err(|e| err("cannot open the state store", e))?;
        let store = StateStore { db: Arc::new(db) };
        // Materialise the table so a read on a never-written store does not fail.
        let tx = store.db.begin_write().map_err(|e| err("cannot begin a write", e))?;
        tx.open_table(AUTH).map_err(|e| err("cannot open the auth table", e))?;
        tx.commit().map_err(|e| err("cannot commit", e))?;
        Ok(store)
    }

    pub fn get(&self, key: &str) -> Result<Option<Vec<u8>>, String> {
        let tx = self.db.begin_read().map_err(|e| err("cannot begin a read", e))?;
        let table = tx.open_table(AUTH).map_err(|e| err("cannot open the auth table", e))?;
        let held = table.get(key).map_err(|e| err(&format!("cannot read {key:?}"), e))?;
        Ok(held.map(|v| v.value().to_vec()))
    }

    pub fn put(&self, key: &str, value: &[u8]) -> Result<(), String> {
        let tx = self.db.begin_write().map_err(|e| err("cannot begin a write", e))?;
        {
            let mut table = tx.open_table(AUTH).map_err(|e| err("cannot open the auth table", e))?;
            table.insert(key, value).map_err(|e| err(&format!("cannot write {key:?}"), e))?;
        }
        tx.commit().map_err(|e| err("cannot commit", e))
    }

    /// Read a value, creating it from `make` if absent — one transaction, so
    /// concurrent callers agree on the value.
    pub fn get_or_init(&self, key: &str, make: impl FnOnce() -> Vec<u8>) -> Result<Vec<u8>, String> {
        let tx = self.db.begin_write().map_err(|e| err("cannot begin a write", e))?;
        let value = {
            let mut table = tx.open_table(AUTH).map_err(|e| err("cannot open the auth table", e))?;
            let existing = table.get(key).map_err(|e| err(&format!("cannot read {key:?}"), e))?.map(|v| v.value().to_vec());
            match existing {
                Some(v) => v,
                None => {
                    let fresh = make();
                    table.insert(key, fresh.as_slice()).map_err(|e| err(&format!("cannot write {key:?}"), e))?;
                    fresh
                }
            }
        };
        tx.commit().map_err(|e| err("cannot commit", e))?;
        Ok(value)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn get_or_init_creates_once_then_keeps_it_across_reopen() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("state.redb");
        let store = StateStore::open(&path).unwrap();
        assert_eq!(store.get("absent").unwrap(), None);
        let first = store.get_or_init("session_key", || vec![1, 2, 3]).unwrap();
        let second = store.get_or_init("session_key", || vec![9]).unwrap();
        assert_eq!(first, second, "a second key would invalidate every cookie the first minted");
        drop(store);
        assert_eq!(StateStore::open(&path).unwrap().get("session_key").unwrap(), Some(vec![1, 2, 3]));
    }
}