Updates to Unsealed, additional tests

Author: coderdan

cipherstash-dynamodb-derive/src/decryptable.rs

@@ -38,7 +38,7 @@ pub(crate) fn derive_decryptable(input: DeriveInput) -> Result<TokenStream, syn:
             let attr_ident = format_ident!("{attr}");
 
             quote! {
-                #attr_ident: ::cipherstash_dynamodb::traits::TryFromTableAttr::try_from_table_attr(unsealed.get_plaintext(#attr))?
+                #attr_ident: ::cipherstash_dynamodb::traits::TryFromTableAttr::try_from_table_attr(unsealed.take_unprotected(#attr))?
             }
         }))
         .chain(skipped_attributes.iter().map(|attr| {

src/crypto/attrs/flattened_protected_attributes.rs

@@ -11,7 +11,7 @@ use itertools::Itertools;
 
 // TODO: This thing is confusingly named - it holds unencrypted attributes that are intended for encryption
 /// Describes a set of flattened protected attributes intended for encryption.
-#[derive(PartialEq, Debug)]
+#[derive(PartialEq)]
 pub(crate) struct FlattenedProtectedAttributes(pub(super) Vec<FlattenedProtectedAttribute>);
 
 impl FlattenedProtectedAttributes {
@@ -68,8 +68,7 @@ impl FromIterator<(Plaintext, String)> for FlattenedProtectedAttributes {
 /// Describes a flattened protected attribute intended for encryption.
 /// It is composed of a [Plaintext] and a [FlattenedKey].
 ///
-// TODO: Only implement Debug in tests
-#[derive(PartialEq, Debug)]
+#[derive(PartialEq)]
 pub(crate) struct FlattenedProtectedAttribute {
     plaintext: Plaintext,
     key: FlattenedAttrName,
@@ -110,8 +109,7 @@ impl From<FlattenedProtectedAttribute> for BytesWithDescriptor {
 ///
 /// The key is composed of a prefix, a key, and an optional subkey.
 /// A Map would have a key and a subkey, while a scalar would only have a key.
-// TODO: Only implement Debug in tests
-#[derive(PartialEq, Hash, Eq, Clone, Debug)]
+#[derive(PartialEq, Hash, Eq, Clone)]
 pub(super) struct FlattenedAttrName {
     // TODO: Use a Cow to avoid copies during decryption
     // We may also never set this to None in which can we can remove the Option
@@ -212,6 +210,26 @@ impl From<(&str, &str)> for FlattenedAttrName {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use std::fmt::Debug;
+
+    impl Debug for FlattenedAttrName {
+        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+            f.debug_struct("FlattenedAttrName")
+                .field("prefix", &self.prefix)
+                .field("name", &self.name)
+                .field("subkey", &self.subkey)
+                .finish()
+        }
+    }
+
+    impl Debug for FlattenedProtectedAttribute {
+        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+            f.debug_struct("FlattenedProtectedAttribute")
+                .field("plaintext", &self.plaintext)
+                .field("key", &self.key)
+                .finish()
+        }
+    }
 
     #[test]
     fn test_flattened_key_from_string() {

src/crypto/attrs/normalized_protected_attributes.rs

@@ -2,7 +2,10 @@ use super::flattened_protected_attributes::{
     FlattenedAttrName, FlattenedProtectedAttribute, FlattenedProtectedAttributes,
 };
 use cipherstash_client::encryption::Plaintext;
-use std::collections::HashMap;
+use std::{
+    collections::HashMap,
+    hash::{Hash, Hasher},
+};
 
 pub(crate) struct NormalizedProtectedAttributes {
     values: HashMap<NormalizedKey, NormalizedValue>,
@@ -110,12 +113,29 @@ impl FromIterator<FlattenedProtectedAttribute> for NormalizedProtectedAttributes
 
 /// Normalized keys are effectively just strings but wrapping them in an enum allows us to
 /// differentiate between scalar and map keys without checking the value.
-#[derive(PartialEq, Debug, Hash, Eq)]
 pub(crate) enum NormalizedKey {
     Scalar(String),
     Map(String),
 }
 
+impl PartialEq for NormalizedKey {
+    fn eq(&self, other: &Self) -> bool {
+        match (self, other) {
+            (Self::Scalar(a) | Self::Map(a), Self::Scalar(b) | Self::Map(b)) => a == b,
+        }
+    }
+}
+
+impl Eq for NormalizedKey {}
+
+impl Hash for NormalizedKey {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        match self {
+            NormalizedKey::Scalar(s) | NormalizedKey::Map(s) => s.hash(state),
+        }
+    }
+}
+
 impl NormalizedKey {
     pub(super) fn new_scalar(key: impl Into<String>) -> Self {
         Self::Scalar(key.into())
@@ -140,7 +160,6 @@ impl From<NormalizedKey> for String {
     }
 }
 
-// TODO: Don't debug or only derive in tests
 #[derive(PartialEq, Debug)]
 pub(crate) enum NormalizedValue {
     Scalar(Plaintext),
@@ -198,6 +217,17 @@ impl NormalizedValue {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use std::fmt::Debug;
+
+    // Only debug in tests
+    impl Debug for NormalizedKey {
+        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+            match self {
+                NormalizedKey::Scalar(s) => write!(f, "Scalar({})", s),
+                NormalizedKey::Map(s) => write!(f, "Map({})", s),
+            }
+        }
+    }
 
     #[test]
     fn test_normalized_key() {
@@ -208,30 +238,6 @@ mod tests {
         assert_eq!(String::from(map), "map");
     }
 
-    #[test]
-    fn test_normalized_key_partial_eq() {
-        assert_eq!(
-            NormalizedKey::Scalar("a".to_string()),
-            NormalizedKey::Scalar("a".to_string())
-        );
-        assert_ne!(
-            NormalizedKey::Scalar("a".to_string()),
-            NormalizedKey::Scalar("b".to_string())
-        );
-        assert_eq!(
-            NormalizedKey::Map("a".to_string()),
-            NormalizedKey::Map("a".to_string())
-        );
-        assert_ne!(
-            NormalizedKey::Map("a".to_string()),
-            NormalizedKey::Map("b".to_string())
-        );
-        assert_ne!(
-            NormalizedKey::Scalar("a".to_string()),
-            NormalizedKey::Map("a".to_string())
-        );
-    }
-
     #[test]
     fn test_normalized_value_into_scalar() {
         let scalar = NormalizedValue::Scalar(Plaintext::from("scalar"));

src/crypto/unsealed.rs

@@ -9,7 +9,9 @@ use crate::{
 use cipherstash_client::encryption::Plaintext;
 use std::collections::HashMap;
 
-/// Wrapper to indicate that a value is NOT encrypted
+/// Wrapper to which values are added prior to being encrypted.
+/// Values added as "protected" (e.g. via [Unsealed::add_protected]) will be encrypted.
+/// Values added as "unprotected" (e.g. via [Unsealed::add_unprotected]) will not be encrypted.
 pub struct Unsealed {
     /// Protected plaintexts with their descriptors
     protected: NormalizedProtectedAttributes,
@@ -38,18 +40,20 @@ impl Unsealed {
         }
     }
 
-    // TODO: Change this to take_unprotected
+    #[deprecated(since = "0.7.3", note = "Use `Unsealed::take_unprotected` instead")]
     pub fn get_plaintext(&self, name: impl Into<AttributeName>) -> TableAttribute {
         self.unprotected
             .get(name)
             .cloned()
             .unwrap_or(TableAttribute::Null)
     }
 
+    /// Add a new protected attribute, `name`, with the given plaintext.
     pub fn add_protected(&mut self, name: impl Into<String>, plaintext: impl Into<Plaintext>) {
         self.protected.insert(name, plaintext.into());
     }
 
+    /// Add a new protected map, `name`, with the given map of plaintexts.
     pub fn add_protected_map(&mut self, name: impl Into<String>, map: HashMap<String, Plaintext>) {
         self.protected.insert_map(name, map);
     }
@@ -68,6 +72,7 @@ impl Unsealed {
             .insert_and_update_map(name, subkey, value.into());
     }
 
+    /// Add a new unprotected attribute, `name`, with the given plaintext.
     pub fn add_unprotected(
         &mut self,
         name: impl Into<AttributeName>,
@@ -76,6 +81,16 @@ impl Unsealed {
         self.unprotected.insert(name, attribute);
     }
 
+    /// Removes and returns the unprotected attribute, `name`.
+    /// See also [TableAttribute].
+    ///
+    /// If the attribute does not exist, `TableAttribute::Null` is returned.
+    pub fn take_unprotected(&mut self, name: impl Into<AttributeName>) -> TableAttribute {
+        self.unprotected
+            .remove(name)
+            .unwrap_or(TableAttribute::Null)
+    }
+
     /// Removes and returns the protected attribute, `name`.
     pub fn take_protected(&mut self, name: &str) -> Option<Plaintext> {
         self.protected.take(name)
@@ -103,63 +118,105 @@ impl Unsealed {
         unsealed
     }
 
+    /// Convert `self` into `T` using the attributes stored in `self`.
+    /// The [Decryptable] trait must be implemented for `T` and this method calls [Decryptable::from_unsealed].
     pub fn into_value<T: Decryptable>(self) -> Result<T, SealError> {
         T::from_unsealed(self)
     }
 }
 
 #[cfg(test)]
 mod tests {
-    /*use std::collections::BTreeMap;
-
     use super::*;
+    use std::collections::BTreeMap;
 
     #[test]
-    fn test_nested_protected() {
+    fn test_protected_field() {
         let mut unsealed = Unsealed::new_with_descriptor("test");
-        unsealed.add_protected("test.a", Plaintext::from("a"));
-        unsealed.add_protected("test.b", Plaintext::from("b"));
-        unsealed.add_protected("test.c", Plaintext::from("c"));
-        unsealed.add_protected("test.d", Plaintext::from("d"));
-
-        let nested = unsealed
-            .nested_protected("test")
-            .collect::<BTreeMap<_, _>>();
+        unsealed.add_protected("test", "value");
 
-        assert_eq!(nested.len(), 4);
-        assert_eq!(nested["a"], Plaintext::from("a"));
-        assert_eq!(nested["b"], Plaintext::from("b"));
-        assert_eq!(nested["c"], Plaintext::from("c"));
-        assert_eq!(nested["d"], Plaintext::from("d"));
+        let plaintext = unsealed.take_protected("test").unwrap();
+        assert_eq!(plaintext, Plaintext::from("value"));
     }
 
     #[test]
-    fn test_flatted_protected_value() {
+    fn test_protected_map() {
+        let mut unsealed = Unsealed::new_with_descriptor("test");
         let mut map = HashMap::new();
         map.insert("a".to_string(), Plaintext::from("value-a"));
         map.insert("b".to_string(), Plaintext::from("value-b"));
         map.insert("c".to_string(), Plaintext::from("value-c"));
-        map.insert("d".to_string(), Plaintext::from("value-d"));
-
-        let protected = NormalizedValue::Map(map, "test".to_string());
-        let flattened = protected.flatten();
-
-        assert_eq!(flattened.len(), 4);
-        assert!(flattened.contains(&NormalizedValue::Scalar(
-            Plaintext::from("value-a"),
-            "test.a".to_string()
-        )));
-        assert!(flattened.contains(&NormalizedValue::Scalar(
-            Plaintext::from("value-b"),
-            "test.b".to_string()
-        )));
-        assert!(flattened.contains(&NormalizedValue::Scalar(
-            Plaintext::from("value-c"),
-            "test.c".to_string()
-        )));
-        assert!(flattened.contains(&NormalizedValue::Scalar(
-            Plaintext::from("value-d"),
-            "test.d".to_string()
-        )));
-    }*/
+        unsealed.add_protected_map("test", map);
+
+        let nested: BTreeMap<String, Plaintext> = unsealed
+            .take_protected_map("test")
+            .unwrap()
+            .into_iter()
+            .collect();
+
+        assert_eq!(nested.len(), 3);
+        assert_eq!(nested["a"], Plaintext::from("value-a"));
+        assert_eq!(nested["b"], Plaintext::from("value-b"));
+        assert_eq!(nested["c"], Plaintext::from("value-c"));
+    }
+
+    #[test]
+    fn test_protected_map_field() {
+        let mut unsealed = Unsealed::new_with_descriptor("test");
+        unsealed.add_protected_map_field("test", "a", "value-a");
+        unsealed.add_protected_map_field("test", "b", "value-b");
+        unsealed.add_protected_map_field("test", "c", "value-c");
+
+        let nested: BTreeMap<String, Plaintext> = unsealed
+            .take_protected_map("test")
+            .unwrap()
+            .into_iter()
+            .collect();
+
+        assert_eq!(nested.len(), 3);
+        assert_eq!(nested["a"], Plaintext::from("value-a"));
+        assert_eq!(nested["b"], Plaintext::from("value-b"));
+        assert_eq!(nested["c"], Plaintext::from("value-c"));
+    }
+
+    #[test]
+    fn test_protected_mixed() {
+        let mut unsealed = Unsealed::new_with_descriptor("test");
+        unsealed.add_protected("test", "value");
+        unsealed.add_protected_map_field("attrs", "a", "value-a");
+        unsealed.add_protected_map_field("attrs", "b", "value-b");
+        unsealed.add_protected_map_field("attrs", "c", "value-c");
+
+        let plaintext = unsealed.take_protected("test").unwrap();
+        assert_eq!(plaintext, Plaintext::from("value"));
+
+        let nested: BTreeMap<String, Plaintext> = unsealed
+            .take_protected_map("attrs")
+            .unwrap()
+            .into_iter()
+            .collect();
+
+        assert_eq!(nested.len(), 3);
+        assert_eq!(nested["a"], Plaintext::from("value-a"));
+        assert_eq!(nested["b"], Plaintext::from("value-b"));
+        assert_eq!(nested["c"], Plaintext::from("value-c"));
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_protected_map_override() {
+        let mut unsealed = Unsealed::new_with_descriptor("test");
+        unsealed.add_protected("test", "value");
+        // Panics because "test" is already a protected scalar
+        unsealed.add_protected_map_field("test", "a", "value-a");
+    }
+
+    #[test]
+    fn test_unprotected() {
+        let mut unsealed = Unsealed::new_with_descriptor("test");
+        unsealed.add_unprotected("test", "value");
+
+        let attribute = unsealed.take_unprotected("test");
+        assert!(attribute == "value".into(), "values do not match");
+    }
 }

src/encrypted_table/table_attributes.rs

@@ -46,6 +46,10 @@ impl TableAttributes {
             .try_insert_map(subkey.into(), value.into())
     }
 
+    pub(crate) fn remove(&mut self, name: impl Into<AttributeName>) -> Option<TableAttribute> {
+        self.0.remove(&name.into())
+    }
+
     // TODO: Add unit tests for this
     /// Partition the attributes into protected and unprotected attributes
     /// given the list of protected keys.

src/traits/mod.rs

@@ -246,7 +246,7 @@ mod tests {
                     unsealed.take_protected("name"),
                 )?,
                 age: TryFromPlaintext::try_from_optional_plaintext(unsealed.take_protected("age"))?,
-                tag: TryFromTableAttr::try_from_table_attr(unsealed.get_plaintext("tag"))?,
+                tag: TryFromTableAttr::try_from_table_attr(unsealed.take_unprotected("tag"))?,
                 attrs: get_attrs(&mut unsealed)?,
             })
         }