Updated methods on encrypted table for dataset awareness

Author: coderdan

docker-compose.yaml

@@ -1,5 +1,3 @@
-version: '3'
-
 services:
   dynamodb:
     build: ./local-dynamodb

src/crypto/attrs/flattened_encrypted_attributes.rs

@@ -1,6 +1,6 @@
 use crate::{
     crypto::{attrs::flattened_protected_attributes::FlattenedAttrName, SealError},
-    encrypted_table::{ScopedZeroKmsCipher, TableAttributes},
+    encrypted_table::{TableAttributes, ZeroKmsCipher},
     traits::TableAttribute,
 };
 use cipherstash_client::{
@@ -34,7 +34,7 @@ impl FlattenedEncryptedAttributes {
     /// Decrypt self, returning a [FlattenedProtectedAttributes].
     pub(crate) async fn decrypt_all(
         self,
-        cipher: &ScopedZeroKmsCipher,
+        cipher: &ZeroKmsCipher,
     ) -> Result<FlattenedProtectedAttributes, SealError> {
         let descriptors = self
             .attrs

src/crypto/sealed.rs

@@ -1,6 +1,6 @@
 use crate::{
     crypto::attrs::FlattenedEncryptedAttributes,
-    encrypted_table::{ScopedZeroKmsCipher, TableEntry},
+    encrypted_table::{TableEntry, ZeroKmsCipher},
     traits::{ReadConversionError, WriteConversionError},
     Decryptable, Identifiable,
 };
@@ -68,7 +68,7 @@ impl SealedTableEntry {
     pub(crate) async fn unseal_all(
         items: Vec<Self>,
         spec: UnsealSpec<'_>,
-        cipher: &ScopedZeroKmsCipher,
+        cipher: &ZeroKmsCipher,
     ) -> Result<Vec<Unsealed>, SealError> {
         let UnsealSpec {
             protected_attributes,
@@ -130,7 +130,7 @@ impl SealedTableEntry {
     pub(crate) async fn unseal(
         self,
         spec: UnsealSpec<'_>,
-        cipher: &ScopedZeroKmsCipher,
+        cipher: &ZeroKmsCipher,
     ) -> Result<Unsealed, SealError> {
         let mut vec = Self::unseal_all(vec![self], spec, cipher).await?;
 
@@ -203,14 +203,13 @@ impl TryFrom<SealedTableEntry> for HashMap<String, AttributeValue> {
 
 #[cfg(test)]
 mod tests {
-    use crate::encrypted_table::{ZeroKmsCipher, ScopedZeroKmsCipher};
+    use crate::encrypted_table::{ZeroKmsCipher};
 
     use super::SealedTableEntry;
     use cipherstash_client::{
         credentials::auto_refresh::AutoRefresh, ConsoleConfig, ZeroKMS, ZeroKMSConfig
     };
     use miette::IntoDiagnostic;
-    use uuid::Uuid;
     use std::{borrow::Cow, sync::Arc};
 
     // FIXME: Use the test cipher from CipherStash Client when that's ready
@@ -240,11 +239,7 @@ mod tests {
             sort_key_prefix: "test".to_string(),
         };
         let cipher = get_cipher().await?;
-        // TODO: Temporary obvs
-        let dataset_id = Uuid::parse_str("93e10481-2692-4d65-a619-37e36a496e64").unwrap();
-        let scoped_cipher = ScopedZeroKmsCipher::init(cipher, dataset_id).await.unwrap();
-
-        let results = SealedTableEntry::unseal_all(vec![], spec, &scoped_cipher)
+        let results = SealedTableEntry::unseal_all(vec![], spec, &cipher)
             .await
             .into_diagnostic()?;
 

src/encrypted_table/mod.rs

@@ -279,20 +279,14 @@ impl<D> EncryptedTable<D> {
     ) -> Result<Vec<T>, DecryptError>
     where T: Decryptable + Identifiable,
     {
-        // TODO: Decryption _may_ not need to be scoped
-         // TODO: Temporary obvs
-         let dataset_id = Uuid::parse_str("93e10481-2692-4d65-a619-37e36a496e64").unwrap();
-         let scoped_cipher = ScopedZeroKmsCipher::init(self.cipher.clone(), dataset_id).await.unwrap();
-
-        decrypt_all(&scoped_cipher, items).await
+        decrypt_all(&self.cipher, items).await
     }
 
     pub async fn create_delete_patch(
         &self,
         delete: PreparedDelete,
+        dataset_id: Option<Uuid>,
     ) -> Result<DynamoRecordPatch, DeleteError> {
-        // TODO: Temporary obvs
-        let dataset_id = Uuid::parse_str("93e10481-2692-4d65-a619-37e36a496e64").unwrap();
         let scoped_cipher = ScopedZeroKmsCipher::init(self.cipher.clone(), dataset_id).await.unwrap();
 
         let PrimaryKeyParts { pk, sk } = encrypt_primary_key_parts(&scoped_cipher, delete.primary_key)?;
@@ -324,14 +318,12 @@ impl<D> EncryptedTable<D> {
     pub async fn create_put_patch(
         &self,
         record: PreparedRecord,
-        dataset_id: Uuid,
+        dataset_id: Option<Uuid>,
         // TODO: Make sure the index_predicate is used correctly
         index_predicate: impl FnMut(&AttributeName, &TableAttribute) -> bool,
     ) -> Result<DynamoRecordPatch, PutError> {
         let mut seen_sk = HashSet::new();
 
-        // TODO: Temporary obvs
-        let dataset_id = Uuid::parse_str("93e10481-2692-4d65-a619-37e36a496e64").unwrap();
         let indexable_cipher = ScopedZeroKmsCipher::init(self.cipher.clone(), dataset_id).await.unwrap();
 
         let PreparedRecord {
@@ -405,13 +397,26 @@ impl EncryptedTable<Dynamo> {
     where
         T: Decryptable + Identifiable,
     {
-        // TODO: Temporary obvs
-        let dataset_id = Uuid::parse_str("93e10481-2692-4d65-a619-37e36a496e64").unwrap();
-        let scoped_cipher = ScopedZeroKmsCipher::init(self.cipher.clone(), dataset_id).await.unwrap();
+        // TODO: Don't unwrap
+        let scoped_cipher = ScopedZeroKmsCipher::init(self.cipher.clone(), None).await.unwrap();
+        self.get_inner(k, scoped_cipher).await
+    }
 
-        let PrimaryKeyParts { pk, sk } =
-            encrypt_primary_key_parts(&scoped_cipher, PreparedPrimaryKey::new::<T>(k))?;
+    pub async fn get_via<T>(&self, k: impl Into<T::PrimaryKey>, dataset_id: Uuid) -> Result<Option<T>, GetError>
+    where
+        T: Decryptable + Identifiable,
+    {
+        // TODO: Don't unwrap
+        let scoped_cipher = ScopedZeroKmsCipher::init(self.cipher.clone(), Some(dataset_id)).await.unwrap();
+        self.get_inner(k, scoped_cipher).await
+    }
 
+    async fn get_inner<T>(&self, k: impl Into<T::PrimaryKey>, cipher: ScopedZeroKmsCipher) -> Result<Option<T>, GetError>
+    where
+        T: Decryptable + Identifiable,
+    {
+        let PrimaryKeyParts { pk, sk } =
+            encrypt_primary_key_parts(&cipher, PreparedPrimaryKey::new::<T>(k))?;
 
         let result = self
             .db
@@ -423,10 +428,8 @@ impl EncryptedTable<Dynamo> {
             .await
             .map_err(|e| GetError::Aws(format!("{e:?}")))?;
 
-        println!("RESULT {:?}", result);
-
         if let Some(item) = result.item {
-            Ok(Some(decrypt(&scoped_cipher, item).await?))
+            Ok(Some(decrypt(&self.cipher, item).await?))
         } else {
             Ok(None)
         }
@@ -435,9 +438,25 @@ impl EncryptedTable<Dynamo> {
     pub async fn delete<E: Searchable + Identifiable>(
         &self,
         k: impl Into<E::PrimaryKey>,
+    ) -> Result<(), DeleteError> {
+        self.delete_inner::<E>(k.into(), None).await
+    }
+
+    pub async fn delete_via<E: Searchable + Identifiable>(
+        &self,
+        k: impl Into<E::PrimaryKey>,
+        dataset_id: Uuid,
+    ) -> Result<(), DeleteError> {
+        self.delete_inner::<E>(k.into(), Some(dataset_id)).await
+    }
+
+    async fn delete_inner<E: Searchable + Identifiable>(
+        &self,
+        k: E::PrimaryKey,
+        dataset_id: Option<Uuid>,
     ) -> Result<(), DeleteError> {
         let transact_items = self
-            .create_delete_patch(PreparedDelete::new::<E>(k))
+            .create_delete_patch(PreparedDelete::new::<E>(k), dataset_id)
             .await?
             .into_transact_write_items(&self.db.table_name)?;
 
@@ -455,6 +474,20 @@ impl EncryptedTable<Dynamo> {
     }
 
     pub async fn put<T>(&self, record: T) -> Result<(), PutError>
+    where
+        T: Searchable + Identifiable,
+    {
+        self.put_inner(record, None).await
+    }
+
+    pub async fn put_via<T>(&self, record: T, dataset_id: Uuid) -> Result<(), PutError>
+    where
+        T: Searchable + Identifiable,
+    {
+        self.put_inner(record, Some(dataset_id)).await
+    }
+
+    async fn put_inner<T>(&self, record: T, dataset_id: Option<Uuid>) -> Result<(), PutError>
     where
         T: Searchable + Identifiable,
     {
@@ -463,6 +496,7 @@ impl EncryptedTable<Dynamo> {
         let transact_items = self
             .create_put_patch(
                 record,
+                dataset_id,
                 // include all records in the indexes
                 |_, _| true,
             )
@@ -502,7 +536,7 @@ fn encrypt_primary_key_parts(
     Ok(PrimaryKeyParts { pk, sk })
 }
 
-async fn decrypt<T>(scoped_cipher: &ScopedZeroKmsCipher, item: HashMap<String, AttributeValue>) -> Result<T, DecryptError>
+async fn decrypt<T>(scoped_cipher: &ZeroKmsCipher, item: HashMap<String, AttributeValue>) -> Result<T, DecryptError>
 where
     T: Decryptable + Identifiable,
 {
@@ -514,7 +548,7 @@ where
 }
 
 async fn decrypt_all<T>(
-    scoped_cipher: &ScopedZeroKmsCipher,
+    scoped_cipher: &ZeroKmsCipher,
     items: impl IntoIterator<Item = HashMap<String, AttributeValue>>,
 ) -> Result<Vec<T>, DecryptError>
 where

src/encrypted_table/query.rs

@@ -6,7 +6,6 @@ use cipherstash_client::{
     },
 };
 use itertools::Itertools;
-use uuid::Uuid;
 use std::{borrow::Cow, collections::HashMap, marker::PhantomData};
 
 use crate::{
@@ -128,19 +127,18 @@ impl<S> QueryBuilder<S, &EncryptedTable<Dynamo>>
 where
     S: Searchable + Identifiable,
 {
+    // TODO: Add load_via
     pub async fn load<T>(self) -> Result<Vec<T>, QueryError>
     where
         T: Decryptable + Identifiable,
     {
-        // TODO: Temporary obvs
-        let dataset_id = Uuid::parse_str("93e10481-2692-4d65-a619-37e36a496e64").unwrap();
-        let scoped_cipher = ScopedZeroKmsCipher::init(self.storage.cipher.clone(), dataset_id).await.unwrap();
+        let scoped_cipher = ScopedZeroKmsCipher::init(self.storage.cipher.clone(), None).await.unwrap();
 
         let storage = self.storage;
         let query = self.build()?;
 
         let items = query.send(storage, &scoped_cipher).await?;
-        let results = super::decrypt_all(&scoped_cipher, items).await?;
+        let results = super::decrypt_all(&storage.cipher, items).await?;
 
         Ok(results)
     }

tests/headless_tests.rs

@@ -65,7 +65,7 @@ async fn test_headless_roundtrip() {
             PreparedRecord::prepare_record(user.clone()).expect("failed to prepare record");
 
         let patch = table
-            .create_put_patch(user_record, |_, _| true)
+            .create_put_patch(user_record, None, |_, _| true)
             .await
             .expect("failed to encrypt");
 

tests/query_builder_direct.rs

@@ -3,9 +3,7 @@ use cipherstash_dynamodb::{
 };
 use itertools::Itertools;
 use serial_test::serial;
-use uuid::Uuid;
 use std::future::Future;
-
 mod common;
 
 #[derive(
@@ -89,8 +87,7 @@ async fn test_query_single_exact() {
             .build()
             .expect("failed to build query");
 
-        let dataset_id = Uuid::parse_str("93e10481-2692-4d65-a619-37e36a496e64").unwrap();
-        let scoped_cipher = ScopedZeroKmsCipher::init(table.cipher(), dataset_id).await.unwrap();
+        let scoped_cipher = ScopedZeroKmsCipher::init(table.cipher(), None).await.unwrap();
 
         let term = query
             .encrypt(&scoped_cipher)
@@ -132,8 +129,7 @@ async fn test_query_single_prefix() {
             .await
             .expect("failed to init table");
 
-        let dataset_id = Uuid::parse_str("93e10481-2692-4d65-a619-37e36a496e64").unwrap();
-        let scoped_cipher = ScopedZeroKmsCipher::init(table.cipher(), dataset_id).await.unwrap();
+        let scoped_cipher = ScopedZeroKmsCipher::init(table.cipher(), None).await.unwrap();
 
         let query = QueryBuilder::<User>::new()
             .starts_with("name", "Dan")
@@ -189,8 +185,7 @@ async fn test_query_compound() {
             .build()
             .expect("failed to build query");
 
-        let dataset_id = Uuid::parse_str("93e10481-2692-4d65-a619-37e36a496e64").unwrap();
-        let scoped_cipher = ScopedZeroKmsCipher::init(table.cipher(), dataset_id).await.unwrap();
+        let scoped_cipher = ScopedZeroKmsCipher::init(table.cipher(), None).await.unwrap();
 
         let term = query
             .encrypt(&scoped_cipher)