cipherstash-client 0.34.1-alpha.1

use self::conversion::IntoOrePlaintext;
use crate::zerokms::IndexKey;
use cipherstash_core::string::orderise_string;
use ore_rs::{scheme::bit2::OreAes128ChaCha20, OreCipher, OreOutput};
use zerokms_protocol::cipherstash_config::column;

use super::{
    errors::EncryptionError,
    indexer::{IndexerInit, Indexes, IndexesForQuery, QueryOp},
    IndexTerm, Plaintext, QueryBuilder, StorageBuilder,
};

pub struct OreIndexer;

pub mod conversion;

/// No options are currently supported
#[derive(Debug, Default)]
pub struct OreIndexerOptions;

impl IndexerInit for OreIndexer {
    type Args = OreIndexerOptions;
    type Error = EncryptionError;

    fn try_init<A>(_opts: A) -> Result<Self, Self::Error>
    where
        Self::Args: TryFrom<A, Error = Self::Error>,
    {
        Ok(Self)
    }
}

impl<'k> Indexes<'k, Plaintext> for OreIndexer {
    fn index(
        &self,
        mut builder: StorageBuilder<'k, Plaintext>,
    ) -> Result<StorageBuilder<'k, Plaintext>, EncryptionError> {
        let index_term = self.encrypt(builder.plaintext(), builder.index_key())?;
        builder.add_index_term(index_term);

        Ok(builder)
    }
}

impl<C> IndexesForQuery<Plaintext, C> for OreIndexer {
    fn query_index(
        &self,
        builder: QueryBuilder<Plaintext, C>,
        _op: QueryOp,
    ) -> Result<IndexTerm, EncryptionError> {
        let index_term = self.encrypt_for_query(builder.plaintext(), builder.index_key())?;
        Ok(index_term)
    }
}

impl Default for OreIndexer {
    fn default() -> Self {
        Self
    }
}

impl TryFrom<&column::IndexType> for OreIndexerOptions {
    type Error = EncryptionError;

    fn try_from(value: &column::IndexType) -> Result<Self, Self::Error> {
        match value {
            column::IndexType::Ore => Ok(Default::default()),
            _ => Err(EncryptionError::IndexingError(
                "MatchIndexerOptions can only be created from a Match index configuration"
                    .to_string(),
            )),
        }
    }
}

impl OreIndexer {
    /// Encrypts the plaintext with an appropriate ORE scheme.
    /// Strings will return an [`IndexTerm::OreArray`].
    /// All other types will return a [`IndexTerm::OreFull`].
    ///
    pub fn encrypt(
        &self,
        value: &Plaintext,
        index_key: &IndexKey,
    ) -> Result<IndexTerm, EncryptionError> {
        let mut k1: [u8; 16] = Default::default();
        let mut k2: [u8; 16] = Default::default();
        k1.copy_from_slice(&index_key.key()[0..16]);
        k2.copy_from_slice(&index_key.key()[16..]);
        let cipher: OreAes128ChaCha20 = OreCipher::init(&k1, &k2)?;

        match value {
            Plaintext::Utf8Str(Some(s)) => self.encrypt_string(s, &cipher),
            Plaintext::Utf8Str(None)
            | Plaintext::BigInt(None)
            | Plaintext::Boolean(None)
            | Plaintext::Decimal(None)
            | Plaintext::Float(None)
            | Plaintext::Int(None)
            | Plaintext::NaiveDate(None)
            | Plaintext::SmallInt(None)
            | Plaintext::Timestamp(None) => Ok(IndexTerm::Null),
            other => {
                let ciphertext = other.to_ore().encrypt(&cipher)?;
                Ok(IndexTerm::OreFull(ciphertext.to_bytes()))
            }
        }
    }

    pub fn encrypt_for_query(
        &self,
        value: &Plaintext,
        index_key: &IndexKey,
    ) -> Result<IndexTerm, EncryptionError> {
        // TODO: Use only the encrypt_left - currently not supported
        // by the pl/pgsql function
        let cipher = get_cipher(index_key)?;

        match value {
            Plaintext::Utf8Str(Some(s)) => self.encrypt_string(s, &cipher),
            Plaintext::Utf8Str(None)
            | Plaintext::BigInt(None)
            | Plaintext::Boolean(None)
            | Plaintext::Decimal(None)
            | Plaintext::Float(None)
            | Plaintext::Int(None)
            | Plaintext::NaiveDate(None)
            | Plaintext::SmallInt(None)
            | Plaintext::Timestamp(None) => Ok(IndexTerm::Null),
            other => {
                let ciphertext = other.to_ore().encrypt(&cipher)?;
                Ok(IndexTerm::OreFull(ciphertext.to_bytes()))
            }
        }
    }

    fn encrypt_string(
        &self,
        input_str: &str,
        cipher: &OreAes128ChaCha20,
    ) -> Result<IndexTerm, EncryptionError> {
        use ore_rs::OreEncrypt;

        let ciphertexts = orderise_string(input_str)?
            .into_iter()
            .map(|value| value.encrypt(cipher).map(|ct| ct.to_bytes()))
            .collect::<Result<Vec<Vec<u8>>, _>>()?;

        Ok(IndexTerm::OreArray(ciphertexts))
    }
}

fn get_cipher(index_key: &IndexKey) -> Result<OreAes128ChaCha20, EncryptionError> {
    let mut k1: [u8; 16] = Default::default();
    let mut k2: [u8; 16] = Default::default();
    k1.copy_from_slice(&index_key.key()[0..16]);
    k2.copy_from_slice(&index_key.key()[16..]);
    OreCipher::init(&k1, &k2).map_err(EncryptionError::from)
}