cipherstash-client 0.42.0

//! The SteVec document — the encrypted, searchable representation of a JSON
//! value and its on-the-wire format.
//!
//! A [`SteVec`] is what an encrypted JSON column stores. Indexing flattens the
//! document into path/value pairs (see the parent [`super`] module) and this
//! module turns those pairs into encrypted [`EncryptedEntry`]s and lays them
//! out for storage and query.
//!
//! # Two entries per node
//!
//! Every JSON node emits a **pair** of entries, so an N-node document has 2N
//! entries:
//!
//! - a **path entry** `{s: SEL(path), c, op?}` — its selector MACs the JSON
//!   path; `c` encrypts the node's value; `op` is the orderable term (present
//!   only on number/string nodes, backing range queries). Ordering is
//!   deliberately lossy (OPE/ORE collation, f64 rounding), so `op` serves
//!   ordering, never equality.
//! - a **value entry** `{s: SEL(tag ‖ path ‖ canonical(value)), c}` — its
//!   selector bakes the value into the MAC, so the selector's *presence* in a
//!   stored document is the equality match for scalar values, while containers
//!   are matched through their flattened containment entries. This avoids the
//!   orderable encoding's text/integer collisions. High-precision decimal
//!   exactness is limited because arbitrary-precision canonicalization is not
//!   currently supported.
//!   Its `c` encrypts a fixed sentinel (see below), not the value — the value
//!   is already committed by the selector.
//!
//! The root entry (index 0) is the whole-document ciphertext; value entries
//! append after all path entries, so the root position is stable.
//!
//! # Envelope wire format
//!
//! A document encrypts all of its entries under **one** ZeroKMS data key. Two
//! consequences shape the wire format:
//!
//! 1. **The key header is hoisted once.** The key-retrieval material — IV, key
//!    tag, descriptor, keyset id — is identical for every entry, so it is
//!    stored a single time in a per-document [`KeyHeader`] ([`SteVec::header`])
//!    rather than repeated inside every entry's ciphertext. An entry's `c` is
//!    therefore the **raw AEAD output only** ([`EncryptedEntry::ciphertext`],
//!    base85 on the wire), not a self-describing record. An entry is not
//!    decryptable on its own: decryption reassembles the header + the entry's
//!    `c` + the entry's nonce into a [`RecordWithNonce`] (see
//!    [`KeyHeader::record_with_selector`]).
//!
//! 2. **Each entry's nonce is derived from its selector.** Because every entry
//!    shares the document's one data key, a shared nonce would make AES-GCM-SIV
//!    encryption deterministic *within the document* — equal values at
//!    different paths, and the constant value-entry plaintext, would produce
//!    byte-identical ciphertexts, leaking equality the value selectors were
//!    designed to hide. Instead this module derives the nonce from
//!    [`TokenizedSelector::aead_nonce`] (`selector[..12]`) and binds the full
//!    selector into the AAD. Ciphertext equality then requires selector
//!    equality *and* plaintext equality — and selector equality is
//!    already public (it is what containment matches), so no equality leaks
//!    that the index does not already announce. As a bonus the nonce is
//!    authenticated (GCM-SIV feeds it into the tag) and re-derived from the
//!    stored selector on decrypt, so a ciphertext grafted onto a different
//!    entry within the same document fails to decrypt.
//!
//! No nonce is stored anywhere — it is always re-derived from `s`.
//!
//! # The value-entry sentinel
//!
//! A value entry's `c` encrypts [`VALUE_ENTRY_SENTINEL`], a fixed versioned
//! byte string, rather than the empty string. `Enc("")` would be ambiguous
//! with a genuine empty-string leaf (`{"a": ""}`), and its reserved variant
//! byte means the sentinel never parses as a legal [`crate::encryption::Plaintext`].
//! The `c` slot is kept (rather than dropped) both to keep the entry shape
//! uniform and to leave room for a future node-local-encryption linkage token.
//!
//! Scalar (non-SteVec) payloads are unaffected by all of the above: they carry
//! one self-describing record whose `iv[..12]` is already a unique nonce.

mod encrypted_entry;
mod encrypted_term;
pub(super) mod priv_state;
mod selector;
mod ste_plaintext_vec;

use crate::{ejsonpath::Selector, encryption::EncryptionError};
use crate::{
    encryption::Plaintext,
    zerokms::{self, DataKeyWithTag, KeyHeader, RecordWithNonce},
};
use aes_gcm_siv::{aead::Payload, Nonce};
use itertools::Itertools;
use serde::{Deserialize, Serialize};
use uuid::Uuid;

pub use encrypted_entry::EncryptedEntry;
pub use ste_plaintext_vec::StePlaintextVec;

pub use encrypted_term::EncryptedSteVecTerm;
pub use selector::TokenizedSelector;

/// The plaintext a **value entry** encrypts: a fixed, versioned sentinel.
///
/// Layout mirrors the [`Plaintext`] byte encoding's `[version, variant, ...]`
/// header: byte 0 is the plaintext-encoding version, byte 1 is the variant —
/// `0x7f`, reserved here and deliberately outside the allocated variant range,
/// so [`Plaintext::from_slice`] can never confuse a value entry with a legal
/// value — and byte 2 is the sentinel's own version.
///
/// Why encrypt a sentinel rather than `""` (the previous scheme) or nothing:
/// - `Enc("")` is ambiguous with a *genuine* empty-string leaf — `{"a": ""}`
///   produces a path entry whose ciphertext has the same plaintext as every
///   value entry, so a decryptor cannot tell "value entry, skip" from "the
///   field really is empty". The sentinel is assertable on decrypt.
/// - Dropping the ciphertext entirely would structurally mark value entries
///   (`{s}` vs `{s, c}`), revealing which selectors are value selectors
///   without any query traffic — and the `c` slot is where a future
///   node-local-encryption linkage token would live.
pub const VALUE_ENTRY_SENTINEL: [u8; 3] = [1, 0x7f, 1];

/// True when `bytes` (a decrypted entry plaintext) is the value-entry
/// sentinel — i.e. the entry was a value-inclusive selector, not a node.
pub fn is_value_entry_sentinel(bytes: &[u8]) -> bool {
    bytes == VALUE_ENTRY_SENTINEL
}

/// Describes a SteVec, which is a vector of encrypted entries used to index JSON.
/// `N` is the length of the selector.
///
/// The envelope wire format: ONE [`KeyHeader`] carries the document data
/// key's retrieval material (iv, tag, descriptor, keyset), and every entry
/// carries only its raw AEAD output — encrypted under that single key with a
/// nonce derived from the entry's tokenized selector
/// ([`TokenizedSelector::aead_nonce`]). Nothing per-entry is self-decryptable
/// on its own: decryption reassembles header + entry `c` + selector nonce
/// (see [`KeyHeader::record_with_selector`]).
#[derive(Debug, Serialize, Deserialize)]
pub struct SteVec<const N: usize> {
    pub(super) header: KeyHeader,
    pub(super) entries: Vec<EncryptedEntry<N>>,
}

impl<const N: usize> SteVec<N> {
    /// The document's key header (stored once at the envelope).
    pub fn header(&self) -> &KeyHeader {
        &self.header
    }

    /// The per-selector encrypted entries. `entries()[0]` is the root entry.
    pub fn entries(&self) -> &[EncryptedEntry<N>] {
        &self.entries
    }
}

// The root-record helpers are inherently tied to the 16-byte selector the
// envelope uses (the reassembled [`RecordWithNonce`] and its AAD binding are
// 16-byte), so they live on the concrete `SteVec<16>`.
impl SteVec<16> {
    /// Consumes the SteVec and returns the root entry's decryptable form:
    /// the key header's record material + the root ciphertext + the root
    /// selector (source of both the nonce and the AAD binding). The root entry
    /// is always first.
    pub fn into_root_record_with_nonce(self) -> Result<RecordWithNonce, EncryptionError> {
        let root = self.entries.into_iter().next().ok_or_else(|| {
            EncryptionError::InvalidValue("SteVec missing root record".to_string())
        })?;
        Ok(self
            .header
            .record_with_selector(root.ciphertext, root.tokenized_selector.as_bytes()))
    }

    /// The root entry's selector-derived AEAD nonce.
    /// If the SteVec is empty, an error is returned.
    pub fn root_nonce(&self) -> Result<[u8; 12], EncryptionError> {
        self.entries
            .first()
            .map(|entry| entry.tokenized_selector.aead_nonce())
            .ok_or_else(|| EncryptionError::InvalidValue("SteVec missing root record".to_string()))
    }

    /// The root entry's full 16-byte tokenized selector, bound into the AAD of
    /// the root ciphertext. If the SteVec is empty, an error is returned.
    pub fn root_selector(&self) -> Result<[u8; 16], EncryptionError> {
        self.entries
            .first()
            .map(|entry| entry.tokenized_selector.as_bytes())
            .ok_or_else(|| EncryptionError::InvalidValue("SteVec missing root record".to_string()))
    }
}

impl IntoIterator for SteVec<16> {
    type Item = EncryptedEntry<16>;
    type IntoIter = std::vec::IntoIter<Self::Item>;

    fn into_iter(self) -> Self::IntoIter {
        self.entries.into_iter()
    }
}

/// Represents a vector of encrypted entries that are pending encryption.
/// It can either be encrypted or turned into a query vector (which will discard the plaintexts, retaining on the selectors and terms).
pub struct SteVecPendingEncryption<const N: usize>(Vec<EntryWithEncryptedTerm<N>>);

impl SteVecPendingEncryption<16> {
    /// Encrypt every entry under the single document `key`, emitting the
    /// envelope form: one [`KeyHeader`] + slim entries. The `descriptor` is
    /// the column's canonical descriptor — the same string the data key was
    /// generated under (it participates in both key retrieval and the AEAD
    /// AAD, so a document grafted onto a different field fails at
    /// retrieve-key).
    pub fn encrypt(
        self,
        key: DataKeyWithTag,
        descriptor: &str,
        keyset_id: Option<Uuid>,
    ) -> Result<SteVec<16>, zerokms::Error> {
        let header = KeyHeader::new(&key, descriptor, keyset_id);
        let entries = self
            .0
            .into_iter()
            .map(|entry| entry.build_final(&key, descriptor))
            .try_collect()?;
        Ok(SteVec { header, entries })
    }
}

impl<const N: usize> SteVecPendingEncryption<N> {
    pub fn into_query(self) -> SteQueryVec<N> {
        SteQueryVec::from(self)
    }
}

/// What one wire entry encrypts.
enum EntryPlaintext {
    /// A path entry encrypts its node's value.
    Node(Plaintext),
    /// A value entry encrypts the fixed [`VALUE_ENTRY_SENTINEL`] — its
    /// selector already commits to (path, value), so the ciphertext exists
    /// only to keep the entry shape uniform and stay assertable on decrypt.
    ValueSentinel,
}

struct EntryWithEncryptedTerm<const N: usize> {
    tokenized_selector: TokenizedSelector<N>,
    /// The orderable term — present only on the path entries of number/string
    /// nodes. Value entries and non-orderable path entries carry no term:
    /// exact matching is selector presence.
    term: Option<EncryptedSteVecTerm>,
    plaintext: EntryPlaintext,
    parent_is_array: bool,
}

impl EntryWithEncryptedTerm<16> {
    fn build_final(
        self,
        key: &DataKeyWithTag,
        descriptor: &str,
    ) -> Result<EncryptedEntry<16>, zerokms::Error> {
        let bytes = match &self.plaintext {
            EntryPlaintext::Node(plaintext) => plaintext.to_vec(),
            EntryPlaintext::ValueSentinel => VALUE_ENTRY_SENTINEL.to_vec(),
        };
        // The entry binds to its selector, NOT the key IV — every entry of the
        // document shares the one key, and the selector-derived nonce is what
        // keeps equal plaintexts at different (path, value) positions from
        // producing equal ciphertexts. The full selector is also bound into the
        // AAD, so the ciphertext is authenticated against the whole selector.
        // SteVec owns this nonce/AAD policy; the general EncryptionTarget
        // remains IV-derived and unaware of JSON selectors.
        let selector = self.tokenized_selector.as_bytes();
        let mut aad = zerokms::build_aad(descriptor, &key.tag);
        aad.extend_from_slice(&selector);
        let nonce = Nonce::from_slice(&selector[..12]);
        let ciphertext = zerokms::encrypt_aead(
            key.key(),
            nonce,
            Payload {
                msg: &bytes,
                aad: &aad,
            },
        )?;
        Ok(EncryptedEntry::new(
            self.tokenized_selector,
            self.term,
            ciphertext,
            self.parent_is_array,
        ))
    }
}

#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
/// Represents one entry in a query Vec — a selector and, for orderable
/// (number/string) path entries, the orderable term. Presence-only entries
/// (value selectors, structural nodes) carry `None`: the DB matches them on
/// selector equality alone.
pub struct QueryEntry<const N: usize>(
    pub(super) TokenizedSelector<N>,
    pub(super) Option<EncryptedSteVecTerm>,
);

/// Converts an [EntryWithEncryptedTerm] into a [QueryEntry], discarding the plaintext.
impl<const N: usize> From<EntryWithEncryptedTerm<N>> for QueryEntry<N> {
    fn from(entry: EntryWithEncryptedTerm<N>) -> Self {
        Self(entry.tokenized_selector, entry.term)
    }
}

// FIXME: This is only Clone because IndexTerm is Clone (which also needs to be fixed)
/// Represents a vector of query entries.
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct SteQueryVec<const N: usize>(pub(super) Vec<QueryEntry<N>>);

impl<const N: usize> From<SteVecPendingEncryption<N>> for SteQueryVec<N> {
    fn from(pending: SteVecPendingEncryption<N>) -> Self {
        // TODO: We need to find the _most_ specific term
        Self(pending.0.into_iter().map(QueryEntry::from).collect())
    }
}

#[cfg(feature = "tokio")]
impl<const N: usize> SteQueryVec<N> {
    /// Build the selector-only containment needle used for exact equality at
    /// one JSON path.
    pub(crate) fn from_value_selector(selector: TokenizedSelector<N>) -> Self {
        Self(vec![QueryEntry(selector, None)])
    }
}

impl<const N: usize> QueryEntry<N> {
    /// Decomposes the entry into its `(selector, optional term)` pair,
    /// consuming it.
    ///
    /// Lets consumers outside this module (e.g. the EQL v3 wire assembler)
    /// build per-entry query shapes without reaching into the tuple fields.
    pub fn into_parts(self) -> (TokenizedSelector<N>, Option<EncryptedSteVecTerm>) {
        (self.0, self.1)
    }
}

impl<const N: usize> IntoIterator for SteQueryVec<N> {
    type Item = QueryEntry<N>;
    type IntoIter = std::vec::IntoIter<QueryEntry<N>>;

    fn into_iter(self) -> Self::IntoIter {
        self.0.into_iter()
    }
}

#[cfg(test)]
mod tests {
    use crate::encryption::json_indexer::{
        prefix_mac::{PrefixMac, UpdatePrefixMac},
        ste_vec::{encrypted_term::EncryptedSteVecTerm, EncryptedEntry, TokenizedSelector},
    };
    use std::fmt::Debug;

    #[macro_export]
    macro_rules! boxed {
        ($($literal:expr),*) => {
            &[$(Box::new($literal)),*]
        };
    }

    #[macro_export]
    macro_rules! assert_was_finalized_with {
        ($macca:expr, $($term:expr),*) => {
            $macca.assert_was_finalized_with($crate::boxed!($($term),*))
        };
    }

    #[derive(Default)]
    pub struct TestPrefixMac {
        pub pending_terms: Vec<Box<dyn Debug>>,
        pub finalized_terms: Vec<Box<dyn Debug>>,
    }

    impl TestPrefixMac {
        pub fn assert_was_finalized_with(&self, terms: &[Box<dyn Debug>]) {
            let x = self
                .finalized_terms
                .iter()
                .map(|term| format!("{term:?}"))
                .collect::<Vec<_>>();

            let y = terms
                .iter()
                .map(|term| format!("{term:?}"))
                .collect::<Vec<_>>();

            assert_eq!(x, y, "Expected MAC updates to be {y:#?}, got {x:#?}");
        }
    }

    impl PrefixMac for TestPrefixMac {
        fn finalize_reset<const N: usize>(&mut self) -> [u8; N] {
            // Accumulate across finalizes: pipelines that tokenize more than
            // one selector per node (path + value) assert the full update
            // sequence in order.
            self.finalized_terms.append(&mut self.pending_terms);
            [0; N]
        }
    }

    impl UpdatePrefixMac<String> for TestPrefixMac {
        fn update(&mut self, value: String) {
            self.pending_terms.push(Box::new(value));
        }
    }

    impl UpdatePrefixMac<&str> for TestPrefixMac {
        fn update(&mut self, value: &str) {
            self.update(value.to_string());
        }
    }

    impl UpdatePrefixMac<[u8; 16]> for TestPrefixMac {
        fn update(&mut self, value: [u8; 16]) {
            self.pending_terms.push(Box::new(value));
        }
    }

    /// Raw AEAD output — under the envelope wire format an entry's `c`
    /// carries no key material.
    fn ciphertext() -> Vec<u8> {
        vec![1; 32]
    }

    #[test]
    fn test_serde_encrypted_entry() {
        use cllw_ore::OpeCllw8VariableV1;

        let entry = EncryptedEntry::new(
            TokenizedSelector([1; 16]),
            Some(EncryptedSteVecTerm::Compat {
                op: OpeCllw8VariableV1::from_bytes(vec![2; 16]),
            }),
            ciphertext(),
            false,
        );

        let serialized = serde_json::to_string(&entry).unwrap();
        assert_eq!(
            serde_json::from_str::<EncryptedEntry<16>>(&serialized).unwrap(),
            entry
        );
    }

    #[test]
    fn test_serde_encrypted_entry_term_less() {
        // A presence-only entry (a value entry, or a structural/bool/null path
        // entry) has no term at all — `{s, c, parent_is_array}` must
        // round-trip with the term absent from the wire, not null.
        let entry = EncryptedEntry::new(TokenizedSelector([3; 16]), None, ciphertext(), false);

        let serialized = serde_json::to_string(&entry).unwrap();
        assert!(
            !serialized.contains("\"hm\""),
            "no MAC term key: {serialized}"
        );
        assert!(
            !serialized.contains("\"op\""),
            "no orderable key: {serialized}"
        );
        assert!(
            !serialized.contains("null"),
            "term must be absent, not null: {serialized}"
        );
        assert_eq!(
            serde_json::from_str::<EncryptedEntry<16>>(&serialized).unwrap(),
            entry
        );
    }
}