cipherstash-client 0.34.1-alpha.4

use crate::zerokms::IndexKey;
use cllw_ore::{CllwOpeEncrypt, Key as OpeKey};
use zerokms_protocol::cipherstash_config::column;

use super::{
    errors::EncryptionError,
    indexer::{IndexerInit, Indexes, IndexesForQuery, QueryOp},
    IndexTerm, Plaintext, QueryBuilder, StorageBuilder,
};

pub struct OpeIndexer;

/// No options are currently supported
#[derive(Debug, Default)]
pub struct OpeIndexerOptions;

impl IndexerInit for OpeIndexer {
    type Args = OpeIndexerOptions;
    type Error = EncryptionError;

    fn try_init<A>(_opts: A) -> Result<Self, Self::Error>
    where
        Self::Args: TryFrom<A, Error = Self::Error>,
    {
        Ok(Self)
    }
}

impl<'k> Indexes<'k, Plaintext> for OpeIndexer {
    fn index(
        &self,
        mut builder: StorageBuilder<'k, Plaintext>,
    ) -> Result<StorageBuilder<'k, Plaintext>, EncryptionError> {
        let index_term = self.encrypt(builder.plaintext(), builder.index_key())?;
        builder.add_index_term(index_term);

        Ok(builder)
    }
}

impl<C> IndexesForQuery<Plaintext, C> for OpeIndexer {
    fn query_index(
        &self,
        builder: QueryBuilder<Plaintext, C>,
        _op: QueryOp,
    ) -> Result<IndexTerm, EncryptionError> {
        let index_term = self.encrypt(builder.plaintext(), builder.index_key())?;
        Ok(index_term)
    }
}

impl Default for OpeIndexer {
    fn default() -> Self {
        Self
    }
}

impl TryFrom<&column::IndexType> for OpeIndexerOptions {
    type Error = EncryptionError;

    fn try_from(value: &column::IndexType) -> Result<Self, Self::Error> {
        match value {
            column::IndexType::Ope => Ok(Default::default()),
            _ => Err(EncryptionError::IndexingError(
                "OpeIndexerOptions can only be created from an Ope index configuration".to_string(),
            )),
        }
    }
}

impl OpeIndexer {
    /// Encrypts the plaintext with the CLWW OPE scheme.
    /// Strings will return an [`IndexTerm::OpeVariable`].
    /// All other types will return an [`IndexTerm::OpeFixed`].
    pub fn encrypt(
        &self,
        value: &Plaintext,
        index_key: &IndexKey,
    ) -> Result<IndexTerm, EncryptionError> {
        let key = derive_ope_key(index_key);

        match value {
            Plaintext::Text(Some(s)) => {
                let ciphertext = s.as_str().encrypt_ope(&key)?;
                Ok(IndexTerm::OpeVariable(ciphertext.as_ref().to_vec()))
            }
            Plaintext::NaiveDate(Some(d)) => {
                let ciphertext = d.encrypt_ope(&key)?;
                Ok(IndexTerm::OpeFixed(ciphertext.as_ref().to_vec()))
            }
            Plaintext::Timestamp(Some(ts)) => {
                let ciphertext = ts.encrypt_ope(&key)?;
                Ok(IndexTerm::OpeFixed(ciphertext.as_ref().to_vec()))
            }
            Plaintext::Decimal(Some(d)) => {
                let ciphertext = d.encrypt_ope(&key)?;
                Ok(IndexTerm::OpeFixed(ciphertext.as_ref().to_vec()))
            }
            Plaintext::BigInt(Some(b)) => {
                let ciphertext = b.encrypt_ope(&key)?;
                Ok(IndexTerm::OpeFixed(ciphertext.as_ref().to_vec()))
            }
            Plaintext::Boolean(Some(b)) => {
                let ciphertext = b.encrypt_ope(&key)?;
                Ok(IndexTerm::OpeFixed(ciphertext.as_ref().to_vec()))
            }
            Plaintext::Float(Some(b)) => {
                let ciphertext = b.encrypt_ope(&key)?;
                Ok(IndexTerm::OpeFixed(ciphertext.as_ref().to_vec()))
            }
            Plaintext::Int(Some(i)) => {
                let ciphertext = i.encrypt_ope(&key)?;
                Ok(IndexTerm::OpeFixed(ciphertext.as_ref().to_vec()))
            }
            Plaintext::SmallInt(Some(i)) => {
                let ciphertext = i.encrypt_ope(&key)?;
                Ok(IndexTerm::OpeFixed(ciphertext.as_ref().to_vec()))
            }
            Plaintext::BigUInt(Some(u)) => {
                let ciphertext = u.encrypt_ope(&key)?;
                Ok(IndexTerm::OpeFixed(ciphertext.as_ref().to_vec()))
            }
            Plaintext::Json(_) => Err(EncryptionError::IndexingError(
                "OPE indexing not supported for JSON documents".into(),
            )),
            Plaintext::Text(None)
            | Plaintext::BigInt(None)
            | Plaintext::BigUInt(None)
            | Plaintext::Boolean(None)
            | Plaintext::Decimal(None)
            | Plaintext::Float(None)
            | Plaintext::Int(None)
            | Plaintext::NaiveDate(None)
            | Plaintext::SmallInt(None)
            | Plaintext::Timestamp(None) => Ok(IndexTerm::Null),
        }
    }
}

/// Derives the OPE-specific encryption key from a column's [`IndexKey`].
///
/// The column's `IndexKey` is also consumed by the ORE indexer (fed
/// straight into `OreCipher::init`). Without an extra derivation step
/// here, ORE and OPE would produce two PRF observations of the same
/// 32-byte secret, voiding the per-scheme key-independence assumption
/// in each scheme's security analysis. We BLAKE3-derive a fresh
/// OPE-only key to keep the two keystreams computationally
/// independent.
///
/// The info string `CIPHERSTASH-OPE-INDEX-V1` is **load-bearing for
/// the wire format**: every OPE ciphertext written from this code
/// path is tied to it. Changing the string invalidates every existing
/// OPE ciphertext, so any future change requires a versioned
/// migration.
///
/// ORE keeps consuming `IndexKey` raw because its wire format is
/// already shipped; this asymmetry is intentional. If/when ORE ever
/// migrates, it should adopt the same pattern with its own info
/// string.
fn derive_ope_key(index_key: &IndexKey) -> OpeKey {
    let mut hasher = blake3::Hasher::new();
    hasher.update(b"CIPHERSTASH-OPE-INDEX-V1");
    hasher.update(index_key.key());
    let mut derived = [0u8; 32];
    hasher.finalize_xof().fill(&mut derived);
    OpeKey::from(derived)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn key() -> IndexKey {
        IndexKey::from([7u8; 32])
    }

    #[test]
    fn encrypts_bigint_to_opefixed_with_65_bytes() {
        let ct = OpeIndexer
            .encrypt(&Plaintext::BigInt(Some(42)), &key())
            .unwrap();
        match ct {
            IndexTerm::OpeFixed(bytes) => assert_eq!(bytes.len(), 65),
            other => panic!("expected OpeFixed, got {other:?}"),
        }
    }

    #[test]
    fn encrypts_text_to_opevariable() {
        let ct = OpeIndexer
            .encrypt(&Plaintext::Text(Some("hello".into())), &key())
            .unwrap();
        assert!(matches!(ct, IndexTerm::OpeVariable(_)));
    }

    #[test]
    fn null_plaintext_yields_null_term() {
        let ct = OpeIndexer
            .encrypt(&Plaintext::BigInt(None), &key())
            .unwrap();
        assert_eq!(ct, IndexTerm::Null);
    }

    #[test]
    fn preserves_order_for_bigint() {
        let a = match OpeIndexer
            .encrypt(&Plaintext::BigInt(Some(10)), &key())
            .unwrap()
        {
            IndexTerm::OpeFixed(b) => b,
            _ => unreachable!(),
        };
        let b = match OpeIndexer
            .encrypt(&Plaintext::BigInt(Some(20)), &key())
            .unwrap()
        {
            IndexTerm::OpeFixed(b) => b,
            _ => unreachable!(),
        };
        assert!(a < b);
    }

    #[test]
    fn preserves_order_across_signed_range() {
        let neg = match OpeIndexer
            .encrypt(&Plaintext::BigInt(Some(-1)), &key())
            .unwrap()
        {
            IndexTerm::OpeFixed(b) => b,
            _ => unreachable!(),
        };
        let pos = match OpeIndexer
            .encrypt(&Plaintext::BigInt(Some(1)), &key())
            .unwrap()
        {
            IndexTerm::OpeFixed(b) => b,
            _ => unreachable!(),
        };
        assert!(neg < pos);
    }

    #[test]
    fn deterministic_for_same_plaintext_and_key() {
        let a = OpeIndexer
            .encrypt(&Plaintext::BigInt(Some(12345)), &key())
            .unwrap();
        let b = OpeIndexer
            .encrypt(&Plaintext::BigInt(Some(12345)), &key())
            .unwrap();
        assert_eq!(a, b);
    }

    #[test]
    fn try_from_rejects_non_ope_index_type() {
        let err = OpeIndexerOptions::try_from(&column::IndexType::Ore);
        assert!(err.is_err());
    }
}