kestrove-core/crates/kv-wasm/src/store.rs

//! The web vault's [`Store`]: an in-memory replica that records which keys changed, so the wasm
//! layer can write just those entries to IndexedDB after each call (and read them all back on
//! start-up). [`Store`] is synchronous and IndexedDB is not, hence the split.
//!
//! Like the SQLite store it holds only ciphertext envelopes, ids and sync bookkeeping. Entries are
//! JSON bytes keyed by:
//!
//! | key                         | value                                       |
//! |-----------------------------|---------------------------------------------|
//! | `meta/<key>`                | raw meta bytes (account record, server, session) |
//! | `vault/<vault_id>`          | [`VaultEntry`]                              |
//! | `item/<vault_id>/<item_id>` | [`ItemEntry`] (working copy + base)         |
//! | `history/<vault_id>/<item_id>` | newest-first list of [`HistoryEntryJson`] |

use std::{
    collections::{BTreeMap, BTreeSet, HashMap},
    sync::Mutex,
};

use kv_proto::{Envelope, ItemRecord};
use kv_store::{HistoryEntry, LocalItem, Result, Store, StoreError, SyncState, VaultRow, HISTORY_LIMIT};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

/// One write for the persistence layer.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Change {
    /// Delete every entry first (sign-out / wipe). Always the first change when present.
    Clear,
    Put(String, Vec<u8>),
    Delete(String),
}

#[derive(Serialize, Deserialize)]
struct VaultEntry {
    /// Insertion order, so `vaults()` lists in the same order after a reload (like SQLite's rowid).
    order: u64,
    vault_id: Uuid,
    sealed_key: Envelope,
    meta: Option<Envelope>,
    cursor: u64,
    manifest_version: u64,
    manifest_hash: Option<[u8; 32]>,
    integrity_alarm: bool,
}

#[derive(Serialize, Deserialize)]
struct ItemEntry {
    record: ItemRecord,
    pending: bool,
    base: Option<ItemRecord>,
}

#[derive(Serialize, Deserialize)]
struct HistoryEntryJson {
    revision: u64,
    wrapped_item_key: Envelope,
    payload: Envelope,
    saved_at: i64,
}

#[derive(Default)]
struct Mem {
    meta: HashMap<String, Vec<u8>>,
    /// (order, row)
    vaults: Vec<(u64, VaultRow)>,
    items: HashMap<Uuid, BTreeMap<Uuid, LocalItem>>,
    history: HashMap<(Uuid, Uuid), Vec<HistoryEntry>>,
    next_order: u64,
    dirty: BTreeSet<String>,
    cleared: bool,
}

#[derive(Default)]
pub struct WebStore {
    inner: Mutex<Mem>,
}

fn meta_key(k: &str) -> String {
    format!("meta/{k}")
}
fn vault_key(v: Uuid) -> String {
    format!("vault/{v}")
}
fn item_key(v: Uuid, i: Uuid) -> String {
    format!("item/{v}/{i}")
}
fn history_key(v: Uuid, i: Uuid) -> String {
    format!("history/{v}/{i}")
}

fn corrupt(key: &str, e: impl std::fmt::Display) -> StoreError {
    StoreError::Io(format!("corrupt entry {key}: {e}"))
}

fn json<T: Serialize>(v: &T) -> Vec<u8> {
    // Only plain data structures are serialized here; this cannot fail.
    serde_json::to_vec(v).unwrap_or_default()
}

fn two_ids(key: &str, rest: &str) -> Result<(Uuid, Uuid)> {
    let (a, b) = rest.split_once('/').ok_or_else(|| corrupt(key, "key"))?;
    Ok((Uuid::parse_str(a).map_err(|e| corrupt(key, e))?, Uuid::parse_str(b).map_err(|e| corrupt(key, e))?))
}

impl WebStore {
    pub fn new() -> Self {
        Self::default()
    }

    fn mem(&self) -> std::sync::MutexGuard<'_, Mem> {
        self.inner.lock().unwrap_or_else(|p| p.into_inner())
    }

    /// Rebuilds the replica from persisted entries (start-up). Nothing is marked dirty.
    pub fn load(entries: impl IntoIterator<Item = (String, Vec<u8>)>) -> Result<Self> {
        let mut m = Mem::default();
        for (key, value) in entries {
            if let Some(k) = key.strip_prefix("meta/") {
                m.meta.insert(k.to_string(), value);
            } else if key.starts_with("vault/") {
                let e: VaultEntry = serde_json::from_slice(&value).map_err(|e| corrupt(&key, e))?;
                m.next_order = m.next_order.max(e.order + 1);
                m.vaults.push((
                    e.order,
                    VaultRow {
                        vault_id: e.vault_id,
                        sealed_key: e.sealed_key,
                        meta: e.meta,
                        cursor: e.cursor,
                        manifest_version: e.manifest_version,
                        manifest_hash: e.manifest_hash,
                        integrity_alarm: e.integrity_alarm,
                    },
                ));
            } else if let Some(rest) = key.strip_prefix("item/") {
                let (v, i) = two_ids(&key, rest)?;
                let e: ItemEntry = serde_json::from_slice(&value).map_err(|e| corrupt(&key, e))?;
                if e.record.item_id != i {
                    return Err(corrupt(&key, "item id"));
                }
                let state = if e.pending { SyncState::Pending } else { SyncState::Clean };
                m.items.entry(v).or_default().insert(i, LocalItem { record: e.record, state, base: e.base });
            } else if let Some(rest) = key.strip_prefix("history/") {
                let ids = two_ids(&key, rest)?;
                let h: Vec<HistoryEntryJson> = serde_json::from_slice(&value).map_err(|e| corrupt(&key, e))?;
                let h = h
                    .into_iter()
                    .map(|h| HistoryEntry {
                        revision: h.revision,
                        wrapped_item_key: h.wrapped_item_key,
                        payload: h.payload,
                        saved_at: h.saved_at,
                    })
                    .collect();
                m.history.insert(ids, h);
            }
            // Unknown prefixes come from a newer version; ignore them.
        }
        m.vaults.sort_by_key(|(o, _)| *o);
        Ok(Self { inner: Mutex::new(m) })
    }

    /// Everything written since the last call, ready for IndexedDB.
    pub fn take_changes(&self) -> Vec<Change> {
        let mut m = self.mem();
        let mut out = Vec::new();
        if std::mem::take(&mut m.cleared) {
            out.push(Change::Clear);
        }
        let dirty = std::mem::take(&mut m.dirty);
        for key in dirty {
            let value = Self::encode(&m, &key);
            out.push(match value {
                Some(v) => Change::Put(key, v),
                None => Change::Delete(key),
            });
        }
        out
    }

    /// True when [`WebStore::take_changes`] would return something.
    pub fn has_changes(&self) -> bool {
        let m = self.mem();
        m.cleared || !m.dirty.is_empty()
    }

    /// Every entry, as [`WebStore::load`] takes them (tests, export of the raw replica).
    pub fn entries(&self) -> Vec<(String, Vec<u8>)> {
        let m = self.mem();
        let mut keys: Vec<String> = m.meta.keys().map(|k| meta_key(k)).collect();
        keys.extend(m.vaults.iter().map(|(_, v)| vault_key(v.vault_id)));
        for (v, items) in &m.items {
            keys.extend(items.keys().map(|i| item_key(*v, *i)));
        }
        keys.extend(m.history.keys().map(|(v, i)| history_key(*v, *i)));
        keys.into_iter().filter_map(|k| Self::encode(&m, &k).map(|v| (k, v))).collect()
    }

    fn encode(m: &Mem, key: &str) -> Option<Vec<u8>> {
        if let Some(k) = key.strip_prefix("meta/") {
            return m.meta.get(k).cloned();
        }
        if let Some(id) = key.strip_prefix("vault/") {
            let id = Uuid::parse_str(id).ok()?;
            let (order, v) = m.vaults.iter().find(|(_, v)| v.vault_id == id)?;
            return Some(json(&VaultEntry {
                order: *order,
                vault_id: v.vault_id,
                sealed_key: v.sealed_key.clone(),
                meta: v.meta.clone(),
                cursor: v.cursor,
                manifest_version: v.manifest_version,
                manifest_hash: v.manifest_hash,
                integrity_alarm: v.integrity_alarm,
            }));
        }
        if let Some(rest) = key.strip_prefix("item/") {
            let (v, i) = two_ids(key, rest).ok()?;
            let li = m.items.get(&v)?.get(&i)?;
            return Some(json(&ItemEntry {
                record: li.record.clone(),
                pending: li.state == SyncState::Pending,
                base: li.base.clone(),
            }));
        }
        if let Some(rest) = key.strip_prefix("history/") {
            let ids = two_ids(key, rest).ok()?;
            let h = m.history.get(&ids)?;
            let h: Vec<HistoryEntryJson> = h
                .iter()
                .map(|h| HistoryEntryJson {
                    revision: h.revision,
                    wrapped_item_key: h.wrapped_item_key.clone(),
                    payload: h.payload.clone(),
                    saved_at: h.saved_at,
                })
                .collect();
            return Some(json(&h));
        }
        None
    }
}

impl Store for WebStore {
    fn get_meta(&self, key: &str) -> Result<Option<Vec<u8>>> {
        Ok(self.mem().meta.get(key).cloned())
    }

    fn put_meta(&self, key: &str, value: &[u8]) -> Result<()> {
        let mut m = self.mem();
        m.meta.insert(key.into(), value.to_vec());
        m.dirty.insert(meta_key(key));
        Ok(())
    }

    fn delete_meta(&self, key: &str) -> Result<()> {
        let mut m = self.mem();
        m.meta.remove(key);
        m.dirty.insert(meta_key(key));
        Ok(())
    }

    fn put_vault(&self, vault: &VaultRow) -> Result<()> {
        let mut m = self.mem();
        match m.vaults.iter_mut().find(|(_, v)| v.vault_id == vault.vault_id) {
            Some((_, v)) => *v = vault.clone(),
            None => {
                let order = m.next_order;
                m.next_order += 1;
                m.vaults.push((order, vault.clone()));
            }
        }
        m.dirty.insert(vault_key(vault.vault_id));
        Ok(())
    }

    fn vault(&self, vault: Uuid) -> Result<Option<VaultRow>> {
        Ok(self.mem().vaults.iter().find(|(_, v)| v.vault_id == vault).map(|(_, v)| v.clone()))
    }

    fn vaults(&self) -> Result<Vec<VaultRow>> {
        Ok(self.mem().vaults.iter().map(|(_, v)| v.clone()).collect())
    }

    fn put(&self, vault: Uuid, item: &LocalItem) -> Result<()> {
        let mut m = self.mem();
        m.items.entry(vault).or_default().insert(item.record.item_id, item.clone());
        m.dirty.insert(item_key(vault, item.record.item_id));
        Ok(())
    }

    fn get(&self, vault: Uuid, item: Uuid) -> Result<Option<LocalItem>> {
        Ok(self.mem().items.get(&vault).and_then(|v| v.get(&item)).cloned())
    }

    fn list(&self, vault: Uuid) -> Result<Vec<LocalItem>> {
        Ok(self.mem().items.get(&vault).map(|v| v.values().cloned().collect()).unwrap_or_default())
    }

    fn remove(&self, vault: Uuid, item: Uuid) -> Result<()> {
        let mut m = self.mem();
        if let Some(v) = m.items.get_mut(&vault) {
            v.remove(&item);
        }
        m.dirty.insert(item_key(vault, item));
        Ok(())
    }

    fn push_history(&self, vault: Uuid, item: Uuid, entry: &HistoryEntry) -> Result<()> {
        let mut m = self.mem();
        let h = m.history.entry((vault, item)).or_default();
        h.insert(0, entry.clone());
        h.truncate(HISTORY_LIMIT);
        m.dirty.insert(history_key(vault, item));
        Ok(())
    }

    fn history(&self, vault: Uuid, item: Uuid) -> Result<Vec<HistoryEntry>> {
        Ok(self.mem().history.get(&(vault, item)).cloned().unwrap_or_default())
    }

    fn wipe(&self) -> Result<()> {
        let mut m = self.mem();
        *m = Mem { cleared: true, ..Mem::default() };
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use kv_proto::{Alg, FORMAT_VERSION};

    use super::*;

    fn env(tag: u8) -> Envelope {
        Envelope {
            v: FORMAT_VERSION,
            alg: Alg::XChaCha20Poly1305,
            kid: Uuid::nil(),
            n: vec![tag; 24],
            ct: vec![tag; 20],
            enc: vec![],
        }
    }

    fn record(id: Uuid, rev: u64, tag: u8) -> ItemRecord {
        ItemRecord {
            item_id: id,
            revision: rev,
            deleted: false,
            wrapped_item_key: Some(env(tag)),
            payload: Some(env(tag + 1)),
            updated_by: Uuid::nil(),
            updated_at: "2026-01-01T00:00:00Z".into(),
        }
    }

    /// Applies changes to a fake IndexedDB object store.
    fn apply(db: &mut BTreeMap<String, Vec<u8>>, changes: Vec<Change>) {
        for c in changes {
            match c {
                Change::Clear => db.clear(),
                Change::Put(k, v) => {
                    db.insert(k, v);
                }
                Change::Delete(k) => {
                    db.remove(&k);
                }
            }
        }
    }

    #[test]
    fn changes_round_trip_through_a_reload() {
        let s = WebStore::new();
        let mut db = BTreeMap::new();
        let (v1, v2) = (Uuid::from_u128(9), Uuid::from_u128(1)); // v2 sorts first but was added second
        s.put_vault(&VaultRow::new(v1, env(1), Some(env(2)))).unwrap();
        s.put_vault(&VaultRow::new(v2, env(3), None)).unwrap();
        s.put_meta("account", b"{}").unwrap();
        let id = Uuid::new_v4();
        s.put_remote(v1, &record(id, 4, 10)).unwrap();
        s.put_local(v1, &record(id, 0, 20)).unwrap();
        s.push_history(v1, id, &HistoryEntry { revision: 4, wrapped_item_key: env(10), payload: env(11), saved_at: 5 })
            .unwrap();
        let gone = Uuid::new_v4();
        s.put_local(v1, &record(gone, 0, 40)).unwrap();
        s.remove(v1, gone).unwrap();
        assert!(s.has_changes());
        apply(&mut db, s.take_changes());
        assert!(!s.has_changes() && s.take_changes().is_empty());
        assert!(!db.contains_key(&item_key(v1, gone)));

        let r = WebStore::load(db.clone()).unwrap();
        assert_eq!(r.vaults().unwrap().iter().map(|v| v.vault_id).collect::<Vec<_>>(), vec![v1, v2]);
        assert_eq!(r.vault(v1).unwrap(), s.vault(v1).unwrap());
        let li = r.get(v1, id).unwrap().unwrap();
        assert_eq!(li.state, SyncState::Pending);
        assert_eq!(li.record.payload, Some(env(21)));
        assert_eq!(li.base.unwrap().revision, 4);
        assert_eq!(r.history(v1, id).unwrap(), s.history(v1, id).unwrap());
        assert_eq!(r.get_meta("account").unwrap().as_deref(), Some(&b"{}"[..]));
        let mut a = r.entries();
        a.sort();
        assert_eq!(a, db.clone().into_iter().collect::<Vec<_>>());

        // A vault added after the reload goes last.
        let v3 = Uuid::from_u128(0);
        r.put_vault(&VaultRow::new(v3, env(5), None)).unwrap();
        r.delete_meta("account").unwrap();
        apply(&mut db, r.take_changes());
        let r2 = WebStore::load(db.clone()).unwrap();
        assert_eq!(r2.vaults().unwrap().last().unwrap().vault_id, v3);
        assert!(r2.get_meta("account").unwrap().is_none());

        r2.wipe().unwrap();
        r2.put_meta("server", b"x").unwrap();
        let ch = r2.take_changes();
        assert_eq!(ch[0], Change::Clear);
        apply(&mut db, ch);
        assert_eq!(db.keys().collect::<Vec<_>>(), vec!["meta/server"]);
    }

    #[test]
    fn rejects_corrupt_entries_and_ignores_unknown_ones() {
        assert!(WebStore::load([("vault/x".to_string(), b"nope".to_vec())]).is_err());
        assert!(WebStore::load([("item/a/b".to_string(), b"{}".to_vec())]).is_err());
        assert!(WebStore::load([("future/x".to_string(), b"?".to_vec())]).unwrap().vaults().unwrap().is_empty());
    }

    #[test]
    fn never_holds_plaintext_fields() {
        let s = WebStore::new();
        let v = Uuid::new_v4();
        s.put_vault(&VaultRow::new(v, env(1), None)).unwrap();
        s.put_local(v, &record(Uuid::new_v4(), 0, 3)).unwrap();
        for (_, value) in s.entries() {
            let text = String::from_utf8(value).unwrap();
            for f in ["title", "username", "password", "\"url", "notes"] {
                assert!(!text.contains(f), "{text}");
            }
        }
    }
}