All tests finally pass

Author: coderdan

src/crypto/attrs/flattened_encrypted_attributes.rs

@@ -19,6 +19,11 @@ impl FlattenedEncryptedAttributes {
         self.0.is_empty()
     }
 
+    // TODO: REmove this
+    pub(crate) fn len(&self) -> usize {
+        self.0.len()
+    }
+
     // TODO: Test this
     /// Decrypt self, returning a [FlattenedProtectedAttributes].
     pub(crate) async fn decrypt_all(

src/crypto/attrs/flattened_protected_attributes.rs

@@ -26,22 +26,20 @@ impl FlattenedProtectedAttributes {
         self.0.into_iter()
     }
 
+    // TODO: Do some more testing with the chunking
     /// Encrypt all attributes in the set and return a list of [FlattenedEncryptedAttributes] objects.
-    /// Each [FlattenedEncryptedAttributes] object contains `chunk_size` encrypted attributes.
+    /// The output is a vec of `chunk_into` [FlattenedEncryptedAttributes] objects.
     pub(crate) async fn encrypt_all(
         self,
         cipher: &Encryption<impl Credentials<Token = ServiceToken>>,
-        chunk_size: usize,
+        chunk_into: usize,
     ) -> Result<Vec<FlattenedEncryptedAttributes>, SealError> {
-        println!("Encrypting all attributes, chunk size: {}", chunk_size);
+        let chunk_size = self.0.len() / chunk_into;
+        println!("Encrypting all attributes, chunk into: {}, chunk_size = {}", chunk_into, chunk_size);
 
-        let x = cipher
+        cipher
             .encrypt(self.0.into_iter())
-            .await?;
-
-        dbg!(&x);
-
-        x
+            .await?
             .into_iter()
             .chunks(chunk_size)
             .into_iter()

src/crypto/sealed.rs

@@ -105,11 +105,15 @@ impl SealedTableEntry {
                 })
                 .collect()
         } else {
+            let chunk_size = protected_items.len() / unprotected_items.len();
+
             let r1 = protected_items
                 .decrypt_all(cipher)
                 .await?
                 .into_iter()
-                .chunks(protected_attributes.len())
+                // FIXME: chunk_size is not the same as protected_attributes.len() when dealing with maps
+                // TODO: Can we make decrypt_all return a Vec of FlattenedProtectedAttributes? (like the mirror of encrypt_all)
+                .chunks(chunk_size)
                 .into_iter()
                 .map(|fpa| fpa.into_iter().collect::<NormalizedProtectedAttributes>())
                 .collect_vec();

src/crypto/sealer.rs

@@ -93,7 +93,13 @@ impl<'a> RecordsWithTerms<'a> {
                 .collect()
 
         } else {
-            let encrypted = protected.encrypt_all(cipher, self.num_protected_attributes).await?;
+            println!("Protecteds before encryption: {:?}", protected);
+            let encrypted = protected.encrypt_all(cipher, num_records).await?;
+
+            println!("Encrypted len = {}", encrypted.len());
+            println!("Unprotected len = {}", unprotecteds.len());
+            println!("Encrypted: {:?}", encrypted);
+            println!("Unprotected: {:?}", unprotecteds);
 
             encrypted
                 .into_iter()
@@ -150,9 +156,6 @@ impl Sealer {
         let protected_attributes = protected_attributes.as_ref();
         let num_protected_attributes = protected_attributes.len();
 
-        dbg!(&protected_attributes);
-        dbg!(num_protected_attributes);
-
         records
             .into_iter()
             .map(|sealer| {

src/encrypted_table/mod.rs

@@ -366,8 +366,6 @@ impl<D> EncryptedTable<D> {
     ) -> Result<DynamoRecordPatch, PutError> {
         let mut seen_sk = HashSet::new();
 
-        dbg!(&record);
-
         let PreparedRecord {
             protected_attributes,
             protected_indexes,

src/encrypted_table/table_attributes.rs

@@ -5,20 +5,6 @@ use std::{borrow::Cow, collections::HashMap};
 #[derive(Debug, Clone)]
 /// Represents a collection of attributes for a table entry.
 /// Attributes are stored as a map of `String` to `TableAttribute`.
-///
-/// ## Namespacing
-///
-/// Attributes can be namespaced by using a dot (`.`) in the attribute name.
-/// For example:
-///
-/// ```rust
-/// use cipherstash_dynamodb::TableAttributes;
-/// use std::collections::HashMap;
-///
-/// // TODO: Add a test for this
-/// ```
-///
-/// TODO: Explain how namespacing works with normalizing and denormalizing
 pub struct TableAttributes(HashMap<String, TableAttribute>);
 
 impl TableAttributes {

tests/nested_tests.rs

@@ -2,8 +2,7 @@ mod common;
 
 // TODO: Use the derive macros for this test
 use std::{borrow::Cow, collections::BTreeMap};
-use tracing_test::traced_test;
-
+use miette::IntoDiagnostic;
 use cipherstash_client::encryption::TypeParseError;
 use cipherstash_dynamodb::{
     crypto::Unsealed,
@@ -55,10 +54,9 @@ impl Identifiable for Test {
     }
 }
 
-// TODO: Make this function consume and return the Unsealed
 fn put_attrs(unsealed: &mut Unsealed, attrs: BTreeMap<String, String>) {
     attrs.into_iter().for_each(|(k, v)| {
-        unsealed.add_protected(format!("attrs.{k}"), Plaintext::from(v));
+        unsealed.add_protected_map_field("attrs", k, Plaintext::from(v));
     })
 }
 
@@ -76,9 +74,8 @@ impl Encryptable for Test {
     }
 
     fn into_unsealed(self) -> Unsealed {
-        // FIXME: This should be a "consuming" method
         let mut unsealed = Unsealed::new_with_descriptor(<Self as Identifiable>::type_name());
-        unsealed.add_protected("pk", Plaintext::from(self.pk));
+        unsealed.add_unprotected("pk", TableAttribute::from(self.pk));
         unsealed.add_unprotected("sk", TableAttribute::from(self.sk));
         unsealed.add_protected("name", Plaintext::from(self.name));
         unsealed.add_protected("age", Plaintext::from(self.age));
@@ -94,7 +91,6 @@ where
 {
     unsealed
         .take_protected_map("attrs")
-        // FIXME: This method should have a better error
         .ok_or(TypeParseError("attrs".to_string()))?
         .into_iter()
         .map(|(k, v)| TryFromPlaintext::try_from_plaintext(v).map(|v| (k, v)))
@@ -103,26 +99,23 @@ where
 
 impl Decryptable for Test {
     fn from_unsealed(mut unsealed: Unsealed) -> Result<Self, SealError> {
-        println!("IN FROM UNSEALED");
         Ok(Self {
-            /*pk: TryFromTableAttr::try_from_table_attr(
+            pk: TryFromTableAttr::try_from_table_attr(
                 unsealed.get_plaintext("pk"),
             )?,
             sk: TryFromTableAttr::try_from_table_attr(
                 unsealed.get_plaintext("sk"),
-            )?,*/
-            pk: String::from("pk-hack"),
-            sk: String::from("sk-hack"),
-            name: dbg!(TryFromPlaintext::try_from_optional_plaintext(
-                dbg!(unsealed.take_protected("name")),
-            ))?,
-            age: dbg!(TryFromPlaintext::try_from_optional_plaintext(
+            )?,
+            name: TryFromPlaintext::try_from_optional_plaintext(
+                unsealed.take_protected("name"),
+            )?,
+            age: TryFromPlaintext::try_from_optional_plaintext(
                 unsealed.take_protected("age"),
-            ))?,
-            tag: dbg!(TryFromTableAttr::try_from_table_attr(
+            )?,
+            tag: TryFromTableAttr::try_from_table_attr(
                 unsealed.get_plaintext("tag"),
-            ))?,
-            attrs: dbg!(get_attrs(&mut unsealed))?,
+            )?,
+            attrs: get_attrs(&mut unsealed)?,
         })
     }
 
@@ -140,8 +133,7 @@ impl Decryptable for Test {
 }
 
 #[tokio::test]
-#[traced_test]
-async fn test_round_trip() {
+async fn test_round_trip() -> Result<(), Box<dyn std::error::Error>> {
     let config = aws_config::from_env()
         .endpoint_url("http://localhost:8000")
         .load()
@@ -154,7 +146,7 @@ async fn test_round_trip() {
 
     let table = EncryptedTable::init(client, table_name)
         .await
-        .expect("Failed to init table");
+        .into_diagnostic()?;
 
     let record = Test {
         pk: "pk".to_string(),
@@ -168,17 +160,14 @@ async fn test_round_trip() {
     table
         .put(record.clone())
         .await
-        .expect("Failed to insert record");
+        .into_diagnostic()?;
 
-    /*let check = table
+    let check = table
         .get::<Test>(("pk", "sk"))
-        .await;
-
-    if let Err(e) = check {
-        panic!("Failed to get record: {:?}", e);
-    }*/
+        .await
+        .into_diagnostic()?;
 
-    assert!(false);
+    assert_eq!(check, Some(record));
 
-    //assert_eq!(check, record);
+    Ok(())
 }