rose-squared-sdk 0.1.0

// Full Search Round-Trip: client → server → client.
//
// Protocol steps (per RO(SE)² §4):
//
//   1. Client calls `prepare_search(keyword)`:
//        - Hashes the keyword to look up its KeywordState.
//        - Generates a SearchToken from the live indices.
//        - Returns the token (a list of pseudorandom tag/key pairs).
//
//   2. Client sends the token's tags to the server.
//      The server fetches each EncValue by tag (O(1) per fetch) and returns them.
//
//   3. Client calls `finalize_search(token, enc_values)`:
//        - Decrypts each EncValue with the corresponding per-entry key.
//        - Verifies the GCM tag — rejects tampered entries (`VaultError::Tampered`).
//        - Extracts the doc_id from each EntryPayload.
//        - Returns the set of document UUIDs matching the keyword.
//
// The server observes only a set of random-looking byte strings.
// It learns: (a) the number of results (result-size leakage, acceptable under
// standard SSE security); and (b) access patterns across queries ONLY if the
// same keyword is searched twice — use token rotation to mitigate (Phase 4).

use uuid::Uuid;

use crate::client::{
    state::ClientStateTable,
    trapdoor::TrapdoorEngine,
    updates::hash_keyword,
};
use crate::crypto::{
    aead,
    kdf::MasterKeySet,
    primitives::{EncValue, EntryPayload, SearchToken},
};
use crate::error::VaultError;
use crate::server::edb::EncryptedStore;

// ── Search coordinator ────────────────────────────────────────────────────────

pub struct SearchProtocol<'a> {
    keys:  &'a MasterKeySet,
    state: &'a ClientStateTable,
}

impl<'a> SearchProtocol<'a> {
    pub fn new(keys: &'a MasterKeySet, state: &'a ClientStateTable) -> Self {
        Self { keys, state }
    }

    // ── Phase 1: generate token ───────────────────────────────────────────────

    /// Generate a single-use search token for `keyword`.
    ///
    /// Returns `Ok(None)` if the keyword has no live results (never written or
    /// all documents deleted), so the caller can skip the network round-trip.
    pub fn prepare_search(&self, keyword: &str) -> Result<Option<SearchToken>, VaultError> {
        let kh = hash_keyword(keyword);

        let state = match self.state.get(&kh) {
            Some(s) => s,
            None    => return Ok(None),  // keyword not in index
        };

        if state.live_indices.is_empty() {
            return Ok(None);  // all results deleted
        }

        let engine = TrapdoorEngine::new(self.keys);
        let token  = engine.generate_search_token(keyword.as_bytes(), state);
        Ok(Some(token))
    }

    // ── Phase 2: server fetch (caller's responsibility) ───────────────────────
    //
    // The caller fetches enc_values from the store using token.pairs.
    // We provide `fetch_token_results` as a convenience that calls the store.

    /// Convenience: fetch all tagged entries for a token from an EDB in one call.
    pub async fn fetch_token_results<S: EncryptedStore>(
        &self,
        token: &SearchToken,
        store: &S,
    ) -> Result<Vec<Option<EncValue>>, VaultError> {
        let tags: Vec<_> = token.pairs.iter().map(|(t, _)| t.clone()).collect();
        store.get_batch(&tags).await
    }

    // ── Phase 3: decrypt and assemble results ─────────────────────────────────

    /// Decrypt and verify the server's response.
    ///
    /// Each `enc_values[i]` corresponds to `token.pairs[i]`.
    /// Entries that the server reports as missing are silently skipped
    /// (the server may have lost them; this is the "robustness" the paper addresses).
    /// Entries that fail GCM verification return `VaultError::Tampered`.
    pub fn finalize_search(
        &self,
        token:      &SearchToken,
        enc_values: Vec<Option<EncValue>>,
    ) -> Result<Vec<SearchResult>, VaultError> {
        let mut results = Vec::new();

        for ((_, key), maybe_enc) in token.pairs.iter().zip(enc_values.into_iter()) {
            let enc = match maybe_enc {
                Some(e) => e,
                None    => continue,   // entry missing from server (robustness case)
            };

            let plaintext = aead::decrypt(key, &enc)?;   // Tampered returns Err here
            let payload: EntryPayload = bincode::deserialize(&plaintext)?;

            results.push(SearchResult {
                doc_id:    payload.doc_id,
                timestamp: payload.timestamp,
            });
        }

        // Sort newest first.
        results.sort_by(|a, b| b.timestamp.cmp(&a.timestamp));
        Ok(results)
    }

    // ── One-shot convenience ──────────────────────────────────────────────────

    /// Full search: prepare token → fetch from store → decrypt results.
    ///
    /// Equivalent to calling `prepare_search` + `fetch_token_results` +
    /// `finalize_search` in sequence.  Useful for single-threaded environments.
    pub async fn search<S: EncryptedStore>(
        &self,
        keyword: &str,
        store:   &S,
    ) -> Result<Vec<SearchResult>, VaultError> {
        let token = match self.prepare_search(keyword)? {
            Some(t) => t,
            None    => return Ok(vec![]),
        };

        let enc_values = self.fetch_token_results(&token, store).await?;
        self.finalize_search(&token, enc_values)
    }
}

// ── Search result ─────────────────────────────────────────────────────────────

/// One document matching the searched keyword.
#[derive(Debug, Clone)]
pub struct SearchResult {
    /// The unique identifier of the matching document.
    pub doc_id:    Uuid,
    /// Write timestamp in ms since UNIX epoch (for ordering, not displayed in clear to server).
    pub timestamp: u64,
}