Merge pull request #15 from cipherstash/feat/new-sort-primary

Change to dynamic sort key, new attrs, primary key abstraction

Author: Bennett Hardwick

cryptonamo-derive/src/encryptable.rs

@@ -33,13 +33,32 @@ pub(crate) fn derive_encryptable(input: DeriveInput) -> Result<TokenStream, syn:
             }
         }));
 
+    let sort_key_impl = if let Some(sort_key_field) = &settings.sort_key_field {
+        let sort_key_attr = format_ident!("{sort_key_field}");
+        quote! { format!("{}#{}", Self::type_name(), self.#sort_key_attr) }
+    } else {
+        quote! { Self::type_name().into() }
+    };
+
+    let primary_key_impl = if settings.sort_key_field.is_some() {
+        quote! { type PrimaryKey = cryptonamo::PkSk; }
+    } else {
+        quote! { type PrimaryKey = cryptonamo::Pk; }
+    };
+
     let expanded = quote! {
         #[automatically_derived]
         impl cryptonamo::traits::Encryptable for #ident {
+            #primary_key_impl
+
             fn type_name() -> &'static str {
                 #type_name
             }
 
+            fn sort_key(&self) -> String {
+                #sort_key_impl
+            }
+
             fn partition_key(&self) -> String {
                 self.#partition_key.to_string()
             }

cryptonamo-derive/src/lib.rs

@@ -10,21 +10,21 @@ mod settings;
 use proc_macro::TokenStream;
 use syn::{parse_macro_input, DeriveInput};
 
-#[proc_macro_derive(Encryptable, attributes(cryptonamo))]
+#[proc_macro_derive(Encryptable, attributes(cryptonamo, sort_key, partition_key))]
 pub fn derive_encryptable(input: TokenStream) -> TokenStream {
     encryptable::derive_encryptable(parse_macro_input!(input as DeriveInput))
         .unwrap_or_else(syn::Error::into_compile_error)
         .into()
 }
 
-#[proc_macro_derive(Decryptable, attributes(cryptonamo))]
+#[proc_macro_derive(Decryptable, attributes(cryptonamo, sort_key, partition_key))]
 pub fn derive_decryptable(input: TokenStream) -> TokenStream {
     decryptable::derive_decryptable(parse_macro_input!(input as DeriveInput))
         .unwrap_or_else(syn::Error::into_compile_error)
         .into()
 }
 
-#[proc_macro_derive(Searchable, attributes(cryptonamo))]
+#[proc_macro_derive(Searchable, attributes(cryptonamo, sort_key, partition_key))]
 pub fn derive_searchable(input: TokenStream) -> TokenStream {
     searchable::derive_searchable(parse_macro_input!(input as DeriveInput))
         .unwrap_or_else(syn::Error::into_compile_error)

cryptonamo-derive/src/settings/builder.rs

@@ -7,7 +7,8 @@ pub(crate) struct SettingsBuilder {
     ident: Ident,
     type_name: String,
     sort_key_prefix: Option<String>,
-    partition_key: Option<String>,
+    sort_key_field: Option<String>,
+    partition_key_field: Option<String>,
     protected_attributes: Vec<String>,
     unprotected_attributes: Vec<String>,
     skipped_attributes: Vec<String>,
@@ -33,7 +34,8 @@ impl SettingsBuilder {
             ident: input.ident.clone(),
             type_name,
             sort_key_prefix: None,
-            partition_key: None,
+            sort_key_field: None,
+            partition_key_field: None,
             protected_attributes: Vec::new(),
             unprotected_attributes: Vec::new(),
             skipped_attributes: Vec::new(),
@@ -92,6 +94,48 @@ impl SettingsBuilder {
 
                     // Parse the meta for the field
                     for attr in &field.attrs {
+                        if attr.path().is_ident("sort_key") {
+                            let field_name = ident
+                                .as_ref()
+                                .ok_or_else(|| {
+                                    syn::Error::new_spanned(
+                                        field,
+                                        "internal error: identifier was not Some",
+                                    )
+                                })?
+                                .to_string();
+
+                            if let Some(f) = &self.sort_key_field {
+                                return Err(syn::Error::new_spanned(
+                                    field,
+                                    format!("sort key was already specified to be '{f}'"),
+                                ));
+                            }
+
+                            self.sort_key_field = Some(field_name);
+                        }
+
+                        if attr.path().is_ident("partition_key") {
+                            let field_name = ident
+                                .as_ref()
+                                .ok_or_else(|| {
+                                    syn::Error::new_spanned(
+                                        field,
+                                        "internal error: identifier was not Some",
+                                    )
+                                })?
+                                .to_string();
+
+                            if let Some(f) = &self.partition_key_field {
+                                return Err(syn::Error::new_spanned(
+                                    field,
+                                    format!("partition key was already specified to be '{f}'"),
+                                ));
+                            }
+
+                            self.partition_key_field = Some(field_name);
+                        }
+
                         if attr.path().is_ident("cryptonamo") {
                             let mut query: Option<(String, String, Span)> = None;
                             let mut compound_index_name: Option<(String, Span)> = None;
@@ -203,7 +247,8 @@ impl SettingsBuilder {
             ident,
             type_name,
             sort_key_prefix,
-            partition_key,
+            sort_key_field,
+            partition_key_field: partition_key,
             protected_attributes,
             unprotected_attributes,
             skipped_attributes,
@@ -222,7 +267,8 @@ impl SettingsBuilder {
         Ok(Settings {
             ident,
             sort_key_prefix,
-            partition_key: Some(partition_key), // TODO: Remove the Some
+            sort_key_field,
+            partition_key_field: Some(partition_key), // TODO: Remove the Some
             protected_attributes,
             unprotected_attributes,
             skipped_attributes,
@@ -236,7 +282,7 @@ impl SettingsBuilder {
     }
 
     pub(crate) fn set_partition_key(&mut self, value: String) -> Result<(), syn::Error> {
-        self.partition_key = Some(value);
+        self.partition_key_field = Some(value);
         Ok(())
     }
 

cryptonamo-derive/src/settings/mod.rs

@@ -16,7 +16,8 @@ pub(crate) enum AttributeMode {
 pub(crate) struct Settings {
     ident: Ident,
     pub(crate) sort_key_prefix: String,
-    pub(crate) partition_key: Option<String>,
+    pub(crate) sort_key_field: Option<String>,
+    pub(crate) partition_key_field: Option<String>,
     protected_attributes: Vec<String>,
     unprotected_attributes: Vec<String>,
 
@@ -70,7 +71,7 @@ impl Settings {
     }
 
     pub(crate) fn get_partition_key(&self) -> Result<String, syn::Error> {
-        self.partition_key.clone().ok_or_else(|| {
+        self.partition_key_field.clone().ok_or_else(|| {
             syn::Error::new_spanned(
                 &self.sort_key_prefix,
                 "No partition key defined for this struct",

dynamo.sh

@@ -8,5 +8,5 @@ aws dynamodb create-table \
         AttributeName=pk,KeyType=HASH \
         AttributeName=sk,KeyType=RANGE \
     --provisioned-throughput ReadCapacityUnits=5,WriteCapacityUnits=5 \
-    --global-secondary-indexes "IndexName=TermIndex,KeySchema=[{AttributeName=term,KeyType=HASH},{AttributeName=pk,KeyType=RANGE}],Projection={ProjectionType=ALL},ProvisionedThroughput={ReadCapacityUnits=5,WriteCapacityUnits=5}" \
+    --global-secondary-indexes "IndexName=TermIndex,KeySchema=[{AttributeName=term,KeyType=HASH}],Projection={ProjectionType=ALL},ProvisionedThroughput={ReadCapacityUnits=5,WriteCapacityUnits=5}" \
     --endpoint-url http://localhost:8000

src/crypto/mod.rs

@@ -44,13 +44,13 @@ pub enum CryptoError {
     Other(String),
 }
 
-pub(crate) fn all_index_keys<E: Searchable + Encryptable>() -> Vec<String> {
+pub(crate) fn all_index_keys<E: Searchable + Encryptable>(sort_key: &str) -> Vec<String> {
     E::protected_indexes()
         .iter()
         .flat_map(|index_name| {
             (0..)
                 .take(MAX_TERMS_PER_INDEX)
-                .map(|i| format!("{}#{}#{}", E::type_name(), index_name, i))
+                .map(|i| format!("{}#{}#{}", sort_key, index_name, i))
                 .collect::<Vec<String>>()
         })
         .collect()

src/crypto/sealer.rs

@@ -60,13 +60,10 @@ impl<T> Sealer<T> {
         T: Searchable,
     {
         let pk = encrypt_partition_key(&self.inner.partition_key(), cipher).unwrap(); // FIXME
+        let sk = self.inner.sort_key();
 
-        let mut table_entry = TableEntry::new_with_attributes(
-            pk.clone(),
-            T::type_name().to_string(),
-            None,
-            self.unsealed.unprotected(),
-        );
+        let mut table_entry =
+            TableEntry::new_with_attributes(pk.clone(), sk, None, self.unsealed.unprotected());
 
         let protected = T::protected_attributes()
             .iter()
@@ -84,6 +81,8 @@ 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()
@@ -126,7 +125,7 @@ impl<T> Sealer<T> {
                         .clone()
                         .set_term(hex::encode(term))
                         // TODO: HMAC the sort key, too (users#index_name#pk)
-                        .set_sk(format!("{}#{}#{}", T::type_name(), index_name, i)),
+                        .set_sk(format!("{}#{}#{}", &sort_key, index_name, i)),
                 )
             })
             .chain(once(Sealed(table_entry.clone())))

src/encrypted_table/mod.rs

@@ -6,7 +6,10 @@ pub use self::{
 };
 use crate::{
     crypto::*,
-    traits::{Decryptable, ReadConversionError, Searchable, WriteConversionError},
+    traits::{
+        Decryptable, PrimaryKey, PrimaryKeyParts, ReadConversionError, Searchable,
+        WriteConversionError,
+    },
 };
 use aws_sdk_dynamodb::{
     types::{AttributeValue, Delete, Put, TransactWriteItem},
@@ -112,12 +115,13 @@ impl EncryptedTable {
         QueryBuilder::new(self)
     }
 
-    pub async fn get<T>(&self, pk: &str) -> Result<Option<T>, GetError>
+    pub async fn get<T>(&self, k: impl Into<T::PrimaryKey>) -> Result<Option<T>, GetError>
     where
         T: Decryptable,
     {
-        let pk = encrypt_partition_key(pk, &self.cipher)?;
-        let sk = T::type_name().to_string();
+        let PrimaryKeyParts { pk, sk } = k.into().into_parts(T::type_name());
+
+        let pk = encrypt_partition_key(&pk, &self.cipher)?;
 
         let result = self
             .db
@@ -138,12 +142,13 @@ impl EncryptedTable {
         }
     }
 
-    pub async fn delete<E: Searchable>(&self, pk: &str) -> Result<(), DeleteError> {
-        let pk = AttributeValue::S(encrypt_partition_key(pk, &self.cipher)?);
+    // TODO: create PrimaryKey abstraction
+    pub async fn delete<E: Searchable>(&self, k: impl Into<E::PrimaryKey>) -> Result<(), DeleteError> {
+        let PrimaryKeyParts { pk, sk } = k.into().into_parts(E::type_name());
+
+        let pk = AttributeValue::S(encrypt_partition_key(&pk, &self.cipher)?);
 
-        let sk_to_delete = [E::type_name().to_string()]
-            .into_iter()
-            .chain(all_index_keys::<E>().into_iter());
+        let sk_to_delete = all_index_keys::<E>(&sk).into_iter().into_iter().chain([sk]);
 
         let transact_items = sk_to_delete.map(|sk| {
             TransactWriteItem::builder()
@@ -177,6 +182,7 @@ 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?;
 
@@ -198,7 +204,7 @@ impl EncryptedTable {
             );
         }
 
-        for index_sk in all_index_keys::<T>() {
+        for index_sk in all_index_keys::<T>(&sk) {
             if seen_sk.contains(&index_sk) {
                 continue;
             }

src/lib.rs

@@ -403,7 +403,7 @@ mod error;
 pub mod traits;
 pub use encrypted_table::{EncryptedTable, QueryBuilder};
 pub use error::Error;
-pub use traits::{Decryptable, Encryptable, Searchable};
+pub use traits::{Decryptable, Encryptable, Searchable, PrimaryKey, PkSk, Pk};
 
 #[doc(hidden)]
 pub use cryptonamo_derive::{Decryptable, Encryptable, Searchable};

src/traits/mod.rs

@@ -6,6 +6,10 @@ pub use cipherstash_client::encryption::{
     },
     Plaintext,
 };
+
+mod primary_key;
+pub use primary_key::*;
+
 use std::fmt::Debug;
 use thiserror::Error;
 
@@ -24,8 +28,15 @@ pub enum WriteConversionError {
 }
 
 pub trait Encryptable: Debug + Sized {
+    type PrimaryKey: PrimaryKey;
+
     // TODO: Add a function indicating that the root should be stored
     fn type_name() -> &'static str;
+
+    fn sort_key(&self) -> String {
+        Self::type_name().into()
+    }
+
     fn partition_key(&self) -> String;
 
     fn protected_attributes() -> Vec<&'static str>;
@@ -56,6 +67,7 @@ pub trait Searchable: Encryptable {
 
 pub trait Decryptable: Encryptable {
     /// Convert an `Unsealed` into a `Self`.
+
     fn from_unsealed(unsealed: Unsealed) -> Result<Self, SealError>;
 
     /// Defines which attributes are decryptable for this type.