rose-squared-sdk 0.1.0

// Client-side state management.
//
// The client is the ONLY entity that ever sees plaintext keyword state.
// This module manages two layers:
//
//   KeywordState   — per-keyword counters and live-index bookkeeping
//   ClientStateTable — the full map, serialised and AES-GCM encrypted before
//                      any persistence (IndexedDB, localStorage, file, etc.)
//
// Design decisions (locked):
//   • Live indices are explicit (Vec<u64>) rather than a single counter.
//     This is required for Backward Security Type-II: on delete we re-derive
//     tags only for *surviving* indices, not all historical ones.
//   • The deletion epoch is per-keyword.  A delete on keyword "invoice" does
//     not force re-keying of keyword "report".
//   • The full state table is encrypted with K_state before leaving the client,
//     so even if the persistence layer is compromised the state is opaque.

use std::collections::HashMap;
use serde::{Deserialize, Serialize};
use zeroize::{Zeroize, ZeroizeOnDrop};
use uuid::Uuid;

use crate::crypto::{aead, primitives::Epoch};
use crate::crypto::kdf::MasterKeySet;
use crate::error::VaultError;

// ── Per-keyword state ─────────────────────────────────────────────────────────

/// Everything the client knows about one keyword.
///
/// Serialised as part of `ClientStateTable` and encrypted with K_state.
#[derive(Clone, Debug, Serialize, Deserialize, Zeroize, ZeroizeOnDrop)]
pub struct KeywordState {
    /// Indices of *live* (non-deleted) entries for this keyword in the EDB.
    /// Each index `i` maps to EDB tag:  PRF(K_tag, keyword || i || epoch)
    /// This is the authoritative source for O(1) search token generation.
    pub live_indices: Vec<u64>,

    /// Maps EDB index → document UUID.  Used during delete to find which
    /// indices to evict from `live_indices`.
    #[zeroize(skip)]
    pub index_to_doc: HashMap<u64, [u8; 16]>,

    /// Monotonically increasing total write counter (never decremented).
    /// Provides unique, fresh indices for every add even after deletes.
    pub total_writes: u64,

    /// Current epoch.  Incremented on every delete for this keyword.
    /// Old tags (from prior epochs) become permanently unreachable — this
    /// is the mechanism behind Backward Security Type-II.
    pub epoch: Epoch,
}

impl KeywordState {
    /// Create fresh state for a keyword that has never been written.
    pub fn new() -> Self {
        Self {
            live_indices: Vec::new(),
            index_to_doc: HashMap::new(),
            total_writes:  0,
            epoch:         0,
        }
    }

    /// Allocate the next fresh index for an add operation.
    /// Returns the index; callers must push it into `live_indices` themselves.
    pub fn next_index(&mut self) -> u64 {
        let idx = self.total_writes;
        self.total_writes += 1;
        idx
    }

    /// Register a newly written index → doc mapping.
    pub fn record_add(&mut self, index: u64, doc_id: Uuid) {
        let bytes: [u8; 16] = *doc_id.as_bytes();
        self.live_indices.push(index);
        self.index_to_doc.insert(index, bytes);
    }

    /// Remove all indices that map to `doc_id`.
    /// Returns `true` if at least one index was removed (doc was indexed here).
    pub fn evict_doc(&mut self, doc_id: Uuid) -> bool {
        let target: [u8; 16] = *doc_id.as_bytes();
        let before = self.live_indices.len();
        self.live_indices.retain(|idx| {
            self.index_to_doc.get(idx) != Some(&target)
        });
        // Clean up the reverse map too.
        self.index_to_doc.retain(|_, v| v != &target);
        // Bump epoch — all old tags are now retired.
        if self.live_indices.len() < before {
            self.epoch += 1;
            true
        } else {
            false
        }
    }
}

impl Default for KeywordState {
    fn default() -> Self { Self::new() }
}

// ── Full state table ──────────────────────────────────────────────────────────

/// The client's complete encrypted state: a map from H(keyword) → KeywordState.
///
/// Stored encrypted on the client device.  The encryption key is K_state from
/// the MasterKeySet, which is derived from the user's password and never leaves
/// the client.
pub struct ClientStateTable {
    /// Keyed by SHA-256(keyword) to avoid leaking keyword lengths at rest.
    inner: HashMap<[u8; 32], KeywordState>,
}

impl ClientStateTable {
    /// Create an empty state table (new vault).
    pub fn new() -> Self {
        Self { inner: HashMap::new() }
    }

    /// Get or create the state for a keyword.
    pub fn get_or_create(&mut self, keyword_hash: [u8; 32]) -> &mut KeywordState {
        self.inner.entry(keyword_hash).or_default()
    }

    /// Get an existing state (read-only).
    pub fn get(&self, keyword_hash: &[u8; 32]) -> Option<&KeywordState> {
        self.inner.get(keyword_hash)
    }

    // ── Persistence ───────────────────────────────────────────────────────────

    /// Serialise + AES-256-GCM encrypt the entire table.
    /// The returned bytes can be stored anywhere (IndexedDB, disk, etc.).
    pub fn export_encrypted(&self, keys: &MasterKeySet) -> Result<Vec<u8>, VaultError> {
        let plaintext = bincode::serialize(&self.inner)?;
        let enc = aead::encrypt(keys.k_state.as_bytes(), &plaintext)?;
        Ok(enc.0)
    }

    /// Decrypt and deserialise a blob produced by `export_encrypted`.
    pub fn import_encrypted(blob: &[u8], keys: &MasterKeySet) -> Result<Self, VaultError> {
        use crate::crypto::primitives::EncValue;
        let plaintext = aead::decrypt(keys.k_state.as_bytes(), &EncValue(blob.to_vec()))?;
        let inner: HashMap<[u8; 32], KeywordState> = bincode::deserialize(&plaintext)?;
        Ok(Self { inner })
    }
}

impl Default for ClientStateTable {
    fn default() -> Self { Self::new() }
}