Merge pull request #61 from cipherstash/fix/seal-all-optional-attributes

Support loading Option<T> from missing attributes

Author: Bennett Hardwick

Cargo.lock

@@ -5,7 +5,7 @@ homepage = "https://cipherstash.com"
 repository = "https://github.com/cipherstash/cipherstash-dynamodb"
 readme = "README.md"
 description = "CipherStash client for storing and querying encrypted data in DynamoDB"
-version = "0.5.0"
+version = "0.6.0"
 edition = "2021"
 
 [package.metadata.docs.rs]
@@ -24,8 +24,8 @@ repository = "https://github.com/cipherstash/cipherstash-dynamodb"
 # and it will keep the alphabetic ordering for you.
 
 [dependencies]
-cipherstash-client = { version = "0.9", registry = "cipherstash" }
-cipherstash-dynamodb-derive = { version = "0.5", path = "cipherstash-dynamodb-derive", registry = "cipherstash" }
+cipherstash-client = { version = "0.10", registry = "cipherstash" }
+cipherstash-dynamodb-derive = { version = "0.6", path = "cipherstash-dynamodb-derive", registry = "cipherstash" }
 
 aws-sdk-dynamodb = "1.3.0"
 async-trait = "0.1.73"

Cargo.toml

@@ -4,7 +4,7 @@ license-file = "../LICENSE.md"
 homepage = "https://cipherstash.com"
 readme = "../README.md"
 description = "Derive macros for the CipherStash client for DynamoDB"
-version = "0.5.0"
+version = "0.6.0"
 edition = "2021"
 
 [lib]

cipherstash-dynamodb-derive/Cargo.toml

@@ -31,14 +31,14 @@ pub(crate) fn derive_decryptable(input: DeriveInput) -> Result<TokenStream, syn:
             let attr_ident = format_ident!("{attr}");
 
             quote! {
-                #attr_ident: ::cipherstash_dynamodb::traits::TryFromPlaintext::try_from_plaintext(unsealed.get_protected(#attr)?.to_owned())?
+                #attr_ident: ::cipherstash_dynamodb::traits::TryFromPlaintext::try_from_optional_plaintext(unsealed.get_protected(#attr).cloned())?
             }
         })
         .chain(plaintext_attributes.iter().map(|attr| {
             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.get_plaintext(#attr))?
             }
         }))
         .chain(skipped_attributes.iter().map(|attr| {

cipherstash-dynamodb-derive/src/decryptable.rs

@@ -0,0 +1,58 @@
+use std::future::Future;
+
+/// Take an input vector and run a callback over only the Somes in the vec.
+///
+/// The callback must return a vector in the same order and equal in length to the Somes.
+pub async fn async_map_somes<T, U, E, F: Future<Output = Result<Vec<U>, E>>>(
+    input: Vec<Option<T>>,
+    callback: impl FnOnce(Vec<T>) -> F,
+) -> Result<Vec<Option<U>>, E> {
+    let mut output = Vec::with_capacity(input.len());
+    output.resize_with(input.len(), || None);
+
+    let (indexes, somes): (Vec<usize>, Vec<T>) = input
+        .into_iter()
+        .enumerate()
+        .filter_map(|(i, x)| x.map(|y| (i, y)))
+        .unzip();
+
+    let somes_len = somes.len();
+
+    let callback_result = callback(somes).await?;
+
+    assert_eq!(
+        callback_result.len(),
+        somes_len,
+        "expected input length to equal output length"
+    );
+
+    for (x, i) in callback_result.into_iter().zip(indexes.into_iter()) {
+        output[i] = Some(x);
+    }
+
+    Ok(output)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    #[should_panic]
+    async fn test_array_different_size() {
+        async_map_somes(vec![Some(10)], |_| async { Ok::<Vec<()>, ()>(vec![]) })
+            .await
+            .unwrap();
+    }
+
+    #[tokio::test]
+    async fn test_maintain_order() {
+        let input = vec![None, Some(1_u8), None, Some(2_u8), None, Some(3_u8)];
+
+        let output = async_map_somes(input.clone(), |x| async { Ok::<_, ()>(x) })
+            .await
+            .unwrap();
+
+        assert_eq!(input, output);
+    }
+}

src/async_map_somes.rs

@@ -1,4 +1,5 @@
 use crate::{
+    async_map_somes::async_map_somes,
     encrypted_table::{TableAttribute, TableEntry},
     traits::{ReadConversionError, WriteConversionError},
     Decryptable,
@@ -62,7 +63,8 @@ impl SealedTableEntry {
             plaintext_attributes,
         } = spec;
 
-        let mut plaintext_items: Vec<Vec<&TableAttribute>> = Vec::with_capacity(items.len());
+        let mut plaintext_items: Vec<Vec<Option<&TableAttribute>>> =
+            Vec::with_capacity(items.len());
         let mut decryptable_items = Vec::with_capacity(items.len() * protected_attributes.len());
 
         for item in items.iter() {
@@ -77,9 +79,11 @@ impl SealedTableEntry {
                         });
 
                         attribute
-                            .ok_or_else(|| SealError::MissingAttribute(name.to_string()))?
-                            .as_encrypted_record()
-                            .ok_or_else(|| SealError::InvalidCiphertext(name.to_string()))
+                            .map(|x| {
+                                x.as_encrypted_record()
+                                    .ok_or_else(|| SealError::InvalidCiphertext(name.to_string()))
+                            })
+                            .transpose()
                     })
                     .collect::<Result<Vec<_>, _>>()?;
 
@@ -97,23 +101,21 @@ impl SealedTableEntry {
                         _ => name,
                     };
 
-                    item.inner()
-                        .attributes
-                        .get(attr)
-                        .ok_or(SealError::MissingAttribute(attr.to_string()))
+                    item.inner().attributes.get(attr)
                 })
-                .collect::<Result<Vec<&TableAttribute>, SealError>>()?;
+                .collect::<Vec<Option<&TableAttribute>>>();
 
             plaintext_items.push(unprotected);
         }
 
-        let decrypted = cipher.decrypt(decryptable_items).await?;
+        let decrypted = async_map_somes(decryptable_items, |items| cipher.decrypt(items)).await?;
 
-        let decrypted_iter: &mut dyn Iterator<Item = &[Plaintext]> =
+        let decrypted_iter: &mut dyn Iterator<Item = &[Option<Plaintext>]> =
             if protected_attributes.len() > 0 {
                 &mut decrypted.chunks_exact(protected_attributes.len())
             } else {
-                &mut std::iter::repeat_with::<&[Plaintext], _>(|| &[]).take(plaintext_items.len())
+                &mut std::iter::repeat_with::<&[Option<Plaintext>], _>(|| &[])
+                    .take(plaintext_items.len())
             };
 
         let unsealed = decrypted_iter
@@ -122,13 +124,17 @@ impl SealedTableEntry {
                 let mut unsealed = Unsealed::new();
 
                 for (name, plaintext) in protected_attributes.iter().zip(decrypted_plaintext) {
-                    unsealed.add_protected(name.to_string(), plaintext.clone());
+                    if let Some(plaintext) = plaintext {
+                        unsealed.add_protected(name.to_string(), plaintext.clone());
+                    }
                 }
 
                 for (name, plaintext) in
                     plaintext_attributes.iter().zip(plaintext_items.into_iter())
                 {
-                    unsealed.add_unprotected(name.to_string(), plaintext.clone());
+                    if let Some(plaintext) = plaintext {
+                        unsealed.add_unprotected(name.to_string(), plaintext.clone());
+                    }
                 }
 
                 unsealed

src/crypto/sealed.rs

@@ -45,20 +45,17 @@ impl Unsealed {
         &self.unprotected
     }
 
-    pub fn get_protected(&self, name: &str) -> Result<&Plaintext, SealError> {
-        let (plaintext, _) = self
-            .protected
-            .get(name)
-            .ok_or_else(|| SealError::MissingAttribute(name.to_string()))?;
+    pub fn get_protected(&self, name: &str) -> Option<&Plaintext> {
+        let (plaintext, _) = self.protected.get(name)?;
 
-        Ok(plaintext)
+        Some(plaintext)
     }
 
-    pub fn get_plaintext(&self, name: &str) -> Result<TableAttribute, SealError> {
+    pub fn get_plaintext(&self, name: &str) -> TableAttribute {
         self.unprotected
             .get(name)
             .cloned()
-            .ok_or_else(|| SealError::MissingAttribute(name.to_string()))
+            .unwrap_or(TableAttribute::Null)
     }
 
     pub fn add_protected(&mut self, name: impl Into<String>, plaintext: Plaintext) {

src/crypto/unsealed.rs

@@ -593,4 +593,6 @@ pub use cipherstash_dynamodb_derive::{Decryptable, Encryptable, Identifiable, Se
 // Re-exports
 pub use cipherstash_client::encryption;
 
+mod async_map_somes;
+
 pub type Key = [u8; 32];

src/lib.rs

@@ -0,0 +1,86 @@
+use cipherstash_dynamodb::{
+    Decryptable, Encryptable, EncryptedTable, Identifiable, Pk, Searchable,
+};
+use serial_test::serial;
+use std::{borrow::Cow, future::Future};
+
+mod common;
+
+#[derive(Encryptable, Decryptable, Searchable, Debug, PartialEq, Ord, PartialOrd, Eq)]
+pub struct User {
+    #[cipherstash(query = "exact")]
+    encrypted: Option<String>,
+    #[cipherstash(plaintext)]
+    plaintext: Option<String>,
+}
+
+impl Identifiable for User {
+    type PrimaryKey = Pk;
+
+    fn get_primary_key(&self) -> Self::PrimaryKey {
+        Pk("user".into())
+    }
+
+    fn type_name() -> Cow<'static, str> {
+        "user".into()
+    }
+
+    fn sort_key_prefix() -> Option<Cow<'static, str>> {
+        None
+    }
+}
+
+#[derive(Encryptable, Decryptable, Searchable, Debug, PartialEq, Ord, PartialOrd, Eq)]
+pub struct Empty {}
+
+impl Identifiable for Empty {
+    type PrimaryKey = Pk;
+
+    fn get_primary_key(&self) -> Self::PrimaryKey {
+        Pk("user".into())
+    }
+
+    fn type_name() -> Cow<'static, str> {
+        "user".into()
+    }
+
+    fn sort_key_prefix() -> Option<Cow<'static, str>> {
+        None
+    }
+}
+
+async fn run_test<F: Future<Output = ()>>(mut f: impl FnMut(EncryptedTable) -> F) {
+    let config = aws_config::from_env()
+        .endpoint_url("http://localhost:8000")
+        .load()
+        .await;
+
+    let client = aws_sdk_dynamodb::Client::new(&config);
+
+    let table_name = "empty-record-load";
+
+    common::create_table(&client, table_name).await;
+
+    let table = EncryptedTable::init(client, table_name)
+        .await
+        .expect("Failed to init table");
+
+    table
+        .put(Empty {})
+        .await
+        .expect("Failed to insert empty record");
+
+    f(table).await;
+}
+
+#[tokio::test]
+#[serial]
+async fn test_load_from_empty() {
+    run_test(|table| async move {
+        table
+            .get::<User>(Pk("user".into()))
+            .await
+            .expect("failed to get user");
+    })
+    .await
+}