Merge pull request #82 from cipherstash/bug/divide-by-zero

Author: coderdan

Cargo.lock

@@ -38,6 +38,7 @@ serde_json = "1.0.117"
 tracing-test = "0.2.5"
 # So we can get backtraces in tests
 miette = { version = "7.2.0", features = ["fancy"] }
+chrono = "0.4.38"
 
 [features]
 default = ["tokio"]

Cargo.toml

@@ -70,7 +70,7 @@ impl SealedTableEntry {
     ///
     /// This should be used over [`Sealed::unseal`] when multiple values need to be unsealed.
     pub(crate) async fn unseal_all(
-        items: Vec<SealedTableEntry>,
+        items: Vec<Self>,
         spec: UnsealSpec<'_>,
         cipher: &Encryption<impl Credentials<Token = ServiceToken>>,
     ) -> Result<Vec<Unsealed>, SealError> {
@@ -79,6 +79,11 @@ impl SealedTableEntry {
             sort_key_prefix,
         } = spec;
 
+        let items_len = items.len();
+        if items_len == 0 {
+            return Ok(Vec::new());
+        }
+
         let mut protected_items = {
             let capacity = items.len() * protected_attributes.len();
             FlattenedEncryptedAttributes::with_capacity(capacity)
@@ -98,16 +103,16 @@ impl SealedTableEntry {
         if protected_items.is_empty() {
             unprotected_items
                 .into_iter()
-                .map(|unprotected| {
-                    // TODO: Create a new_from_unprotected method
-                    Ok(Unsealed::new_from_parts(
-                        NormalizedProtectedAttributes::new(),
-                        unprotected,
-                    ))
-                })
+                .map(|unprotected| Ok(Unsealed::new_from_unprotected(unprotected)))
                 .collect()
         } else {
-            let chunk_size = protected_items.len() / unprotected_items.len();
+            let chunk_size =
+                protected_items
+                    .len()
+                    .checked_div(items_len)
+                    .ok_or(SealError::AssertionFailed(
+                        "Division by zero when calculating chunk size".to_string(),
+                    ))?;
 
             protected_items
                 .decrypt_all(cipher)
@@ -199,3 +204,52 @@ impl TryFrom<SealedTableEntry> for HashMap<String, AttributeValue> {
         Ok(map)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::SealedTableEntry;
+    use cipherstash_client::{
+        credentials::{auto_refresh::AutoRefresh, service_credentials::ServiceCredentials},
+        encryption::Encryption,
+        ConsoleConfig, ZeroKMS, ZeroKMSConfig,
+    };
+    use miette::IntoDiagnostic;
+    use std::borrow::Cow;
+
+    type Cipher = Encryption<AutoRefresh<ServiceCredentials>>;
+
+    // FIXME: Use the test cipher from CipherStash Client when that's ready
+    async fn get_cipher() -> Result<Cipher, Box<dyn std::error::Error>> {
+        let console_config = ConsoleConfig::builder().with_env().build()?;
+        let zero_kms_config = ZeroKMSConfig::builder()
+            .decryption_log(true)
+            .with_env()
+            .console_config(&console_config)
+            .build_with_client_key()?;
+
+        let zero_kms_client = ZeroKMS::new_with_client_key(
+            &zero_kms_config.base_url(),
+            AutoRefresh::new(zero_kms_config.credentials()),
+            zero_kms_config.decryption_log_path().as_deref(),
+            zero_kms_config.client_key(),
+        );
+
+        let config = zero_kms_client.load_dataset_config().await?;
+        Ok(Encryption::new(config.index_root_key, zero_kms_client))
+    }
+
+    #[tokio::test]
+    async fn test_unseal_all_empty() -> Result<(), Box<dyn std::error::Error>> {
+        let spec = super::UnsealSpec {
+            protected_attributes: Cow::Borrowed(&[]),
+            sort_key_prefix: "test".to_string(),
+        };
+        let cipher = get_cipher().await?;
+        let results = SealedTableEntry::unseal_all(vec![], spec, &cipher)
+            .await
+            .into_diagnostic()?;
+        assert!(results.is_empty());
+
+        Ok(())
+    }
+}

src/crypto/sealed.rs

@@ -40,6 +40,24 @@ impl Unsealed {
         }
     }
 
+    /// Create a new Unsealed from the protected and unprotected attributes.
+    pub(crate) fn new_from_parts(
+        protected: NormalizedProtectedAttributes,
+        unprotected: TableAttributes,
+    ) -> Self {
+        let mut unsealed = Self::new();
+        unsealed.protected = protected;
+        unsealed.unprotected = unprotected;
+        unsealed
+    }
+
+    /// Create a new Unsealed from the protected and unprotected attributes.
+    pub(crate) fn new_from_unprotected(unprotected: TableAttributes) -> Self {
+        let mut unsealed = Self::new();
+        unsealed.unprotected = unprotected;
+        unsealed
+    }
+
     #[deprecated(since = "0.7.3", note = "Use `Unsealed::take_unprotected` instead")]
     pub fn get_plaintext(&self, name: impl Into<AttributeName>) -> TableAttribute {
         self.unprotected
@@ -107,17 +125,6 @@ impl Unsealed {
         (self.protected.flatten(), self.unprotected)
     }
 
-    /// Create a new Unsealed from the protected and unprotected attributes.
-    pub(crate) fn new_from_parts(
-        protected: NormalizedProtectedAttributes,
-        unprotected: TableAttributes,
-    ) -> Self {
-        let mut unsealed = Self::new();
-        unsealed.protected = protected;
-        unsealed.unprotected = unprotected;
-        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> {

src/crypto/unsealed.rs

@@ -12,7 +12,7 @@ pub trait TryFromTableAttr: Sized {
     fn try_from_table_attr(value: TableAttribute) -> Result<Self, ReadConversionError>;
 }
 
-#[derive(Clone, PartialEq)]
+#[derive(Clone, PartialEq, Debug)]
 pub enum TableAttribute {
     String(String),
     Number(String),

src/encrypted_table/table_attribute.rs

@@ -4,7 +4,7 @@ use std::{
     collections::{hash_map::IntoIter, HashMap},
 };
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 /// Represents a collection of attributes for a table entry.
 /// Attributes are stored as a map of `String` to `TableAttribute`.
 pub struct TableAttributes(HashMap<AttributeName, TableAttribute>);
@@ -28,7 +28,8 @@ impl TableAttributes {
         value: impl Into<TableAttribute>,
     ) {
         let name: AttributeName = name.into();
-        self.0.insert(name, value.into());
+        let attr: TableAttribute = value.into();
+        self.0.insert(name, attr);
     }
 
     /// Attempts to insert a value into a map with key `subkey` where the map is stored at `key`.
@@ -57,7 +58,6 @@ impl TableAttributes {
         let (protected, unprotected): (HashMap<_, _>, HashMap<_, _>) =
             self.0.into_iter().partition(|(k, _)| {
                 let check = k.as_external_name();
-                //protected_keys.iter().any(|key| match_key(check, key))
                 protected_keys.iter().any(|key| check == key)
             });
 

src/encrypted_table/table_attributes.rs

@@ -0,0 +1,122 @@
+use chrono::NaiveDate;
+use cipherstash_dynamodb::{Decryptable, Encryptable, EncryptedTable, Identifiable, Searchable};
+use serial_test::serial;
+use std::future::Future;
+
+mod common;
+
+#[derive(
+    Identifiable, Encryptable, Decryptable, Searchable, Debug, PartialEq, Ord, PartialOrd, Eq,
+)]
+#[cipherstash(sort_key_prefix = "user")]
+pub struct User {
+    #[cipherstash(query = "exact", compound = "email#name")]
+    #[cipherstash(query = "exact")]
+    #[partition_key]
+    pub email: String,
+
+    #[cipherstash(query = "prefix", compound = "email#name")]
+    #[cipherstash(query = "prefix")]
+    pub name: String,
+
+    pub dob: NaiveDate,
+}
+
+impl User {
+    pub fn new(email: impl Into<String>, name: impl Into<String>, dob: NaiveDate) -> Self {
+        Self {
+            name: name.into(),
+            email: email.into(),
+            dob,
+        }
+    }
+}
+
+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 = "test-users-plaintext-email";
+
+    common::create_table(&client, table_name).await;
+    let table = EncryptedTable::init(client, table_name)
+        .await
+        .expect("Failed to init table");
+
+    table
+        .put(User::new(
+            "[email protected]",
+            "Dan Draper",
+            NaiveDate::from_ymd_opt(2000, 1, 10).unwrap(),
+        ))
+        .await
+        .expect("Failed to insert Dan");
+
+    table
+        .put(User::new(
+            "[email protected]",
+            "Jane Smith",
+            NaiveDate::from_ymd_opt(1990, 2, 20).unwrap(),
+        ))
+        .await
+        .expect("Failed to insert Jane");
+
+    table
+        .put(User::new(
+            "[email protected]",
+            "Daniel Johnson",
+            NaiveDate::from_ymd_opt(1980, 3, 30).unwrap(),
+        ))
+        .await
+        .expect("Failed to insert Daniel");
+
+    f(table).await;
+}
+
+#[tokio::test]
+#[serial]
+async fn test_get() {
+    run_test(|table| async move {
+        let res: Option<User> = table
+            .get("[email protected]")
+            .await
+            .expect("Failed to get user");
+
+        assert_eq!(
+            res,
+            Some(User::new(
+                "[email protected]",
+                "Dan Draper",
+                NaiveDate::from_ymd_opt(2000, 1, 10).unwrap()
+            ))
+        );
+    })
+    .await;
+}
+
+#[tokio::test]
+#[serial]
+async fn test_query_single_exact() {
+    run_test(|table| async move {
+        let res: Vec<User> = table
+            .query()
+            .eq("email", "[email protected]")
+            .send()
+            .await
+            .expect("Failed to query");
+
+        assert_eq!(
+            res,
+            vec![User::new(
+                "[email protected]",
+                "Dan Draper",
+                NaiveDate::from_ymd_opt(2000, 1, 10).unwrap()
+            )]
+        );
+    })
+    .await;
+}