cipherstash-client 0.34.1-alpha.1

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 or if the underlying data is
    /// an [`SteVec`].
    ///
    /// Encrypted::SteVec(_) cannot currently have a descriptor set on it, so if one were to be encrypted using a
    /// descriptor then decryption would use an empty string which would cause decryption to fail.
    ///
    /// This is a limitation of the implementation: an `SteVec` *should* have a descriptor but changing it without an
    /// EQL upgrade with a mandatory migration process would break existing encrypted data.
    pub fn descriptor(&self) -> &str {
        if self.ste_vec.is_none() {
            self.descriptor.as_deref().unwrap_or("")
        } else {
            ""
        }
    }

    // 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 {
            EncryptableInner::SteVec(ste_vec)
        } 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]>),
    SteVec(SteVecPendingEncryption<16>), // TODO: This should also have a descriptor on it
}

#[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) => ste_vec_pending_encryption
                .encrypt(datakey_with_tag)
                .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),
            Encrypted::SteVec(stevec) => stevec.into_root_ciphertext(),
        }
    }

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

    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) => stevec.root_ciphertext().and_then(|root| {
                root.retrieve_key_payload()
                    .map_err(|_| EncryptionError::Infallible)
            }),
        }
    }
}

#[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(())
    }
}