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::{ScopedCipherWithCreds, 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: &ScopedCipherWithCreds,
+        cipher: &ZeroKmsCipher,
     ) -> Result<FlattenedProtectedAttributes, SealError> {
         let descriptors = self
             .attrs

src/crypto/attrs/flattened_protected_attributes.rs

@@ -2,7 +2,7 @@ use super::{
     flattened_encrypted_attributes::FlattenedEncryptedAttributes,
     normalized_protected_attributes::NormalizedKey,
 };
-use crate::{crypto::SealError, encrypted_table::{AttributeName, ScopedCipherWithCreds}};
+use crate::{crypto::SealError, encrypted_table::{AttributeName, ScopedZeroKmsCipher}};
 use cipherstash_client::{
     encryption::{BytesWithDescriptor, Plaintext}, zerokms::EncryptPayload,
 };
@@ -31,7 +31,7 @@ impl FlattenedProtectedAttributes {
     /// The output is a vec of `chunk_into` [FlattenedEncryptedAttributes] objects.
     pub(crate) async fn encrypt_all(
         self,
-        cipher: &ScopedCipherWithCreds,
+        cipher: &ScopedZeroKmsCipher,
         chunk_into: usize,
     ) -> Result<Vec<FlattenedEncryptedAttributes>, SealError> {
         let chunk_size = self.0.len() / chunk_into;

src/crypto/sealed.rs

@@ -1,6 +1,6 @@
 use crate::{
     crypto::attrs::FlattenedEncryptedAttributes,
-    encrypted_table::{ScopedCipherWithCreds, 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: &ScopedCipherWithCreds,
+        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: &ScopedCipherWithCreds,
+        cipher: &ZeroKmsCipher,
     ) -> Result<Unsealed, SealError> {
         let mut vec = Self::unseal_all(vec![self], spec, cipher).await?;
 
@@ -203,18 +203,17 @@ impl TryFrom<SealedTableEntry> for HashMap<String, AttributeValue> {
 
 #[cfg(test)]
 mod tests {
-    use crate::encrypted_table::{Cipher, ScopedCipherWithCreds};
+    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
-    async fn get_cipher() -> Result<Arc<Cipher>, Box<dyn std::error::Error>> {
+    async fn get_cipher() -> Result<Arc<ZeroKmsCipher>, Box<dyn std::error::Error>> {
         let console_config = ConsoleConfig::builder().with_env().build().into_diagnostic()?;
         let zero_kms_config = ZeroKMSConfig::builder()
             .decryption_log(true)
@@ -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 = ScopedCipherWithCreds::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/crypto/sealer.rs

@@ -2,7 +2,7 @@ use super::{
     attrs::FlattenedProtectedAttributes, b64_encode, format_term_key, SealError, SealedTableEntry, Unsealed, MAX_TERMS_PER_INDEX
 };
 use crate::{
-    encrypted_table::{AttributeName, ScopedCipherWithCreds, TableAttribute, TableAttributes, TableEntry},
+    encrypted_table::{AttributeName, ScopedZeroKmsCipher, TableAttribute, TableAttributes, TableEntry},
     traits::PrimaryKeyParts,
     IndexType,
 };
@@ -53,7 +53,7 @@ impl RecordsWithTerms {
 
     async fn encrypt(
         self,
-        cipher: &ScopedCipherWithCreds,
+        cipher: &ScopedZeroKmsCipher,
     ) -> Result<Vec<Sealed>, SealError> {
         let num_records = self.records.len();
         let mut pksks = Vec::with_capacity(num_records);
@@ -134,7 +134,7 @@ impl Sealer {
     fn index_all_terms<'a>(
         records: impl IntoIterator<Item = Sealer>,
         protected_attributes: impl AsRef<[Cow<'a, str>]>,
-        cipher: &ScopedCipherWithCreds,
+        cipher: &ScopedZeroKmsCipher,
         // FIXME: This might need to be a const generic
         term_length: usize,
     ) -> Result<RecordsWithTerms, SealError> {
@@ -210,7 +210,7 @@ impl Sealer {
     pub(crate) async fn seal_all<'a>(
         records: impl IntoIterator<Item = Sealer>,
         protected_attributes: impl AsRef<[Cow<'a, str>]>,
-        cipher: &ScopedCipherWithCreds,
+        cipher: &ScopedZeroKmsCipher,
         term_length: usize,
     ) -> Result<Vec<Sealed>, SealError> {
         Self::index_all_terms(records, protected_attributes, &cipher, term_length)?
@@ -221,7 +221,7 @@ impl Sealer {
     pub(crate) async fn seal<'a>(
         self,
         protected_attributes: impl AsRef<[Cow<'a, str>]>,
-        cipher: &ScopedCipherWithCreds,
+        cipher: &ScopedZeroKmsCipher,
         term_length: usize,
     ) -> Result<Sealed, SealError> {
         let mut vec = Self::seal_all([self], protected_attributes, cipher, term_length).await?;

src/encrypted_table/mod.rs

@@ -49,16 +49,16 @@ impl Deref for Dynamo {
     }
 }
 
-pub type Cipher = ZeroKMSWithClientKey<AutoRefresh<ServiceCredentials>>;
-pub type ScopedCipherWithCreds = ScopedCipher<AutoRefresh<ServiceCredentials>>;
+pub type ZeroKmsCipher = ZeroKMSWithClientKey<AutoRefresh<ServiceCredentials>>;
+pub type ScopedZeroKmsCipher = ScopedCipher<AutoRefresh<ServiceCredentials>>;
 
 pub struct EncryptedTable<D = Dynamo> {
     db: D,
-    cipher: Arc<Cipher>,
+    cipher: Arc<ZeroKmsCipher>,
 }
 
 impl<D> EncryptedTable<D> {
-    pub fn cipher(&self) -> Arc<Cipher> {
+    pub fn cipher(&self) -> Arc<ZeroKmsCipher> {
         self.cipher.clone()
     }
 }
@@ -279,20 +279,15 @@ impl<D> EncryptedTable<D> {
     ) -> Result<Vec<T>, DecryptError>
     where T: Decryptable + Identifiable,
     {
-         // TODO: Temporary obvs
-         let dataset_id = Uuid::parse_str("93e10481-2692-4d65-a619-37e36a496e64").unwrap();
-         let scoped_cipher = ScopedCipherWithCreds::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 = ScopedCipherWithCreds::init(self.cipher.clone(), dataset_id).await.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)?;
 
@@ -323,14 +318,13 @@ impl<D> EncryptedTable<D> {
     pub async fn create_put_patch(
         &self,
         record: PreparedRecord,
+        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 = ScopedCipherWithCreds::init(self.cipher.clone(), dataset_id).await.unwrap();
+        let indexable_cipher = ScopedZeroKmsCipher::init(self.cipher.clone(), dataset_id).await.unwrap();
 
         let PreparedRecord {
             protected_attributes,
@@ -403,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 = ScopedCipherWithCreds::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
@@ -421,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)
         }
@@ -433,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)?;
 
@@ -453,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,
     {
@@ -461,6 +496,7 @@ impl EncryptedTable<Dynamo> {
         let transact_items = self
             .create_put_patch(
                 record,
+                dataset_id,
                 // include all records in the indexes
                 |_, _| true,
             )
@@ -484,7 +520,7 @@ impl EncryptedTable<Dynamo> {
 /// Take a prepared primary key and encrypt it to get the [`PrimaryKeyParts`] which can be used
 /// for retrieval.
 fn encrypt_primary_key_parts(
-    scoped_cipher: &ScopedCipherWithCreds,
+    scoped_cipher: &ScopedZeroKmsCipher,
     prepared_primary_key: PreparedPrimaryKey,
 ) -> Result<PrimaryKeyParts, PrimaryKeyError> {
     let PrimaryKeyParts { mut pk, mut sk } = prepared_primary_key.primary_key_parts;
@@ -500,7 +536,7 @@ fn encrypt_primary_key_parts(
     Ok(PrimaryKeyParts { pk, sk })
 }
 
-async fn decrypt<T>(scoped_cipher: &ScopedCipherWithCreds, 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,
 {
@@ -512,7 +548,7 @@ where
 }
 
 async fn decrypt_all<T>(
-    scoped_cipher: &ScopedCipherWithCreds,
+    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::{
@@ -15,7 +14,7 @@ use crate::{
 };
 use cipherstash_client::encryption::IndexTerm;
 
-use super::{Dynamo, EncryptedTable, ScopedCipherWithCreds, QueryError, SealError};
+use super::{Dynamo, EncryptedTable, ScopedZeroKmsCipher, QueryError, SealError};
 
 /// A builder for a query operation which returns records of type `S`.
 /// `B` is the storage backend used to store the data.
@@ -35,7 +34,7 @@ pub struct PreparedQuery {
 impl PreparedQuery {
     pub async fn encrypt(
         self,
-        scoped_cipher: &ScopedCipherWithCreds,
+        scoped_cipher: &ScopedZeroKmsCipher,
     ) -> Result<AttributeValue, QueryError> {
         let PreparedQuery {
             index_name,
@@ -62,7 +61,7 @@ impl PreparedQuery {
     pub async fn send(
         self,
         table: &EncryptedTable<Dynamo>,
-        scoped_cipher: &ScopedCipherWithCreds,
+        scoped_cipher: &ScopedZeroKmsCipher,
     ) -> Result<Vec<HashMap<String, AttributeValue>>, QueryError> {
         let term = self.encrypt(scoped_cipher).await?;
 
@@ -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 = ScopedCipherWithCreds::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)
     }