cipherstash-client 0.34.1-alpha.1

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

use crate::zerokms::{EncryptPayload, EncryptedRecord, IndexKey, ZeroKMSWithClientKey};
use compound_indexer::{
    accumulator::Accumulator, composable_plaintext::ComposablePlaintext, ComposableIndex,
    CompoundIndex,
};
use serde::Serialize;
use stack_auth::AuthStrategy;
use zerokms_protocol::cipherstash_config::{column::IndexType, operator::Operator, ColumnType};

// Re-exports
pub use builder::{Encryptable, Encrypted, QueryBuilder, Queryable, StorageBuilder};
pub use errors::{EncryptionError, TypeParseError};
pub use indexer::IndexerInit;
pub use indexer::QueryOp;
pub use json_indexer::{
    EncryptedEntry, EncryptedSteVecTerm, JsonIndexer, JsonIndexerOptions, SteQueryVec, SteVec,
    TokenizedSelector,
};
pub use jsonb_indexer::*;
pub use match_indexer::MatchIndexer;
pub use ore_indexer::OreIndexer;
pub use plaintext::{
    BytesWithDescriptor, Plaintext, PlaintextNullVariant, PlaintextTarget, TryFromPlaintext,
};
pub use scoped_cipher::{DecryptOptions, ScopedCipher};
pub use unique_indexer::UniqueIndexer;

pub struct Encryption<C = stack_auth::AutoStrategy> {
    index_key: IndexKey,
    client: ZeroKMSWithClientKey<C>,
}

impl<Creds> Encryption<Creds>
where
    Creds: Send + Sync + 'static,
    for<'a> &'a Creds: AuthStrategy,
{
    pub fn new(root_key: IndexKey, client: ZeroKMSWithClientKey<Creds>) -> Self {
        Self {
            index_key: root_key,
            client,
        }
    }

    pub async fn encrypt<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 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 decrypted = self
            .client
            .decrypt_single(ciphertext, None, None, None)
            .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<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_cbor_bytes(&bytes).ok())
                    {
                        target.push(rec);
                        all.push(true);
                    } else {
                        all.push(false);
                    }
                    (all, target)
                });

        let decrypted = self.client.decrypt(records.1, None, None, None).await?;
        let mut results = decrypted
            .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 decrypted = self.client.decrypt(ciphertexts, None, None, None).await?;
        Ok(decrypted
            .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::try_init(index_type)?.encrypt(value, &self.index_key),
            IndexType::Unique { .. } => {
                UniqueIndexer::try_init(index_type)?.encrypt(value, &self.index_key)
            }
            IndexType::Match { .. } => {
                MatchIndexer::try_init(index_type)?.encrypt(value, &self.index_key)
            }
            IndexType::SteVec { .. } => Err(EncryptionError::IndexingError(
                "SteVec not supported via direct Encryption API: use Pipeline or Builder instead"
                    .to_string(),
            )),
        }
    }

    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.index_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.index_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.encrypt_for_query(value, &self.index_key),
            // 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, PartialEq, Serialize)]
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>),
    SteVecSelector(TokenizedSelector<16>),
    SteVecTerm(EncryptedSteVecTerm),
    SteQueryVec(SteQueryVec<16>),
    /// 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,
        }
    }
}