cipherstash-client 0.41.1

use super::builder::StorageBuilder;
use super::indexer::{IndexerInit, Indexes, IndexesForQuery, QueryOp};
use super::text::{char_filter_prefix_and_suffix, TokenFilter, Tokenizer};
use super::QueryBuilder;
use super::{errors::EncryptionError, plaintext::Plaintext, IndexTerm};
use crate::zerokms::IndexKey;
use cipherstash_config::column;
use cipherstash_core::bloom_filter::{BloomFilter, BloomFilterOps};

impl IndexerInit for MatchIndexer {
    type Args = MatchIndexerOptions;
    type Error = EncryptionError;

    fn try_init<A>(opts: A) -> Result<Self, Self::Error>
    where
        Self::Args: TryFrom<A, Error = Self::Error>,
    {
        let opts = MatchIndexerOptions::try_from(opts)?;
        Ok(Self::new(opts))
    }
}

impl<'k> Indexes<'k, Plaintext> for MatchIndexer {
    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 MatchIndexer {
    fn query_index(
        &self,
        builder: QueryBuilder<Plaintext, C>,
        _op: QueryOp,
    ) -> Result<IndexTerm, EncryptionError> {
        // QUERY inputs (unlike stored values) must produce at least one
        // token: a zero-token bloom filter is empty, the EQL inclusion check
        // (query bits ⊆ row bits) is then vacuously true, and the query
        // silently matches EVERY row — turning e.g. a two-character LIKE
        // needle into SELECT *. Stored values stay unvalidated on purpose:
        // a short stored string is legal, it just can't be found via match.
        if let Plaintext::Text(Some(value)) = builder.plaintext() {
            self.validate_query_input(value)?;
        }
        let index_term = self.encrypt(builder.plaintext(), builder.index_key())?;
        Ok(index_term)
    }
}

pub struct MatchIndexerOptions {
    pub tokenizer: Tokenizer,
    pub token_filters: Vec<TokenFilter>,
    pub filter_opts: BloomFilterOps,
}

impl Default for MatchIndexerOptions {
    fn default() -> Self {
        Self {
            tokenizer: Tokenizer::Ngram { token_length: 3 },
            token_filters: vec![TokenFilter::Downcase],
            filter_opts: Default::default(),
        }
    }
}

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

    fn try_from(value: &column::IndexType) -> Result<Self, Self::Error> {
        match value {
            column::IndexType::Match {
                tokenizer,
                token_filters,
                k,
                m,
                ..
            } => Ok(Self {
                tokenizer: Tokenizer::from(*tokenizer),
                token_filters: token_filters.iter().copied().map(From::from).collect(),
                filter_opts: BloomFilterOps::default()
                    .with_filter_size(*m as u32)
                    .with_hash_function_count(*k),
            }),
            _ => Err(EncryptionError::IndexingError(
                "MatchIndexerOptions can only be created from a Match index configuration"
                    .to_string(),
            )),
        }
    }
}

impl Default for MatchIndexer {
    fn default() -> Self {
        Self::new(Default::default())
    }
}

pub struct MatchIndexer {
    tokenizer: Tokenizer,
    token_filters: Vec<TokenFilter>,
    filter_opts: BloomFilterOps,
}

impl MatchIndexer {
    pub fn new(
        MatchIndexerOptions {
            tokenizer,
            token_filters,
            filter_opts,
        }: MatchIndexerOptions,
    ) -> Self {
        Self {
            tokenizer,
            token_filters,
            filter_opts,
        }
    }

    /// Run `value` through the pipeline that derives bloom terms: the
    /// `%`/`_` char filter, the tokenizer, then the token filters. Both
    /// `encrypt` and `validate_query_input` use this — they must never
    /// disagree about what tokenizes to nothing.
    ///
    /// On the char filter: it removes '%' and '_' operators from the
    /// beginning and end of the string. This is a short term solution, so
    /// the most common of LIKE/ILIKE queries will work without any code
    /// changes (based off what is being used in the demo and current
    /// clients codebases).
    ///
    /// This covers the below LIKE/ILIKE queries: %value%, %value, value%
    ///
    /// Until we do a proper implementation, the below LIKE/ILIKE queries
    /// are not handled correctly: a%e, a_e, value_, _a%
    ///
    /// Also, this could mean that plaintext values that genuinely have
    /// these chars (%, _) will be stripped of those and returned results
    /// will be incomplete.
    ///
    /// A proper implementation of Like/ILike queries will be handled in
    /// this card:
    /// https://www.notion.so/cipherstash/WIP-Driver-more-robust-LIKE-op-handling-7ccf85c873374fb68ad651816f6bd9f6?pvs=4
    fn terms(&self, value: &str) -> Vec<String> {
        let filtered_output = char_filter_prefix_and_suffix(value, &['%', '_']);
        let tokens = self.tokenizer.process(filtered_output);
        self.token_filters
            .iter()
            .fold(tokens, |tokens, filter| filter.process(tokens))
    }

    /// The smallest input (in chars, after the `%`/`_` char filter) that can
    /// produce a token under this index's tokenizer, when that minimum is a
    /// simple length. `None` for tokenizers without a length threshold
    /// (`Standard`) — used only to make the validation error actionable.
    fn min_token_length(&self) -> Option<usize> {
        match self.tokenizer {
            Tokenizer::Ngram { token_length } => Some(token_length),
            Tokenizer::EdgeNgram { min_gram, .. } => Some(min_gram),
            Tokenizer::Standard => None,
        }
    }

    /// Reject a match QUERY input that would tokenize to nothing.
    ///
    /// A zero-token input yields an empty bloom filter, and the EQL match
    /// comparison (query bits ⊆ row bits) is vacuously true for an empty
    /// query — so the term silently matches every row. That's never what a
    /// query means, so callers minting match *query* terms must reject the
    /// input instead. The most common trigger is an input shorter than an
    /// `Ngram` tokenizer's `token_length` (default 3), but this also catches
    /// e.g. a `Stop` token filter swallowing every token.
    ///
    /// Storage-side indexing must NOT use this: storing a short string is
    /// legal (the value simply isn't findable via the match index).
    ///
    /// The error message never contains the input itself — only its length
    /// after the char filter, and the index's minimum.
    pub fn validate_query_input(&self, value: &str) -> Result<(), EncryptionError> {
        if !self.terms(value).is_empty() {
            return Ok(());
        }
        // The message's length is measured AFTER the %/_ char filter —
        // '%ab%' is four chars but only two survive trimming — and the
        // claim is "tokenizes to nothing", not "too short": token filters
        // (e.g. Stop) can empty the token list at any input length.
        let filtered_len = char_filter_prefix_and_suffix(value, &['%', '_'])
            .chars()
            .count();
        let msg = match self.min_token_length() {
            Some(min) => format!(
                "match query input tokenizes to nothing ({filtered_len} chars after wildcard trimming; the index's minimum token length is {min}): it would match every row",
            ),
            None => "match query input tokenizes to nothing (all tokens filtered out): it would match every row".to_string(),
        };
        Err(EncryptionError::InvalidValue(msg))
    }

    pub fn encrypt(
        &self,
        plaintext: &Plaintext,
        index_key: &IndexKey,
    ) -> Result<IndexTerm, EncryptionError> {
        match plaintext {
            Plaintext::Text(Some(value)) => {
                let terms = self.terms(value.as_str());

                // FIXME: Bloomfilter should be moved out of cipherstash-core and into client so that we can use the IndexKey and avoid this clone
                // Bloomfilter keys won't be zeroized at present
                // See https://linear.app/cipherstash/issue/CIP-844/wip-move-bloomfilter-out-of-cipherstash-core
                let mut filter =
                    BloomFilter::new(*index_key.key(), self.filter_opts).map_err(|e| {
                        EncryptionError::IndexingError(format!(
                            "Bloom Filter init failed with error {e}"
                        ))
                    })?;
                filter.add_terms(terms);

                Ok(IndexTerm::BitMap(filter.into_vec()))
            }
            Plaintext::Text(None) => Ok(IndexTerm::Null),
            _ => Err(EncryptionError::IndexingError(format!(
                "{plaintext:?} is not supported by match indexes"
            ))),
        }
    }
}

/// Converts a schema (cipherstash-config) Tokenizer (which has no impl)
/// into the type used in the `text` module which *is* implemented 😅
impl From<column::Tokenizer> for Tokenizer {
    fn from(value: column::Tokenizer) -> Self {
        match value {
            column::Tokenizer::Standard => Self::Standard,
            column::Tokenizer::Ngram { token_length } => Self::Ngram { token_length },
            column::Tokenizer::EdgeNgram { min_gram, max_gram } => {
                Self::EdgeNgram { min_gram, max_gram }
            }
        }
    }
}

/// Converts a schema (cipherstash-config) TokenFilter (which has no impl)
/// into the type used in the `text` module which *is* implemented 😅
impl From<column::TokenFilter> for TokenFilter {
    fn from(value: column::TokenFilter) -> Self {
        match value {
            column::TokenFilter::Downcase => TokenFilter::Downcase,
            column::TokenFilter::Upcase => TokenFilter::Upcase,
            column::TokenFilter::Stemmer => TokenFilter::Stemmer,
            column::TokenFilter::Stop => TokenFilter::Stop,
        }
    }
}

#[cfg(test)]
mod tests {
    use column::{ColumnConfig, Index};
    use zerokms_protocol::cipherstash_config::operator;

    use super::*;

    #[test]
    fn test_encrypt_term() -> Result<(), Box<dyn std::error::Error>> {
        let config = ColumnConfig::build("name").add_index(Index::new_match());
        let index = config
            .index_for_operator(&operator::Operator::Like)
            .unwrap();

        let index_key = IndexKey::from([0u8; 32]);
        let match_indexer_opts = MatchIndexerOptions::try_from(&index.index_type)?;
        let indexer = MatchIndexer::new(match_indexer_opts);

        let term = indexer.encrypt(&"Dan Draper".into(), &index_key)?;
        assert!(matches!(term, IndexTerm::BitMap(_)));

        Ok(())
    }

    // A query input shorter than the ngram token length tokenizes to
    // NOTHING, and an empty bloom filter matches every row (the EQL
    // inclusion check is vacuously true). `validate_query_input` is what
    // stands between a two-character LIKE needle and SELECT *.

    #[test]
    fn short_ngram_query_input_is_rejected() {
        // Default tokenizer: Ngram { token_length: 3 }.
        let indexer = MatchIndexer::default();

        for needle in ["a", "bp", "xy", ""] {
            let err = indexer
                .validate_query_input(needle)
                .expect_err("a needle shorter than the ngram must be rejected");
            let msg = err.to_string();
            assert!(
                msg.contains("minimum token length is 3"),
                "error should name the minimum, got: {msg}"
            );
        }

        // The message must describe the input, never contain it ("zq" is
        // chosen not to collide with any word in the message text).
        let msg = indexer
            .validate_query_input("zq")
            .expect_err("rejected")
            .to_string();
        assert!(!msg.contains("zq"), "error must not leak the needle: {msg}");

        assert!(indexer.validate_query_input("abc").is_ok());
        assert!(indexer.validate_query_input("Dan Draper").is_ok());
    }

    #[test]
    fn char_filtered_operators_do_not_count_toward_the_minimum() {
        // '%ab%' LIKE-strips to 'ab' — still too short for ngram(3).
        let indexer = MatchIndexer::default();
        assert!(indexer.validate_query_input("%ab%").is_err());
        assert!(indexer.validate_query_input("%abc%").is_ok());
    }

    #[test]
    fn edge_ngram_minimum_is_min_gram() {
        let indexer = MatchIndexer::new(MatchIndexerOptions {
            tokenizer: Tokenizer::EdgeNgram {
                min_gram: 2,
                max_gram: 5,
            },
            token_filters: vec![TokenFilter::Downcase],
            filter_opts: Default::default(),
        });
        assert!(indexer.validate_query_input("a").is_err());
        assert!(indexer.validate_query_input("ab").is_ok());
    }

    #[test]
    fn standard_tokenizer_accepts_single_characters() {
        // Standard splits on separators and has no length minimum — a
        // one-character word is a real token and must stay queryable.
        let indexer = MatchIndexer::new(MatchIndexerOptions {
            tokenizer: Tokenizer::Standard,
            token_filters: vec![TokenFilter::Downcase],
            filter_opts: Default::default(),
        });
        assert!(indexer.validate_query_input("a").is_ok());
    }

    #[test]
    fn all_tokens_filtered_out_is_also_rejected() {
        // The vacuous-query hazard isn't only length: a Stop filter that
        // swallows every token leaves the same empty bloom filter.
        let indexer = MatchIndexer::new(MatchIndexerOptions {
            tokenizer: Tokenizer::Standard,
            token_filters: vec![TokenFilter::Downcase, TokenFilter::Stop],
            filter_opts: Default::default(),
        });
        assert!(indexer.validate_query_input("the").is_err());
        assert!(indexer.validate_query_input("laplace").is_ok());
    }

    #[test]
    fn storing_a_short_value_still_succeeds() {
        // Storage is deliberately NOT validated: a short stored string is
        // legal — it simply can't be found via the match index. Erroring
        // here would reject legitimate writes.
        let indexer = MatchIndexer::default();
        let index_key = IndexKey::from([0u8; 32]);
        let term = indexer
            .encrypt(&"bp".into(), &index_key)
            .expect("storing a short value must not error");
        assert!(matches!(term, IndexTerm::BitMap(_)));
    }
}