Change to prepared record

Author: Bennett Hardwick

cipherstash-dynamodb-derive/src/encryptable.rs

@@ -37,14 +37,14 @@ pub(crate) fn derive_encryptable(input: DeriveInput) -> Result<TokenStream, syn:
             let attr_ident = format_ident!("{attr}");
 
             quote! {
-                .add_protected(#attr, |x| cipherstash_dynamodb::traits::Plaintext::from(x.#attr_ident.to_owned()))
+                unsealed.add_protected(#attr, cipherstash_dynamodb::traits::Plaintext::from(self.#attr_ident.to_owned()));
             }
         })
         .chain(plaintext_attributes.iter().map(|attr| {
             let attr_ident = format_ident!("{attr}");
 
             quote! {
-                .add_plaintext(#attr, |x| cipherstash_dynamodb::traits::TableAttribute::from(x.#attr_ident.clone()))
+                unsealed.add_unprotected(#attr, cipherstash_dynamodb::traits::TableAttribute::from(self.#attr_ident.clone()));
             }
         }));
 
@@ -70,9 +70,12 @@ pub(crate) fn derive_encryptable(input: DeriveInput) -> Result<TokenStream, syn:
             }
 
             #[allow(clippy::needless_question_mark)]
-            fn into_sealer(self) -> Result<cipherstash_dynamodb::crypto::Sealer<Self>, cipherstash_dynamodb::crypto::SealError> {
-                Ok(cipherstash_dynamodb::crypto::Sealer::new_with_descriptor(self, <Self as cipherstash_dynamodb::traits::Encryptable>::type_name())
-                    #(#into_unsealed_impl?)*)
+            fn into_unsealed(self) -> cipherstash_dynamodb::crypto::Unsealed {
+                let mut unsealed = cipherstash_dynamodb::crypto::Unsealed::new_with_descriptor(<Self as cipherstash_dynamodb::traits::Encryptable>::type_name());
+
+                #(#into_unsealed_impl)*
+
+                unsealed
             }
         }
     };

src/crypto/mod.rs

@@ -3,6 +3,8 @@ mod sealed;
 mod sealer;
 mod unsealed;
 
+use std::borrow::Cow;
+
 use crate::{
     traits::{
         Encryptable, PrimaryKeyError, PrimaryKeyParts, ReadConversionError, Searchable,
@@ -66,8 +68,12 @@ pub fn format_term_key(
     format!("{sort_key}#{index_name}#{index_type}#{counter}")
 }
 
-pub(crate) fn all_index_keys<E: Searchable + Encryptable>(sort_key: &str) -> Vec<String> {
-    E::protected_indexes()
+pub(crate) fn all_index_keys<'a>(
+    sort_key: &str,
+    protected_indexes: impl AsRef<[(Cow<'a, str>, IndexType)]>,
+) -> Vec<String> {
+    protected_indexes
+        .as_ref()
         .iter()
         .flat_map(|(index_name, index_type)| {
             (0..)

src/crypto/sealer.rs

@@ -1,6 +1,6 @@
 use super::{
-    b64_encode, encrypt_primary_key, format_term_key, hmac, SealError, SealedTableEntry, Unsealed,
-    MAX_TERMS_PER_INDEX,
+    b64_encode, encrypt_primary_key, format_term_key, hmac, SealError, SealedTableEntry,
+    UnsealSpec, Unsealed, MAX_TERMS_PER_INDEX,
 };
 use crate::{
     encrypted_table::{TableAttribute, TableEntry},
@@ -9,15 +9,32 @@ use crate::{
 };
 use cipherstash_client::{
     credentials::{service_credentials::ServiceToken, Credentials},
-    encryption::{compound_indexer::CompoundIndex, Encryption, IndexTerm, Plaintext},
+    encryption::{
+        compound_indexer::{ComposableIndex, ComposablePlaintext, CompoundIndex},
+        Encryption, IndexTerm, Plaintext,
+    },
 };
 use itertools::Itertools;
 use std::{borrow::Cow, collections::HashMap, ops::Deref};
 
 /// Builder pattern for sealing a record of type, `T`.
-pub struct Sealer<T> {
-    inner: T,
-    unsealed: Unsealed,
+pub struct Sealer {
+    pub(crate) pk: String,
+    pub(crate) sk: String,
+
+    pub(crate) is_pk_encrypted: bool,
+    pub(crate) is_sk_encrypted: bool,
+
+    pub(crate) type_name: Cow<'static, str>,
+
+    pub(crate) unsealed_indexes: Vec<(
+        ComposablePlaintext,
+        Box<dyn ComposableIndex>,
+        Cow<'static, str>,
+        IndexType,
+    )>,
+
+    pub(crate) unsealed: Unsealed,
 }
 
 struct Term {
@@ -82,94 +99,110 @@ impl Sealed {
     }
 }
 
-impl<T> Sealer<T> {
-    pub fn new(inner: T) -> Self {
-        Self {
-            inner,
-            unsealed: Unsealed::new(),
-        }
-    }
+impl Sealer {
+    // pub fn new<T>(inner: T) -> Self {
+    //     Self {
+    //         inner,
+    //         unsealed: Unsealed::new(),
+    //     }
+    // }
+
+    // pub fn new_with_descriptor<T>(inner: T, descriptor: impl Into<String>) -> Self {
+    //     Self {
+    //         inner,
+    //         unsealed: Unsealed::new_with_descriptor(descriptor),
+    //     }
+    // }
+
+    // pub fn add_protected<F>(mut self, name: impl Into<String>, f: F) -> Result<Self, SealError>
+    // where
+    //     F: FnOnce(&T) -> Plaintext,
+    // {
+    //     let name: String = name.into();
+
+    //     self.unsealed.add_protected(name, f(&self.inner));
+    //     Ok(self)
+    // }
+
+    // pub fn add_plaintext<F>(mut self, name: impl Into<String>, f: F) -> Result<Self, SealError>
+    // where
+    //     F: FnOnce(&T) -> TableAttribute,
+    // {
+    //     let name: String = name.into();
+    //     self.unsealed.add_unprotected(name, f(&self.inner));
+    //     Ok(self)
+    // }
+
+    pub(crate) async fn seal_all<'a>(
+        records: impl IntoIterator<Item = Sealer>,
+        protected_attributes: impl AsRef<[Cow<'a, str>]>,
+        cipher: &Encryption<impl Credentials<Token = ServiceToken>>,
+        term_length: usize,
+    ) -> Result<Vec<Sealed>, SealError> {
+        let records = records.into_iter();
+        let records_len = records.size_hint().1.unwrap_or(1);
 
-    pub fn new_with_descriptor(inner: T, descriptor: impl Into<String>) -> Self {
-        Self {
-            inner,
-            unsealed: Unsealed::new_with_descriptor(descriptor),
-        }
-    }
+        let protected_attributes = protected_attributes.as_ref();
 
-    pub fn add_protected<F>(mut self, name: impl Into<String>, f: F) -> Result<Self, SealError>
-    where
-        F: FnOnce(&T) -> Plaintext,
-    {
-        let name: String = name.into();
+        let mut protected = Vec::with_capacity(records_len * protected_attributes.len());
+        let mut table_entries = Vec::with_capacity(records_len);
 
-        self.unsealed.add_protected(name, f(&self.inner));
-        Ok(self)
-    }
+        for mut record in records {
+            let mut pk = record.pk;
+            let mut sk = record.sk;
 
-    pub fn add_plaintext<F>(mut self, name: impl Into<String>, f: F) -> Result<Self, SealError>
-    where
-        F: FnOnce(&T) -> TableAttribute,
-    {
-        let name: String = name.into();
-        self.unsealed.add_unprotected(name, f(&self.inner));
-        Ok(self)
-    }
+            if record.is_pk_encrypted {
+                pk = b64_encode(hmac(&pk, None, cipher)?);
+            }
 
-    pub(crate) async fn seal_all<S: Searchable + Identifiable>(
-        records: impl AsRef<[Sealer<S>]>,
-        cipher: &Encryption<impl Credentials<Token = ServiceToken>>,
-        term_length: usize,
-    ) -> Result<Vec<Sealed>, SealError> {
-        let records = records.as_ref();
-        let protected_attributes = S::protected_attributes();
-        let protected_indexes = S::protected_indexes();
+            if record.is_sk_encrypted {
+                sk = b64_encode(hmac(&sk, Some(pk.as_str()), cipher)?);
+            }
 
-        let mut protected = Vec::with_capacity(records.len() * protected_attributes.len());
-        let mut table_entries = Vec::with_capacity(records.len());
+            let type_name = &record.type_name;
 
-        for record in records.iter() {
-            let PrimaryKeyParts { pk, sk } = encrypt_primary_key::<S>(
-                record.inner.get_primary_key(),
-                &S::type_name(),
-                S::sort_key_prefix().as_deref(),
-                cipher,
-            )?;
+            // let PrimaryKeyParts { pk, sk } = encrypt_primary_key::<S>(
+            //     record.inner.get_primary_key(),
+            //     &S::type_name(),
+            //     S::sort_key_prefix().as_deref(),
+            //     cipher,
+            // )?;
 
             for attr in protected_attributes.iter() {
-                protected.push(record.unsealed.protected_with_descriptor(attr)?);
+                protected.push(record.unsealed.remove_protected_with_descriptor(attr)?);
             }
 
-            let terms: Vec<(&Cow<'_, str>, IndexType, Vec<u8>)> = protected_indexes
-                .iter()
-                .map(|(index_name, index_type)| {
-                    record
-                        .inner
-                        .attribute_for_index(index_name, *index_type)
-                        .and_then(|attr| {
-                            S::index_by_name(index_name, *index_type)
-                                .map(|index| (attr, index, index_name, index_type))
-                        })
-                        .ok_or(SealError::MissingAttribute(index_name.to_string()))
-                        .and_then(|(attr, index, index_name, index_type)| {
-                            let term = cipher.compound_index(
-                                &CompoundIndex::new(index),
-                                attr,
-                                Some(format!("{}#{}", S::type_name(), index_name)),
-                                term_length,
-                            )?;
-
-                            Ok::<_, SealError>((index_name, index_type, term))
-                        })
+            let terms: Vec<(Cow<'_, str>, IndexType, Vec<u8>)> = record
+                .unsealed_indexes
+                .into_iter()
+                // .iter()
+                // .map(|(index_name, index_type)| {
+                //     record
+                //         .inner
+                //         .attribute_for_index(index_name, *index_type)
+                //         .and_then(|attr| {
+                //             S::index_by_name(index_name, *index_type)
+                //                 .map(|index| (attr, index, index_name, index_type))
+                //         })
+                //         .ok_or(SealError::MissingAttribute(index_name.to_string()))
+                .map(|(attr, index, index_name, index_type)| {
+                    let term = cipher.compound_index(
+                        &CompoundIndex::new(index),
+                        attr,
+                        Some(format!("{}#{}", type_name, index_name)),
+                        term_length,
+                    )?;
+
+                    Ok::<_, SealError>((index_name, index_type, term))
                 })
                 .map(|index_term| match index_term {
                     Ok((index_name, index_type, IndexTerm::Binary(x))) => {
-                        Ok(vec![(index_name, *index_type, x)])
+                        Ok(vec![(index_name, index_type, x)])
                     }
                     Ok((index_name, index_type, IndexTerm::BinaryVec(x))) => Ok(x
                         .into_iter()
                         .take(MAX_TERMS_PER_INDEX)
-                        .map(|x| (index_name, *index_type, x))
+                        .map(|x| (index_name.clone(), index_type, x))
                         .collect()),
                     _ => Err(SealError::InvalidCiphertext("Invalid index term".into())),
                 })
@@ -183,7 +216,9 @@ impl<T> Sealer<T> {
             ));
         }
 
-        let encrypted = cipher.encrypt(protected).await?;
+        let encrypted = cipher
+            .encrypt(protected.iter().map(|(a, b)| (a, b.as_str())))
+            .await?;
 
         for (encrypted, (_, attributes, _)) in encrypted
             .chunks_exact(protected_attributes.len())
@@ -207,7 +242,7 @@ impl<T> Sealer<T> {
             }
         }
 
-        let mut output = Vec::with_capacity(records.len());
+        let mut output = Vec::with_capacity(table_entries.len());
 
         for (PrimaryKeyParts { pk, sk }, attributes, terms) in table_entries.into_iter() {
             let terms = terms
@@ -235,16 +270,13 @@ impl<T> Sealer<T> {
         Ok(output)
     }
 
-    pub(crate) async fn seal<C>(
+    pub(crate) async fn seal<'a>(
         self,
-        cipher: &Encryption<C>,
+        protected_attributes: impl AsRef<[Cow<'a, str>]>,
+        cipher: &Encryption<impl Credentials<Token = ServiceToken>>,
         term_length: usize,
-    ) -> Result<Sealed, SealError>
-    where
-        C: Credentials<Token = ServiceToken>,
-        T: Searchable + Identifiable,
-    {
-        let mut vec = Self::seal_all([self], cipher, term_length).await?;
+    ) -> Result<Sealed, SealError> {
+        let mut vec = Self::seal_all([self], protected_attributes, cipher, term_length).await?;
 
         if vec.len() != 1 {
             let actual = vec.len();

src/crypto/unsealed.rs

@@ -15,15 +15,15 @@ pub struct Unsealed {
 }
 
 impl Unsealed {
-    pub(super) fn new() -> Self {
+    pub fn new() -> Self {
         Self {
             descriptor: None,
             protected: Default::default(),
             unprotected: Default::default(),
         }
     }
 
-    pub(super) fn new_with_descriptor(descriptor: impl Into<String>) -> Self {
+    pub fn new_with_descriptor(descriptor: impl Into<String>) -> Self {
         Self {
             descriptor: Some(descriptor.into()),
             protected: Default::default(),
@@ -47,30 +47,31 @@ impl Unsealed {
             .ok_or_else(|| SealError::MissingAttribute(name.to_string()))
     }
 
-    pub(super) fn add_protected(&mut self, name: impl Into<String>, plaintext: Plaintext) {
+    pub fn add_protected(&mut self, name: impl Into<String>, plaintext: Plaintext) {
         let name = name.into();
         let descriptor = format!("{}/{}", self.descriptor.as_deref().unwrap_or(""), &name);
         self.protected.insert(name, (plaintext, descriptor));
     }
 
-    pub(super) fn add_unprotected(&mut self, name: impl Into<String>, attribute: TableAttribute) {
+    pub fn add_unprotected(&mut self, name: impl Into<String>, attribute: TableAttribute) {
         self.unprotected.insert(name.into(), attribute);
     }
 
     pub(crate) fn unprotected(&self) -> HashMap<String, TableAttribute> {
         self.unprotected.clone()
     }
 
-    /// Returns a reference to the protected value along with its descriptor.
-    pub(crate) fn protected_with_descriptor(
-        &self,
+    /// Remove and return a protected value along with its descriptor.
+    pub(crate) fn remove_protected_with_descriptor(
+        &mut self,
         name: &str,
-    ) -> Result<(&Plaintext, &str), SealError> {
+    ) -> Result<(Plaintext, String), SealError> {
         let out = self
             .protected
-            .get(name)
+            .remove(name)
             .ok_or(SealError::MissingAttribute(name.to_string()))?;
-        Ok((&out.0, &out.1))
+
+        Ok(out)
     }
 
     pub fn into_value<T: Decryptable>(self) -> Result<T, SealError> {