Added derive macro for custom type handling, updated tests

Author: coderdan

cipherstash-dynamodb-derive/src/decryptable.rs

@@ -9,7 +9,7 @@ pub(crate) fn derive_decryptable(input: DeriveInput) -> Result<TokenStream, syn:
         .field_attributes(&input)?
         .build()?;
 
-    let protected_attributes = settings.protected_attributes();
+    let protected_excluding_handlers = settings.protected_attributes_excluding_handlers();
     let plaintext_attributes = settings.plaintext_attributes();
 
     let protected_attributes_cow = settings
@@ -25,7 +25,7 @@ pub(crate) fn derive_decryptable(input: DeriveInput) -> Result<TokenStream, syn:
     let skipped_attributes = settings.skipped_attributes();
     let ident = settings.ident();
 
-    let from_unsealed_impl = protected_attributes
+    let from_unsealed_impl = protected_excluding_handlers
         .iter()
         .map(|attr| {
             let attr_ident = format_ident!("{attr}");
@@ -47,6 +47,13 @@ pub(crate) fn derive_decryptable(input: DeriveInput) -> Result<TokenStream, syn:
             quote! {
                 #attr_ident: Default::default()
             }
+        }))
+        .chain(settings.decrypt_handlers().iter().map(|(attr, handler)| {
+            let attr_ident = format_ident!("{attr}");
+
+            quote! {
+                #attr_ident: #handler(&mut unsealed)?
+            }
         }));
 
     let expanded = quote! {

cipherstash-dynamodb-derive/src/encryptable.rs

@@ -9,7 +9,7 @@ pub(crate) fn derive_encryptable(input: DeriveInput) -> Result<TokenStream, syn:
         .field_attributes(&input)?
         .build()?;
 
-    let protected_attributes = settings.protected_attributes();
+    let protected_excluding_handlers = settings.protected_attributes_excluding_handlers();
     let plaintext_attributes = settings.plaintext_attributes();
 
     let protected_attributes_cow = settings
@@ -24,20 +24,27 @@ pub(crate) fn derive_encryptable(input: DeriveInput) -> Result<TokenStream, syn:
 
     let ident = settings.ident();
 
-    let into_unsealed_impl = protected_attributes
+    let into_unsealed_impl = protected_excluding_handlers
         .iter()
         .map(|attr| {
             let attr_ident = format_ident!("{attr}");
 
             quote! {
-                unsealed.add_protected(#attr, cipherstash_dynamodb::traits::Plaintext::from(self.#attr_ident.to_owned()));
+                unsealed.add_protected(#attr, self.#attr_ident);
             }
         })
         .chain(plaintext_attributes.iter().map(|attr| {
             let attr_ident = format_ident!("{attr}");
 
             quote! {
-                unsealed.add_unprotected(#attr, cipherstash_dynamodb::traits::TableAttribute::from(self.#attr_ident.clone()));
+                unsealed.add_unprotected(#attr, self.#attr_ident);
+            }
+        }))
+        .chain(settings.encrypt_handlers().iter().map(|(attr, handler)| {
+            let attr_ident = format_ident!("{attr}");
+
+            quote! {
+                #handler(&mut unsealed, self.#attr_ident);
             }
         }));
 

cipherstash-dynamodb-derive/src/settings/builder.rs

@@ -1,7 +1,7 @@
 use super::{index_type::IndexType, AttributeMode, Settings};
 use proc_macro2::{Ident, Span};
 use std::collections::HashMap;
-use syn::{Data, DeriveInput, Fields, LitStr};
+use syn::{Data, DeriveInput, Fields, LitStr, ExprPath};
 
 enum SortKeyPrefix {
     Default,
@@ -31,6 +31,8 @@ pub(crate) struct SettingsBuilder {
     unprotected_attributes: Vec<String>,
     skipped_attributes: Vec<String>,
     indexes: Vec<IndexType>,
+    encrypt_handlers: HashMap<String, ExprPath>,
+    decrypt_handlers: HashMap<String, ExprPath>,
 }
 
 impl SettingsBuilder {
@@ -58,6 +60,8 @@ impl SettingsBuilder {
             unprotected_attributes: Vec::new(),
             skipped_attributes: Vec::new(),
             indexes: Vec::new(),
+            encrypt_handlers: HashMap::new(),
+            decrypt_handlers: HashMap::new(),
         }
     }
 
@@ -288,6 +292,18 @@ impl SettingsBuilder {
 
                                     Ok(())
                                 }
+                                Some("encryptable_with") => {
+                                    let value = meta.value()?;
+                                    let handler = value.parse::<ExprPath>()?;
+                                    self.encrypt_handlers.insert(field_name.clone(), handler);
+                                    Ok(())
+                                }
+                                Some("decryptable_with") => {
+                                    let value = meta.value()?;
+                                    let handler = value.parse::<ExprPath>()?;
+                                    self.decrypt_handlers.insert(field_name.clone(), handler);
+                                    Ok(())
+                                }
                                 _ => Err(meta.error("unsupported field attribute")),
                             }
                         })?;
@@ -308,7 +324,10 @@ impl SettingsBuilder {
                                 }
 
                                 (None, Some((compound_index_name, span))) => {
-                                    return Err(syn::Error::new(span,  format!("Compound attribute was specified but no query options were. Specify how this field should be queried with the attribute #[cipherstash(query = <option>, compound = \"{compound_index_name}\")]")));
+                                    return Err(syn::Error::new(
+                                        span,
+                                        format!("Compound attribute was specified but no query options were. Specify how this field should be queried with the attribute #[cipherstash(query = <option>, compound = \"{compound_index_name}\")]"))
+                                    );
                                 }
 
                                 (None, None) => {}
@@ -345,6 +364,8 @@ impl SettingsBuilder {
             unprotected_attributes,
             skipped_attributes,
             indexes,
+            encrypt_handlers,
+            decrypt_handlers,
         } = self;
 
         let sort_key_prefix = sort_key_prefix.into_prefix(&type_name);
@@ -359,6 +380,8 @@ impl SettingsBuilder {
             unprotected_attributes,
             skipped_attributes,
             indexes,
+            encrypt_handlers,
+            decrypt_handlers,
         })
     }
 
@@ -499,4 +522,4 @@ impl SettingsBuilder {
 
         Ok(())
     }
-}
+}

cipherstash-dynamodb-derive/src/settings/mod.rs

@@ -1,9 +1,11 @@
 mod builder;
 pub mod index_type;
+use std::collections::HashMap;
+
 use self::{builder::SettingsBuilder, index_type::IndexType};
 use itertools::Itertools;
 use proc_macro2::Ident;
-use syn::DeriveInput;
+use syn::{DeriveInput, ExprPath};
 
 pub(crate) enum AttributeMode {
     Protected,
@@ -20,6 +22,12 @@ pub(crate) struct Settings {
     protected_attributes: Vec<String>,
     unprotected_attributes: Vec<String>,
 
+    /// Map of attribute names to the encryption handler to use.
+    encrypt_handlers: HashMap<String, ExprPath>,
+
+    /// Map of attribute names to the decryption handler to use.
+    decrypt_handlers: HashMap<String, ExprPath>,
+
     /// Skipped attributes are never encrypted by the `DecryptedRecord` trait will
     /// use these to reconstruct the struct via `Default` (like serde).
     skipped_attributes: Vec<String>,
@@ -43,6 +51,24 @@ impl Settings {
             .collect::<Vec<_>>()
     }
 
+    pub(crate) fn protected_attributes_excluding_handlers(&self) -> Vec<&str> {
+        self.protected_attributes
+            .iter()
+            .filter(|s| !self.encrypt_handlers.contains_key(s.as_str()))
+            .filter(|s| !self.decrypt_handlers.contains_key(s.as_str()))
+            .map(|s| s.as_str())
+            .sorted()
+            .collect::<Vec<_>>()
+    }
+
+    pub(crate) fn encrypt_handlers(&self) -> &HashMap<String, ExprPath> {
+        &self.encrypt_handlers
+    }
+
+    pub(crate) fn decrypt_handlers(&self) -> &HashMap<String, ExprPath> {
+        &self.decrypt_handlers
+    }
+
     pub(crate) fn plaintext_attributes(&self) -> Vec<&str> {
         self.unprotected_attributes
             .iter()

src/crypto/attrs/normalized_protected_attributes.rs

@@ -96,18 +96,17 @@ impl FromIterator<(NormalizedKey, NormalizedValue)> for NormalizedProtectedAttri
 
 impl FromIterator<FlattenedProtectedAttribute> for NormalizedProtectedAttributes {
     fn from_iter<T: IntoIterator<Item = FlattenedProtectedAttribute>>(iter: T) -> Self {
-        iter.into_iter()
-            .fold(Self::new(), |mut acc, fpa| {
-                match fpa.normalize_into_parts() {
-                    (plaintext, key, Some(subkey)) => {
-                        acc.insert_and_update_map(key, subkey, plaintext);
-                    }
-                    (plaintext, key, None) => {
-                        acc.insert(key, plaintext);
-                    }
+        iter.into_iter().fold(Self::new(), |mut acc, fpa| {
+            match fpa.normalize_into_parts() {
+                (plaintext, key, Some(subkey)) => {
+                    acc.insert_and_update_map(key, subkey, plaintext);
                 }
-                acc
-            })
+                (plaintext, key, None) => {
+                    acc.insert(key, plaintext);
+                }
+            }
+            acc
+        })
     }
 }
 

tests/nested_tests.rs

@@ -1,16 +1,14 @@
 mod common;
-
-// TODO: Use the derive macros for this test
 use cipherstash_client::encryption::TypeParseError;
 use cipherstash_dynamodb::{
     crypto::Unsealed,
     errors::SealError,
-    traits::{Plaintext, TryFromPlaintext, TryFromTableAttr},
-    Decryptable, Encryptable, EncryptedTable, Identifiable, PkSk,
+    traits::{Plaintext, TryFromPlaintext},
+    Decryptable, Encryptable, EncryptedTable, Identifiable,
 };
 use cipherstash_dynamodb_derive::Searchable;
 use miette::IntoDiagnostic;
-use std::{borrow::Cow, collections::BTreeMap};
+use std::collections::BTreeMap;
 
 fn make_btree_map() -> BTreeMap<String, String> {
     let mut map = BTreeMap::new();
@@ -20,116 +18,40 @@ fn make_btree_map() -> BTreeMap<String, String> {
     map
 }
 
-#[derive(Debug, Clone, PartialEq, Searchable)]
+#[derive(Debug, Clone, PartialEq, Searchable, Encryptable, Decryptable, Identifiable)]
 struct Test {
     #[partition_key]
     pub pk: String,
     #[sort_key]
     pub sk: String,
     pub name: String,
     pub age: i16,
+    #[cipherstash(plaintext)]
     pub tag: String,
+    #[cipherstash(encryptable_with = put_attrs, decryptable_with = get_attrs)]
     pub attrs: BTreeMap<String, String>,
 }
 
-impl Identifiable for Test {
-    type PrimaryKey = PkSk;
-
-    fn get_primary_key(&self) -> Self::PrimaryKey {
-        PkSk(self.pk.to_string(), self.sk.to_string())
-    }
-    #[inline]
-    fn type_name() -> Cow<'static, str> {
-        std::borrow::Cow::Borrowed("test")
-    }
-    #[inline]
-    fn sort_key_prefix() -> Option<Cow<'static, str>> {
-        None
-    }
-    fn is_pk_encrypted() -> bool {
-        false
-    }
-    fn is_sk_encrypted() -> bool {
-        false
-    }
-}
-
 fn put_attrs(unsealed: &mut Unsealed, attrs: BTreeMap<String, String>) {
     attrs.into_iter().for_each(|(k, v)| {
         unsealed.add_protected_map_field("attrs", k, Plaintext::from(v));
     })
 }
 
-impl Encryptable for Test {
-    fn protected_attributes() -> Cow<'static, [Cow<'static, str>]> {
-        Cow::Borrowed(&[
-            Cow::Borrowed("name"),
-            Cow::Borrowed("age"),
-            Cow::Borrowed("attrs"),
-        ])
-    }
-
-    fn plaintext_attributes() -> Cow<'static, [Cow<'static, str>]> {
-        Cow::Borrowed(&[
-            Cow::Borrowed("tag"),
-            Cow::Borrowed("pk"),
-            Cow::Borrowed("sk"),
-        ])
-    }
-
-    fn into_unsealed(self) -> Unsealed {
-        let mut unsealed = Unsealed::new_with_descriptor(<Self as Identifiable>::type_name());
-        unsealed.add_unprotected("pk", self.pk);
-        unsealed.add_unprotected("sk", self.sk);
-        unsealed.add_protected("name", self.name);
-        unsealed.add_protected("age", self.age);
-        unsealed.add_unprotected("tag", self.tag);
-        put_attrs(&mut unsealed, self.attrs);
-        println!("Encryption: {:?}", unsealed);
-        unsealed
-    }
-}
-
-fn get_attrs<T>(unsealed: &mut Unsealed) -> Result<T, TypeParseError>
+fn get_attrs<T>(unsealed: &mut Unsealed) -> Result<T, SealError>
 where
     T: FromIterator<(String, String)>,
 {
     unsealed
         .take_protected_map("attrs")
         .ok_or(TypeParseError("attrs".to_string()))?
         .into_iter()
-        .map(|(k, v)| TryFromPlaintext::try_from_plaintext(v).map(|v| (k, v)))
-        .collect()
-}
-
-impl Decryptable for Test {
-    fn from_unsealed(mut unsealed: Unsealed) -> Result<Self, SealError> {
-        println!("{:?}", unsealed);
-        Ok(Self {
-            pk: TryFromTableAttr::try_from_table_attr(unsealed.get_plaintext("pk"))?,
-            sk: TryFromTableAttr::try_from_table_attr(unsealed.get_plaintext("sk"))?,
-            name: TryFromPlaintext::try_from_optional_plaintext(unsealed.take_protected("name"))?,
-            age: TryFromPlaintext::try_from_optional_plaintext(unsealed.take_protected("age"))?,
-            tag: TryFromTableAttr::try_from_table_attr(unsealed.get_plaintext("tag"))?,
-            attrs: get_attrs(&mut unsealed)?,
+        .map(|(k, v)| {
+            TryFromPlaintext::try_from_plaintext(v)
+                .map(|v| (k, v))
+                .map_err(SealError::from)
         })
-    }
-
-    fn protected_attributes() -> Cow<'static, [Cow<'static, str>]> {
-        Cow::Borrowed(&[
-            Cow::Borrowed("name"),
-            Cow::Borrowed("age"),
-            Cow::Borrowed("attrs"),
-        ])
-    }
-
-    fn plaintext_attributes() -> Cow<'static, [Cow<'static, str>]> {
-        Cow::Borrowed(&[
-            Cow::Borrowed("tag"),
-            Cow::Borrowed("pk"),
-            Cow::Borrowed("sk"),
-        ])
-    }
+        .collect()
 }
 
 #[tokio::test]