HMAC sort key

Author: Bennett Hardwick

cryptonamo-derive/src/encryptable.rs

@@ -24,6 +24,11 @@ pub(crate) fn derive_encryptable(input: DeriveInput) -> Result<TokenStream, syn:
     let ident = settings.ident();
 
     let is_partition_key_encrypted = protected_attributes.contains(&partition_key_field.as_str());
+    let is_sort_key_encrypted = settings
+        .sort_key_field
+        .as_ref()
+        .map(|x| protected_attributes.contains(&x.as_str()))
+        .unwrap_or(true);
 
     let into_unsealed_impl = protected_attributes
         .iter()
@@ -67,6 +72,7 @@ pub(crate) fn derive_encryptable(input: DeriveInput) -> Result<TokenStream, syn:
         impl cryptonamo::traits::Encryptable for #ident {
             #primary_key_impl
 
+            #[inline]
             fn type_name() -> &'static str {
                 #type_name
             }
@@ -75,14 +81,21 @@ pub(crate) fn derive_encryptable(input: DeriveInput) -> Result<TokenStream, syn:
                 #sort_key_impl
             }
 
+            #[inline]
             fn sort_key_prefix() -> Option<&'static str> {
                 #sort_key_prefix
             }
 
+            #[inline]
             fn is_partition_key_encrypted() -> bool {
                 #is_partition_key_encrypted
             }
 
+            #[inline]
+            fn is_sort_key_encrypted() -> bool {
+                #is_sort_key_encrypted
+            }
+
             fn partition_key(&self) -> String {
                 self.#partition_key.to_string()
             }

src/crypto/sealer.rs

@@ -1,6 +1,7 @@
 use super::{encrypt_partition_key, SealError, Sealed, Unsealed, MAX_TERMS_PER_INDEX};
 use crate::{
     encrypted_table::{TableAttribute, TableEntry},
+    traits::PrimaryKeyParts,
     Searchable,
 };
 use cipherstash_client::{
@@ -54,22 +55,28 @@ impl<T> Sealer<T> {
         self,
         cipher: &Encryption<C>,
         term_length: usize, // TODO: SealError
-    ) -> Result<(String, Vec<Sealed>), SealError>
+    ) -> Result<(PrimaryKeyParts, Vec<Sealed>), SealError>
     where
         C: Credentials<Token = ViturToken>,
         T: Searchable,
     {
         let mut pk = self.inner.partition_key();
+        let mut sk = self.inner.sort_key();
 
         if T::is_partition_key_encrypted() {
-            // FIXME
-            pk = encrypt_partition_key(&self.inner.partition_key(), cipher).unwrap();
+            pk = encrypt_partition_key(&pk, cipher)?;
         }
 
-        let sk = self.inner.sort_key();
+        if T::is_sort_key_encrypted() {
+            sk = encrypt_partition_key(&sk, cipher)?;
+        }
 
-        let mut table_entry =
-            TableEntry::new_with_attributes(pk.clone(), sk, None, self.unsealed.unprotected());
+        let mut table_entry = TableEntry::new_with_attributes(
+            pk.clone(),
+            sk.clone(),
+            None,
+            self.unsealed.unprotected(),
+        );
 
         let protected = T::protected_attributes()
             .iter()
@@ -91,8 +98,6 @@ impl<T> Sealer<T> {
                 }
             });
 
-        let sort_key = self.inner.sort_key();
-
         let protected_indexes = T::protected_indexes();
         let terms: Vec<(&&str, Vec<u8>)> = protected_indexes
             .iter()
@@ -135,13 +140,13 @@ impl<T> Sealer<T> {
                         .clone()
                         .set_term(hex::encode(term))
                         // TODO: HMAC the sort key, too (users#index_name#pk)
-                        .set_sk(format!("{}#{}#{}", &sort_key, index_name, i)),
+                        .set_sk(format!("{}#{}#{}", &sk, index_name, i)),
                 )
             })
             .chain(once(Sealed(table_entry.clone())))
             .collect();
 
-        Ok((pk, table_entries))
+        Ok((PrimaryKeyParts { pk, sk }, table_entries))
     }
 
     #[allow(dead_code)]

src/encrypted_table/mod.rs

@@ -120,12 +120,16 @@ impl EncryptedTable {
         &self,
         k: impl Into<T::PrimaryKey>,
     ) -> Result<PrimaryKeyParts, EncryptionError> {
-        let PrimaryKeyParts { mut pk, sk } = k.into().into_parts::<T>();
+        let PrimaryKeyParts { mut pk, mut sk } = k.into().into_parts::<T>();
 
         if T::is_partition_key_encrypted() {
             pk = encrypt_partition_key(&pk, &self.cipher)?;
         }
 
+        if T::is_sort_key_encrypted() {
+            sk = encrypt_partition_key(&sk, &self.cipher)?;
+        }
+
         Ok(PrimaryKeyParts { pk, sk })
     }
 
@@ -160,7 +164,7 @@ impl EncryptedTable {
     ) -> Result<(), DeleteError> {
         let PrimaryKeyParts { pk, sk } = self.get_primary_key_parts::<E>(k)?;
 
-        let sk_to_delete = all_index_keys::<E>(&sk).into_iter().into_iter().chain([sk]);
+        let sk_to_delete = all_index_keys::<E>(&sk).into_iter().chain([sk]);
 
         let transact_items = sk_to_delete.map(|sk| {
             TransactWriteItem::builder()
@@ -194,9 +198,8 @@ impl EncryptedTable {
     {
         let mut seen_sk = HashSet::new();
 
-        let sk = record.sort_key();
         let sealer: Sealer<T> = record.into_sealer()?;
-        let (pk, sealed) = sealer.seal(&self.cipher, 12).await?;
+        let (PrimaryKeyParts { pk, sk }, sealed) = sealer.seal(&self.cipher, 12).await?;
 
         // TODO: Use a combinator
         let mut items: Vec<TransactWriteItem> = Vec::with_capacity(sealed.len());

src/traits/mod.rs

@@ -37,6 +37,8 @@ pub trait Encryptable: Debug + Sized {
 
     fn is_partition_key_encrypted() -> bool;
 
+    fn is_sort_key_encrypted() -> bool;
+
     fn sort_key(&self) -> String {
         Self::type_name().into()
     }

src/traits/primary_key.rs

@@ -2,7 +2,7 @@ use crate::Encryptable;
 
 pub struct PrimaryKeyParts {
     pub(crate) pk: String,
-    pub(crate) sk: String,
+    pub(crate) sk: String
 }
 
 pub trait PrimaryKey: private::Sealed {