cipherstash-client 0.34.1-alpha.1

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::{BytesWithDescriptor, Plaintext},
    zerokms::{self, encrypt, DataKeyWithTag, EncryptedRecord},
};
use itertools::Itertools;
use serde::{Deserialize, Serialize};

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

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

/// Describes a SteVec, which is a vector of encrypted entries used for index JSON.
/// `N` is the length of the selector.
#[derive(Debug, Serialize, Deserialize)]
pub struct SteVec<const N: usize>(pub(super) Vec<EncryptedEntry<N>>);

impl<const N: usize> SteVec<N> {
    /// Consumes the SteVec and returns the root ciphertext.
    /// If the SteVec is empty, an error is returned.
    pub fn into_root_ciphertext(self) -> Result<EncryptedRecord, EncryptionError> {
        // The root record will always be the first record in the vector
        self.0
            .into_iter()
            .next()
            .map(|enc| enc.record)
            .ok_or_else(|| EncryptionError::InvalidValue("SteVec missing root record".to_string()))
    }

    /// Returns a reference to the root ciphertext.
    /// If the SteVec is empty, an error is returned.
    pub fn root_ciphertext(&self) -> Result<&EncryptedRecord, EncryptionError> {
        // The root record will always be the first record in the vector
        self.0
            .first()
            .map(|enc| &enc.record)
            .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.0.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<const N: usize> SteVecPendingEncryption<N> {
    pub fn encrypt(self, key: DataKeyWithTag) -> Result<SteVec<N>, zerokms::Error> {
        self.0
            .into_iter()
            .map(|entry| entry.build_final(key.clone()))
            .try_collect()
            .map(SteVec)
    }

    pub fn into_query(self) -> SteQueryVec<N> {
        SteQueryVec::from(self)
    }
}

struct EntryWithEncryptedTerm<const N: usize> {
    tokenized_selector: TokenizedSelector<N>,
    term: EncryptedSteVecTerm,
    plaintext: Plaintext,
    parent_is_array: bool,
}

impl<const N: usize> EntryWithEncryptedTerm<N> {
    fn build_final(self, key: DataKeyWithTag) -> Result<EncryptedEntry<N>, zerokms::Error> {
        // TODO: The descriptor is currently set to an empty string - it should be the context (i.e. table and column name)
        let bwd = BytesWithDescriptor::from((&self.plaintext, ""));
        // FIXME: If we move keyset_id to DataKeyWithTag, this becomes a lot cleaner
        let ciphertext = encrypt(&bwd, key, None)?;
        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.
pub struct QueryEntry<const N: usize>(
    pub(super) TokenizedSelector<N>,
    pub(super) 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(test)]
mod tests {
    use crate::encryption::json_indexer::{
        prefix_mac::{PrefixMac, UpdatePrefixMac},
        ste_vec::{
            encrypted_term::{EncryptedSteVecTerm, Mac},
            EncryptedEntry, TokenizedSelector,
        },
    };
    use recipher::key::Iv;
    use std::fmt::Debug;
    use uuid::Uuid;

    #[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] {
            self.finalized_terms = self.pending_terms.drain(..).collect();
            [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));
        }
    }

    #[test]
    fn test_serde_encrypted_entry() {
        use crate::zerokms::EncryptedRecord;

        let ciphertext = EncryptedRecord {
            iv: Iv::default(),
            ciphertext: vec![1; 32],
            tag: vec![1; 16],
            descriptor: "users/name".to_string(),
            keyset_id: Some(Uuid::new_v4()),
        };

        let entry = EncryptedEntry::new(
            TokenizedSelector([1; 16]),
            EncryptedSteVecTerm::Mac(Mac::new(vec![2; 16])),
            ciphertext,
            false,
        );

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