rose-squared-sdk 0.1.0

// PrivacyVault — The high-level API.
//
// This is the only type frontend developers need to interact with.
// It composes all lower-level modules behind a clean interface:
//
//   vault.add_document(keywords, doc_id)
//   vault.search(keyword)             → Vec<Uuid>
//   vault.delete_document(keyword, doc_id)
//   vault.export_state()              → Vec<u8>   (encrypted, safe to persist)
//   vault.import_state(blob)
//
// Thread safety: `PrivacyVault` is NOT `Send` in WASM (single-threaded runtime).
// In native builds with `std`, wrap in `Arc<Mutex<>>` for multi-threaded use.

use uuid::Uuid;

use crate::client::{
    state::ClientStateTable,
    updates::{UpdateEngine, hash_keyword},
};
use crate::crypto::kdf::MasterKeySet;
use crate::error::VaultError;
use crate::protocol::{
    search::{SearchProtocol, SearchResult},
    swissse::VolumeConfig,
};
use crate::server::edb::EncryptedStore;

// ── PrivacyVault ──────────────────────────────────────────────────────────────

pub struct PrivacyVault {
    pub keys:   MasterKeySet,
    pub state:  ClientStateTable,
    pub config: VolumeConfig,
}

impl PrivacyVault {
    // ── Construction ──────────────────────────────────────────────────────────

    /// Create a new vault from a user password and a 16-byte random salt.
    ///
    /// The salt must be generated once and stored alongside the encrypted
    /// state blob — it is NOT secret, but it must be consistent across sessions.
    ///
    /// ```rust
    /// use rose_squared_sdk::PrivacyVault;
    /// use rand::RngCore;
    /// let mut salt = [0u8; 16];
    /// rand::thread_rng().fill_bytes(&mut salt);
    /// let vault = PrivacyVault::new("correct-horse-battery-staple", &salt, Default::default());
    /// ```
    pub fn new(
        password: &str,
        salt:     &[u8; 16],
        config:   VolumeConfig,
    ) -> Result<Self, VaultError> {
        let keys  = MasterKeySet::derive(password, salt)?;
        let state = ClientStateTable::new();
        Ok(Self { keys, state, config })
    }

    /// Restore a vault from a previously exported state blob.
    pub fn from_exported(
        password: &str,
        salt:     &[u8; 16],
        blob:     &[u8],
        config:   VolumeConfig,
    ) -> Result<Self, VaultError> {
        let keys  = MasterKeySet::derive(password, salt)?;
        let state = ClientStateTable::import_encrypted(blob, &keys)?;
        Ok(Self { keys, state, config })
    }

    // ── Write operations ──────────────────────────────────────────────────────

    /// Index a document under one or more keywords.
    ///
    /// Sends one EDB entry per keyword to the store.
    /// The store never sees the keywords or the document ID in plaintext.
    ///
    /// `doc_id` should be your application's stable identifier for the document
    /// (e.g., the UUID of the file in your encrypted document store).
    pub async fn add_document<S: EncryptedStore>(
        &mut self,
        keywords: &[&str],
        doc_id:   Uuid,
        store:    &S,
    ) -> Result<(), VaultError> {
        let engine = UpdateEngine::new(&self.keys);

        for &keyword in keywords {
            let kh    = hash_keyword(keyword);
            let kw_st = self.state.get_or_create(kh);
            let entry = engine.prepare_add(keyword.as_bytes(), doc_id, kw_st)?;

            // Use padded write for SWiSSSE volume hiding.
            store.padded_put_batch(vec![entry], self.config.n_max).await?;
        }

        Ok(())
    }

    /// Remove a document from one keyword's result set.
    ///
    /// This performs a Backward-Security Type-II delete:
    ///   • The epoch for this keyword is bumped.
    ///   • All surviving entries are atomically re-written under the new epoch.
    ///   • Old epoch entries are deleted.
    ///
    /// After this call, any previously issued search tokens for this keyword
    /// are invalid — they address old-epoch tags which are now gone.
    pub async fn delete_document<S: EncryptedStore>(
        &mut self,
        keyword: &str,
        doc_id:  Uuid,
        store:   &S,
    ) -> Result<(usize, usize), VaultError> {
        let engine = UpdateEngine::new(&self.keys);
        let kh     = hash_keyword(keyword);
        let kw_st  = self.state.get_or_create(kh);

        let batch = engine.prepare_delete(keyword.as_bytes(), doc_id, kw_st)?;

        let removes_count = batch.removes.len();
        let adds_count = batch.adds.len();

        // Atomic: remove old tags, write new-epoch entries.
        store.atomic_update(batch.adds, batch.removes).await?;

        Ok((removes_count, adds_count))
    }

    // ── Search ────────────────────────────────────────────────────────────────

    /// Search for all documents indexed under `keyword`.
    ///
    /// The keyword never leaves the client in plaintext.
    /// The server returns opaque ciphertexts; the client decrypts them here.
    ///
    /// Returns document UUIDs sorted newest-first.
    ///
    /// With `SWiSSSE` enabled (default), every search fetches exactly
    /// `n_max` tags from the server, hiding the true result count.
    pub async fn search<S: EncryptedStore>(
        &self,
        keyword: &str,
        store:   &S,
    ) -> Result<Vec<Uuid>, VaultError> {
        let results = self.search_with_metadata(keyword, store).await?;
        Ok(results.into_iter().map(|r| r.doc_id).collect())
    }

    /// Search and return full `SearchResult` (doc_id + timestamp).
    pub async fn search_with_metadata<S: EncryptedStore>(
        &self,
        keyword: &str,
        store:   &S,
    ) -> Result<Vec<SearchResult>, VaultError> {
        let proto = SearchProtocol::new(&self.keys, &self.state);
        
        let token = match proto.prepare_search(keyword)? {
            Some(t) => t,
            None    => {
                // To hide the fact that a keyword doesn't exist, we SHOULD 
                // perform a dummy search. For now, we'll return empty.
                return Ok(vec![]);
            }
        };

        // SWiSSSE: Pad the token to n_max tags to hide volume.
        let padded_token = crate::protocol::swissse::pad_search_token(token, &self.config)?;
        
        // Fetch results (batch)
        let enc_values = proto.fetch_token_results(&padded_token, store).await?;
        
        // Decrypt and finalize
        proto.finalize_search(&padded_token, enc_values)
    }

    // ── State persistence ─────────────────────────────────────────────────────

    /// Export the encrypted client state as a byte blob.
    ///
    /// Store this in IndexedDB, a file, or any persistent medium.
    /// It is AES-256-GCM encrypted with K_state — safe to store in the cloud.
    pub fn export_state(&self) -> Result<Vec<u8>, VaultError> {
        self.state.export_encrypted(&self.keys)
    }
}