cipherstash-client 0.42.0

mod pipeline;
mod query;

use std::borrow::Cow;

use super::{
    indexer::Indexes,
    json_indexer::{SteVec, SteVecPendingEncryption},
    EncryptionError, IndexTerm,
};
use crate::zerokms::{Decryptable, EncryptedRecord, IndexKey, RetrieveKeyPayload};
use uuid::Uuid;

#[cfg(feature = "tokio")]
use super::BytesWithDescriptor;

#[cfg(feature = "tokio")]
use crate::zerokms::DataKeyWithTag;

pub use pipeline::Encryptable;
pub use query::{QueryBuilder, Queryable};
use zerokms_protocol::Context;

pub struct StorageBuilder<'k, T> {
    source: T,
    key: &'k IndexKey,
    index_terms: Vec<IndexTerm>,
    descriptor: Option<String>,
    context: Cow<'k, [Context]>,
    keyset_id: Uuid,
    ste_vec: Option<SteVecPendingEncryption<16>>,
}

impl<'k, T> StorageBuilder<'k, T> {
    pub fn new(source: T, key: &'k IndexKey, keyset_id: Uuid) -> Self {
        Self {
            source,
            key,
            keyset_id,
            index_terms: Default::default(),
            descriptor: None,
            context: Default::default(),
            ste_vec: None,
        }
    }

    // TODO: Should this error if the same index is added twice?
    pub fn index_with<I>(self, index: I) -> Result<Self, EncryptionError>
    where
        for<'t> I: Indexes<'t, T>,
    {
        index.index(self)
    }

    pub fn with_descriptor(mut self, descriptor: impl Into<String>) -> Self {
        self.descriptor = Some(descriptor.into());
        self
    }

    /// Adds context to the record.
    pub fn with_context(mut self, context: Context) -> Self {
        self.context.to_mut().push(context);
        self
    }

    /// Sets the context for the record (replacing any existing context).
    pub fn set_context(mut self, context: Cow<'k, [Context]>) -> Self {
        self.context = context;
        self
    }

    pub(crate) fn plaintext(&self) -> &T {
        &self.source
    }

    pub(crate) fn index_key(&self) -> &IndexKey {
        self.key
    }

    pub(crate) fn set_ste_vec(&mut self, ste_vec: SteVecPendingEncryption<16>) {
        self.ste_vec = Some(ste_vec);
    }

    pub(crate) fn add_index_term(&mut self, index_term: IndexTerm) {
        self.index_terms.push(index_term);
    }

    pub fn keyset_id(&self) -> Uuid {
        self.keyset_id
    }

    /// Gets the descriptor. An empty string is returned if a descriptor has
    /// not been set.
    ///
    /// SteVec targets use the descriptor like every other target: the
    /// document's data key is generated under it, and the envelope's
    /// [`crate::zerokms::KeyHeader`] stores it once for decryption. (The old
    /// wire format hard-coded `""` for SteVec because the descriptor was
    /// repeated inside every entry and changing it would have broken stored
    /// data; the envelope format is a sanctioned wire break, closing that
    /// limitation.)
    pub fn descriptor(&self) -> &str {
        self.descriptor.as_deref().unwrap_or("")
    }

    // An SteVec cannot exist alongside other index types
    #[cfg(feature = "tokio")]
    pub(crate) fn build_for_encryption(self) -> PendingEncryption<'k>
    where
        BytesWithDescriptor: From<(T, String)>,
    {
        let encryptable = if let Some(ste_vec) = self.ste_vec {
            // The descriptor rides along so the SteVec envelope's KeyHeader
            // carries the same descriptor the data key was generated under.
            EncryptableInner::SteVec(ste_vec, self.descriptor.unwrap_or_default())
        } else {
            let bytes =
                BytesWithDescriptor::from((self.source, self.descriptor.unwrap_or_default()));
            EncryptableInner::Plaintext(bytes, self.index_terms, self.context)
        };

        PendingEncryption {
            encryptable,
            keyset_id: self.keyset_id,
        }
    }

    pub fn to_index_terms(self) -> Vec<IndexTerm> {
        self.index_terms
    }
}

#[cfg(feature = "tokio")]
enum EncryptableInner<'a> {
    Plaintext(BytesWithDescriptor, Vec<IndexTerm>, Cow<'a, [Context]>),
    /// A pending SteVec plus the column's descriptor (the KeyHeader stores
    /// it once for the whole document).
    SteVec(SteVecPendingEncryption<16>, String),
}

#[cfg(feature = "tokio")]
pub struct PendingEncryption<'a> {
    encryptable: EncryptableInner<'a>,
    keyset_id: Uuid,
}

#[cfg(feature = "tokio")]
impl<'a> PendingEncryption<'a> {
    // FIXME: We probably can change EncryptPayload to OWN the plaintext rather than borrow it
    pub fn encrypt(self, datakey_with_tag: DataKeyWithTag) -> Result<Encrypted, EncryptionError> {
        match self.encryptable {
            EncryptableInner::Plaintext(bytes, index_terms, context) => {
                use crate::zerokms::{self, EncryptPayload};

                let payload = EncryptPayload::from(&bytes).set_context(context.clone());
                let keyset_id = Some(self.keyset_id);
                zerokms::encrypt(payload, datakey_with_tag, keyset_id)
                    .map(|encrypted| Encrypted::Record(encrypted, index_terms))
                    .map_err(EncryptionError::from)
            }
            EncryptableInner::SteVec(ste_vec_pending_encryption, descriptor) => {
                ste_vec_pending_encryption
                    .encrypt(datakey_with_tag, &descriptor, Some(self.keyset_id))
                    .map(Encrypted::SteVec)
                    .map_err(EncryptionError::from)
            }
        }
    }
}

/// Represents an encrypted payload that is ready to be stored.
#[derive(Debug)]
pub enum Encrypted {
    SteVec(SteVec<16>),
    Record(EncryptedRecord, Vec<IndexTerm>),
}

impl Decryptable for Encrypted {
    type Error = EncryptionError;

    fn into_encrypted_record(self) -> Result<EncryptedRecord, Self::Error> {
        match self {
            Encrypted::Record(record, _) => Ok(record),
            // The root entry's decryptable form: header key material + root
            // ciphertext. Its selector-derived nonce is surfaced via
            // `nonce_override` (the pipelines capture it before this
            // conversion).
            Encrypted::SteVec(stevec) => {
                stevec.into_root_record_with_nonce().map(|root| root.record)
            }
        }
    }

    fn keyset_id(&self) -> Option<Uuid> {
        match self {
            Encrypted::Record(encrypted_record, ..) => encrypted_record.keyset_id,
            Encrypted::SteVec(stevec) => stevec.header().keyset_id,
        }
    }

    fn retrieve_key_payload(&self) -> Result<RetrieveKeyPayload<'_>, Self::Error> {
        match self {
            Encrypted::Record(encrypted_record, ..) => encrypted_record
                .retrieve_key_payload()
                .map_err(|_| EncryptionError::Infallible),
            Encrypted::SteVec(stevec) => Ok(stevec.header().retrieve_key_payload()),
        }
    }

    /// The decrypt path reads this to learn the AEAD nonce for the root
    /// ciphertext. A classic [`EncryptedRecord`] derives its nonce from the
    /// data-key IV, so it returns `None` and the default IV-derived nonce is
    /// used. A SteVec document has no self-contained root record: its root
    /// ciphertext is the first `sv` entry, whose nonce is derived from that
    /// entry's stored selector, not the IV ([`SteVec::root_nonce`]). Returning
    /// it here lets the shared decrypt path reassemble the root exactly as it
    /// was sealed (`Err` — an empty document with no root entry — degrades to
    /// `None`, and decryption then fails on the missing ciphertext, not here).
    fn nonce_override(&self) -> Option<[u8; 12]> {
        match self {
            Encrypted::Record(..) => None,
            Encrypted::SteVec(stevec) => stevec.root_nonce().ok(),
        }
    }

    /// Companion to [`Self::nonce_override`]: the root ciphertext is bound to
    /// the root entry's full selector in the AAD, so decryption must supply the
    /// same 16 bytes. Classic records add nothing (`None`).
    fn aad_selector(&self) -> Option<[u8; 16]> {
        match self {
            Encrypted::Record(..) => None,
            Encrypted::SteVec(stevec) => stevec.root_selector().ok(),
        }
    }
}

#[cfg(all(test, feature = "tokio"))]
mod tests {
    use super::*;
    use crate::encryption::{JsonIndexer, MatchIndexer, Plaintext, UniqueIndexer};
    use crate::zerokms::{self, DataKeyWithTag, IndexKey};
    use serde_json::json;

    impl EncryptableInner<'_> {
        fn bytes_with_descriptor(&self) -> Option<&BytesWithDescriptor> {
            match self {
                EncryptableInner::Plaintext(bytes, ..) => Some(bytes),
                _ => None,
            }
        }

        fn index_terms(&self) -> Option<&[IndexTerm]> {
            match self {
                EncryptableInner::Plaintext(_, terms, ..) => Some(terms),
                _ => None,
            }
        }
    }

    #[test]
    fn test_single_standard_index() -> Result<(), EncryptionError> {
        let plaintext = Plaintext::from("hello world");
        let index_key = IndexKey::from([0; 32]);
        let keyset_id = Uuid::new_v4();

        let for_encryption = StorageBuilder::new(plaintext, &index_key, keyset_id)
            .index_with(UniqueIndexer::default())?
            .with_descriptor("descriptor")
            .build_for_encryption();

        assert!(matches!(
            for_encryption.encryptable,
            EncryptableInner::Plaintext(..)
        ));
        assert_eq!(for_encryption.encryptable.index_terms().unwrap().len(), 1);
        assert!(for_encryption.encryptable.bytes_with_descriptor().is_some());
        assert_eq!(for_encryption.keyset_id, keyset_id);

        Ok(())
    }

    #[test]
    fn test_multiple_standard_indexes() -> Result<(), EncryptionError> {
        let plaintext = Plaintext::from("hello world");
        let index_key = IndexKey::from([0; 32]);
        let keyset_id = Uuid::new_v4();

        let for_encryption = StorageBuilder::new(plaintext, &index_key, keyset_id)
            .index_with(UniqueIndexer::default())?
            .index_with(MatchIndexer::default())?
            .with_descriptor("descriptor")
            // TODO: Add ORE here, too
            .build_for_encryption();

        assert!(matches!(
            for_encryption.encryptable,
            EncryptableInner::Plaintext(..)
        ));
        assert_eq!(for_encryption.encryptable.index_terms().unwrap().len(), 2);
        assert_eq!(for_encryption.keyset_id, keyset_id);

        Ok(())
    }

    #[test]
    fn test_ste_vec_index() -> Result<(), EncryptionError> {
        let plaintext = Plaintext::from(json!({"hello": {}}));
        let index_key = IndexKey::from([0; 32]);
        let keyset_id = Uuid::new_v4();

        let for_encryption = StorageBuilder::new(plaintext, &index_key, keyset_id)
            .index_with(JsonIndexer::default())?
            .with_descriptor("descriptor")
            .build_for_encryption();

        assert!(matches!(
            for_encryption.encryptable,
            EncryptableInner::SteVec(_, _)
        ));
        assert_eq!(for_encryption.encryptable.index_terms(), None);
        assert!(for_encryption.encryptable.bytes_with_descriptor().is_none());
        assert_eq!(for_encryption.keyset_id, keyset_id);

        Ok(())
    }

    #[test]
    fn test_standard_encrypt() -> Result<(), EncryptionError> {
        let plaintext = Plaintext::from("hello world");
        let index_key = IndexKey::from([0; 32]);
        let datakey_with_tag = DataKeyWithTag::default();
        let keyset_id = Uuid::new_v4();

        let encrypted = StorageBuilder::new(plaintext, &index_key, keyset_id)
            .index_with(UniqueIndexer::default())?
            .with_descriptor("descriptor")
            .build_for_encryption()
            .encrypt(datakey_with_tag)?;

        assert!(matches!(encrypted, Encrypted::Record(_, _)));

        let result = zerokms::decrypt(encrypted, DataKeyWithTag::default().key).unwrap();
        let result = Plaintext::from_slice(&result)?;
        assert_eq!(result, *"hello world");

        Ok(())
    }

    #[test]
    fn test_standard_encrypt_with_context() -> Result<(), EncryptionError> {
        let plaintext = Plaintext::from("hello world");
        let index_key = IndexKey::from([0; 32]);
        let datakey_with_tag = DataKeyWithTag::default();
        let keyset_id = Uuid::new_v4();

        let encrypted = StorageBuilder::new(plaintext, &index_key, keyset_id)
            .index_with(UniqueIndexer::default())?
            .with_descriptor("descriptor")
            .with_context(Context::new_tag("sensitive"))
            .with_context(Context::new_identity_claim("sub"))
            .build_for_encryption()
            .encrypt(datakey_with_tag)?;

        assert!(matches!(encrypted, Encrypted::Record(_, _)));

        Ok(())
    }

    #[test]
    fn test_ste_vec_encrypt() -> Result<(), EncryptionError> {
        let plaintext = Plaintext::from(json!({"hello": {}}));
        let index_key = IndexKey::from([0; 32]);
        let datakey_with_tag = DataKeyWithTag::default();
        let keyset_id = Uuid::new_v4();

        let encrypted = StorageBuilder::new(plaintext, &index_key, keyset_id)
            .index_with(JsonIndexer::default())?
            .with_descriptor("descriptor")
            .build_for_encryption()
            .encrypt(datakey_with_tag)?;

        assert!(matches!(encrypted, Encrypted::SteVec(_)));

        Ok(())
    }
}