cipherstash-client 0.12.5

//! Module for CipherStash encryption schemes and indexers
pub mod compound_indexer;
mod errors;
mod json_indexer;
mod match_indexer;
mod ore_indexer;
mod plaintext;
mod text;
mod unique_indexer;

use self::compound_indexer::{
    accumulator::Accumulator, composable_plaintext::ComposablePlaintext, ComposableIndex,
    CompoundIndex,
};
use crate::{
    credentials::{
        service_credentials::{ServiceCredentials, ServiceToken},
        Credentials,
    },
    zerokms::{EncryptPayload, EncryptedRecord, ZeroKMSWithClientKey},
};
use zerokms_protocol::cipherstash_config::{column::IndexType, operator::Operator, ColumnType};

// Re-exports
pub use self::{
    errors::{EncryptionError, TypeParseError},
    json_indexer::containment::JsonContainmentIndexer,
    match_indexer::MatchIndexer,
    ore_indexer::OreIndexer,
    plaintext::{
        BytesWithDescriptor, Plaintext, PlaintextNullVariant, PlaintextTarget, TryFromPlaintext,
    },
    unique_indexer::UniqueIndexer,
};

pub struct Encryption<C: Credentials<Token = ServiceToken> = ServiceCredentials> {
    // This field is public in order for the Driver to be able to cache
    // configuration and avoid a round trip to Vitur for every database
    // connection.
    pub root_key: [u8; 32],
    client: ZeroKMSWithClientKey<C>,
}

impl<Creds: Credentials<Token = ServiceToken>> Encryption<Creds> {
    pub fn new(root_key: [u8; 32], client: ZeroKMSWithClientKey<Creds>) -> Self {
        Self { root_key, client }
    }

    pub async fn encrypt<T: Into<BytesWithDescriptor>>(
        &self,
        items: impl IntoIterator<Item = T>,
    ) -> Result<Vec<EncryptedRecord>, EncryptionError> {
        let timer = cipherstash_stats::ENCRYPTION_DURATION.start_timer();

        let result = self.encrypt_impl(items).await;

        match result {
            Ok(ref output) => {
                timer.stop_and_record();
                cipherstash_stats::ENCRYPTIONS.inc_by(output.len() as u64);
            }
            Err(_) => {
                cipherstash_stats::ENCRYPTION_ERRORS.inc();
            }
        };

        result
    }

    #[inline(always)]
    pub async fn encrypt_impl<T: Into<BytesWithDescriptor>>(
        &self,
        items: impl IntoIterator<Item = T>,
    ) -> Result<Vec<EncryptedRecord>, EncryptionError> {
        let payloads: Vec<BytesWithDescriptor> = items.into_iter().map(Into::into).collect();

        Ok(self
            .client
            .encrypt(payloads.iter().map(EncryptPayload::from), None)
            .await?)
    }

    pub async fn encrypt_single(
        &self,
        target: PlaintextTarget,
    ) -> Result<EncryptedRecord, EncryptionError> {
        let timer = cipherstash_stats::ENCRYPTION_DURATION.start_timer();

        let result = self.encrypt_single_impl(target).await;

        match result {
            Ok(_) => {
                timer.stop_and_record();
                cipherstash_stats::ENCRYPTIONS.inc();
            }
            Err(_) => {
                cipherstash_stats::ENCRYPTION_ERRORS.inc();
            }
        };
        result
    }

    #[inline(always)]
    pub async fn encrypt_single_impl(
        &self,
        target: PlaintextTarget,
    ) -> Result<EncryptedRecord, EncryptionError> {
        let payload = target.payload();

        let ciphertext = self
            .client
            .encrypt_single(EncryptPayload::from(&payload), None)
            .await?;

        Ok(ciphertext)
    }

    pub async fn decrypt_single(
        &self,
        ciphertext: EncryptedRecord,
    ) -> Result<Plaintext, EncryptionError> {
        let timer = cipherstash_stats::DECRYPTION_DURATION.start_timer();

        let result = self.decrypt_single_impl(ciphertext).await;
        match result {
            Ok(_) => {
                timer.stop_and_record();
                cipherstash_stats::DECRYPTIONS.inc();
            }
            Err(_) => {
                cipherstash_stats::DECRYPTION_ERRORS.inc();
            }
        };
        result
    }

    #[inline(always)]
    pub async fn decrypt_single_impl(
        &self,
        ciphertext: EncryptedRecord,
    ) -> Result<Plaintext, EncryptionError> {
        let decrypted = self.client.decrypt_single(ciphertext).await?;
        Ok(Plaintext::from_slice(&decrypted)?)
    }

    /// Like `decrypt` but doesn't expect all values to be decryptable.
    /// This only means that a given input is `None` or the slice is not a
    /// serialized [`EncryptedRecord`].
    /// In the future this could also cover cases
    /// where the caller is not _authorized_ to decrypt a given value.
    ///
    /// As it stands, this function will return an Error if any valid ciphertexts
    /// fail to decrypt.
    ///
    /// Items in the returned vec wil be in the same order as the input
    /// but any values that are unable to be decrypted will be returned as `None`.
    ///
    /// Encrypted records must be encoded with hex.
    pub async fn maybe_decrypt_hex<C>(
        &self,
        ciphertexts: impl IntoIterator<Item = Option<C>>,
    ) -> Result<Vec<Option<Plaintext>>, EncryptionError>
    where
        C: AsRef<[u8]>,
    {
        let timer = cipherstash_stats::DECRYPTION_DURATION.start_timer();

        let result = self.maybe_decrypt_hex_impl(ciphertexts).await;

        match result {
            Ok(ref output) => {
                timer.stop_and_record();
                cipherstash_stats::DECRYPTIONS.inc_by(output.len() as u64);
            }
            Err(_) => {
                cipherstash_stats::DECRYPTION_ERRORS.inc();
            }
        };

        result
    }

    #[inline(always)]
    pub async fn maybe_decrypt_hex_impl<I, C>(
        &self,
        ciphertexts: I,
    ) -> Result<Vec<Option<Plaintext>>, EncryptionError>
    where
        I: IntoIterator<Item = Option<C>>,
        C: AsRef<[u8]>,
    {
        let records: (Vec<bool>, Vec<EncryptedRecord>) =
            ciphertexts
                .into_iter()
                .fold(Default::default(), |(mut all, mut target), hex_str| {
                    if let Some(rec) = hex_str
                        .map(hex::decode)
                        .transpose()
                        .unwrap_or(None)
                        .and_then(|bytes| EncryptedRecord::from_slice(&bytes).ok())
                    {
                        target.push(rec);
                        all.push(true);
                    } else {
                        all.push(false);
                    }
                    (all, target)
                });

        let mut results = self
            .client
            .decrypt(records.1)
            .await?
            .into_iter()
            .map(|bytes| Plaintext::from_slice(&bytes));

        Ok(records
            .0
            .iter()
            .map(|valid| {
                if *valid {
                    results.next().transpose()
                } else {
                    Ok(None)
                }
            })
            .collect::<Result<Vec<Option<Plaintext>>, _>>()?)
    }

    pub async fn decrypt(
        &self,
        ciphertexts: impl IntoIterator<Item = EncryptedRecord>,
    ) -> Result<Vec<Plaintext>, EncryptionError> {
        let timer = cipherstash_stats::DECRYPTION_DURATION.start_timer();

        let result = self.decrypt_impl(ciphertexts).await;

        match result {
            Ok(ref output) => {
                timer.stop_and_record();
                cipherstash_stats::DECRYPTIONS.inc_by(output.len() as u64);
            }
            Err(_) => {
                cipherstash_stats::DECRYPTION_ERRORS.inc();
            }
        };
        result
    }

    #[inline(always)]
    pub async fn decrypt_impl(
        &self,
        ciphertexts: impl IntoIterator<Item = EncryptedRecord>,
    ) -> Result<Vec<Plaintext>, EncryptionError> {
        Ok(self
            .client
            .decrypt(ciphertexts)
            .await?
            .iter()
            .map(|bytes| Plaintext::from_slice(bytes))
            .collect::<Result<Vec<Plaintext>, _>>()?)
    }

    pub fn index(
        &self,
        value: &Plaintext,
        index_type: &IndexType,
    ) -> Result<IndexTerm, EncryptionError> {
        match index_type {
            IndexType::Ore => OreIndexer::new(self.root_key)?.encrypt(value),
            IndexType::Unique { token_filters } => {
                UniqueIndexer::new(self.root_key, token_filters.clone()).encrypt(value)
            }
            IndexType::Match {
                tokenizer,
                token_filters,
                k,
                m,
                ..
            } => MatchIndexer::new(
                self.root_key,
                tokenizer.clone(),
                token_filters.to_vec(),
                *k,
                *m,
            )
            .encrypt(value),
            IndexType::SteVec { prefix } => {
                JsonContainmentIndexer::new(self.root_key, prefix.clone()).encrypt(
                    value.into_json().ok_or(EncryptionError::IndexingError(
                        "expected JSONB plaintext".into(),
                    ))?,
                )
            }
        }
    }

    pub fn index_all(&self, target: &PlaintextTarget) -> Result<Vec<IndexTerm>, EncryptionError> {
        let mut indexes = vec![];

        for index in target.config().indexes.iter() {
            indexes.push(self.index(&target.plaintext, &index.index_type)?);
        }

        Ok(indexes)
    }

    pub fn compound_index(
        &self,
        index: &CompoundIndex<impl ComposableIndex + Send>,
        input: impl Into<ComposablePlaintext>,
        salt: Option<impl AsRef<[u8]>>,
        term_length: usize,
    ) -> Result<IndexTerm, EncryptionError> {
        let accumulator = salt
            .map(|s| Accumulator::from_salt(s.as_ref()))
            .unwrap_or_else(Accumulator::empty);

        let term = index
            .compose_index(self.root_key, input.into(), accumulator)?
            .truncate(term_length)?;

        Ok(term.into())
    }

    pub fn compound_query(
        &self,
        index: &CompoundIndex<impl ComposableIndex + Send>,
        input: impl Into<ComposablePlaintext>,
        salt: Option<impl AsRef<[u8]>>,
        term_length: usize,
    ) -> Result<IndexTerm, EncryptionError> {
        let accumulator = salt
            .map(|s| Accumulator::from_salt(s.as_ref()))
            .unwrap_or_else(Accumulator::empty);

        let term = index
            .compose_query(self.root_key, input.into(), accumulator)?
            .exactly_one()?
            .truncate(term_length)?;

        Ok(term.try_into()?)
    }

    pub fn index_for_operator(
        &self,
        value: &Plaintext,
        index_type: &IndexType,
        operator: &Operator,
        cast_type: &ColumnType,
    ) -> Result<IndexTerm, EncryptionError> {
        // Check if index supports op
        if !index_type.supports(operator, cast_type) {
            return Err(EncryptionError::IndexingError(format!(
                "Unsupported operator ({}) for Index {:?}",
                operator.as_str(),
                index_type
            )));
        }
        match index_type {
            IndexType::Ore => OreIndexer::new(self.root_key)?.encrypt_for_query(value),
            // Unique and Match don't work any differently for queries
            IndexType::Unique { .. } => self.index(value, index_type),
            IndexType::Match { .. } => self.index(value, index_type),
            IndexType::SteVec { .. } => self.index(value, index_type),
        }
    }
}

#[derive(Debug, Eq, Clone, PartialEq)]
pub enum IndexTerm {
    Binary(Vec<u8>),
    BinaryVec(Vec<Vec<u8>>),
    BitMap(Vec<u16>),
    /// Represents a full ORE Ciphertext (both left and right)
    OreFull(Vec<u8>),
    /// Array of FullOre terms
    OreArray(Vec<Vec<u8>>),
    /// Represents a Left ORE Ciphertext
    OreLeft(Vec<u8>),
    /// NULL index field
    Null,
}

impl IndexTerm {
    pub fn as_binary(self) -> Option<Vec<u8>> {
        if let Self::Binary(x) = self {
            Some(x)
        } else {
            None
        }
    }

    /// Get the index term as a vector of binary terms.
    /// If the term is a single binary term, it will be wrapped in a vec.
    pub fn as_binary_vec(self) -> Option<Vec<Vec<u8>>> {
        match self {
            Self::BinaryVec(x) => Some(x),
            Self::Binary(x) => Some(vec![x]),
            _ => None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{config::zero_kms_config::ZeroKMSConfig, zerokms::ZeroKMS};
    use zerokms_protocol::cipherstash_config::ColumnConfig;

    fn create_test_encryption() -> Encryption {
        let root_key = [0; 32];

        let config = ZeroKMSConfig::builder()
            .with_env()
            .build_with_client_key()
            .expect("Unable to load Vitur config");

        let vitur_client = ZeroKMS::new_with_client_key(
            &config.base_url(),
            config.credentials(),
            config.decryption_log_path().as_deref(),
            config.client_key(),
        );

        Encryption::new(root_key, vitur_client)
    }

    // Ignore for now because these require a real Vitur instance running
    #[ignore]
    #[tokio::test]
    async fn test_round_trip_single() -> Result<(), Box<dyn std::error::Error>> {
        let encryption = create_test_encryption();
        let value = "hello cipher";
        let target = PlaintextTarget::new(value, ColumnConfig::build("name"), None);

        let ciphertext = encryption.encrypt_single(target).await?;
        assert_eq!(
            Plaintext::new(value),
            encryption.decrypt_single(ciphertext).await?
        );

        Ok(())
    }

    // Ignore for now because these require a real Vitur instance running
    /*#[ignore]
    #[tokio::test]
    async fn test_round_trip_bulk_decrypt() -> Result<(), Box<dyn std::error::Error>> {
        let encryption = create_test_encryption();

        let plaintexts = vec!["a".into(), "b".into(), "c".into()];

        let mut ciphertexts = vec![];

        for (i, plaintext) in plaintexts.iter().enumerate() {
            ciphertexts.push(
                encryption
                    .encrypt_single(plaintext, &format!("value-{i}"))
                    .await?,
            );
        }

        assert_eq!(plaintexts, encryption.decrypt(ciphertexts).await?);

        Ok(())
    }

    // Ignore for now because these require a real Vitur instance running
    #[ignore]
    #[tokio::test]
    async fn test_round_trip_bulk_maybe_decrypt() -> Result<(), Box<dyn std::error::Error>> {
        let encryption = create_test_encryption();

        let p1 = "a".into();
        let p2 = "b".into();
        let p3 = "c".into();

        let mut ciphertexts = vec![];

        for (i, plaintext) in vec![&p1, &p2, &p3].into_iter().enumerate() {
            ciphertexts.push(Some(hex::encode(
                encryption
                    .encrypt_single(plaintext, &format!("value-{i}"))
                    .await?
                    .to_vec()
                    .unwrap(),
            )));

            ciphertexts.push(Some(format!("not-encrypted-{i}")));
        }

        ciphertexts.push(None);

        assert_eq!(
            vec![Some(p1), None, Some(p2), None, Some(p3), None, None],
            encryption.maybe_decrypt_hex(ciphertexts).await?
        );

        Ok(())
    }*/

    // TODO: Test the other functions and indexers
}