use super::builder::StorageBuilder;
use super::indexer::{IndexerInit, Indexes, IndexesForQuery, QueryOp};
use super::text::TokenFilter;
use super::QueryBuilder;
use super::{errors::EncryptionError, plaintext::Plaintext, IndexTerm};
use crate::zerokms::IndexKey;
use cipherstash_config::column;
use hmac::{Hmac, Mac};
use sha2::Sha256;
type HmacSha256 = Hmac<Sha256>;
impl IndexerInit for UniqueIndexer {
type Args = UniqueIndexerOptions;
type Error = EncryptionError;
fn try_init<A>(args: A) -> Result<Self, Self::Error>
where
Self::Args: TryFrom<A, Error = Self::Error>,
{
let args = UniqueIndexerOptions::try_from(args)?;
Ok(Self::new(args))
}
}
impl<'k> Indexes<'k, Plaintext> for UniqueIndexer {
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 UniqueIndexer {
fn query_index(
&self,
builder: QueryBuilder<Plaintext, C>,
_op: QueryOp,
) -> Result<IndexTerm, EncryptionError> {
let index_term = self.encrypt(builder.plaintext(), builder.index_key())?;
Ok(index_term)
}
}
pub struct UniqueIndexerOptions {
pub token_filters: Vec<TokenFilter>,
}
impl TryFrom<&column::IndexType> for UniqueIndexerOptions {
type Error = EncryptionError;
fn try_from(value: &column::IndexType) -> Result<Self, Self::Error> {
match value {
column::IndexType::Unique { token_filters, .. } => Ok(Self {
token_filters: token_filters.iter().copied().map(From::from).collect(),
}),
_ => Err(EncryptionError::IndexingError(
"UniqueIndexerOptions can only be created from a Unique Index configuration"
.to_string(),
)),
}
}
}
impl Default for UniqueIndexerOptions {
fn default() -> Self {
Self {
token_filters: vec![TokenFilter::Downcase],
}
}
}
impl Default for UniqueIndexer {
fn default() -> Self {
Self::new(Default::default())
}
}
#[derive(Debug)]
pub struct UniqueIndexer {
token_filters: Vec<TokenFilter>,
}
impl UniqueIndexer {
pub fn new(opts: UniqueIndexerOptions) -> Self {
Self {
token_filters: opts.token_filters,
}
}
pub(super) fn create_hmac(&self, index_key: &IndexKey) -> Result<HmacSha256, EncryptionError> {
Ok(HmacSha256::new_from_slice(index_key.key())?)
}
pub(super) fn encrypt_into_hmac(
&self,
mac: &mut HmacSha256,
plaintext: &Plaintext,
) -> Result<(), EncryptionError> {
let plaintext_bytes = match plaintext {
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) => return Ok(()),
Plaintext::Utf8Str(Some(utf_str)) => {
let filtered_string = self
.token_filters
.iter()
.try_fold(utf_str.to_string(), |s, filter| filter.process_single(s));
Plaintext::from(filtered_string).to_vec()
}
x => x.to_vec(),
};
mac.update(&plaintext_bytes);
Ok(())
}
pub fn encrypt(
&self,
plaintext: &Plaintext,
index_key: &IndexKey,
) -> Result<IndexTerm, EncryptionError> {
if plaintext.is_null() {
Ok(IndexTerm::Null)
} else {
let mut mac = self.create_hmac(index_key)?;
self.encrypt_into_hmac(&mut mac, plaintext)?;
Ok(IndexTerm::Binary(mac.finalize().into_bytes().to_vec()))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_should_case_insensitive_compare() {
let index_key = [1; 32].into();
let indexer = UniqueIndexer::new(Default::default());
let first = indexer
.encrypt(&"hello WORLD".into(), &index_key)
.expect("Failed to encrypt");
let second = indexer
.encrypt(&"HELLO world".into(), &index_key)
.expect("Failed to encrypt");
assert_eq!(first, second);
}
}