Merge pull request #86 from cipherstash/feat/per-dataset-ops

Feat/per dataset ops

Author: coderdan

Cargo.lock

@@ -13,7 +13,7 @@ keywords = ["cryptography", "security", "databases", "encryption", "dynamodb"]
 categories = ["cryptography", "database"]
 
 [dependencies]
-cipherstash-client = { version = "0.12.5" }
+cipherstash-client = { version = "0.13.0-pre.1" }
 cipherstash-dynamodb-derive = { version = "0.8", path = "cipherstash-dynamodb-derive" }
 
 aws-sdk-dynamodb = "1.3.0"
@@ -26,6 +26,7 @@ hex = "0.4.3"
 tracing = { version = "0.1", features = ["log"] }
 tracing-subscriber = { version = "0.3", default-features = false, features = ["fmt"]}
 miette = "7.2.0"
+uuid = "1.10.0"
 
 [dev-dependencies]
 tokio = { version = "1", features = ["full"] }

Cargo.toml

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

docker-compose.yaml

@@ -1,13 +1,9 @@
 use crate::{
     crypto::{attrs::flattened_protected_attributes::FlattenedAttrName, SealError},
-    encrypted_table::TableAttributes,
+    encrypted_table::{TableAttributes, ZeroKmsCipher},
     traits::TableAttribute,
 };
-use cipherstash_client::{
-    credentials::{service_credentials::ServiceToken, Credentials},
-    encryption::Encryption,
-    zerokms::EncryptedRecord,
-};
+use cipherstash_client::{encryption::Plaintext, zerokms::EncryptedRecord};
 use itertools::Itertools;
 
 use super::FlattenedProtectedAttributes;
@@ -36,7 +32,7 @@ impl FlattenedEncryptedAttributes {
     /// Decrypt self, returning a [FlattenedProtectedAttributes].
     pub(crate) async fn decrypt_all(
         self,
-        cipher: &Encryption<impl Credentials<Token = ServiceToken>>,
+        cipher: &ZeroKmsCipher,
     ) -> Result<FlattenedProtectedAttributes, SealError> {
         let descriptors = self
             .attrs
@@ -47,8 +43,16 @@ impl FlattenedEncryptedAttributes {
         cipher
             .decrypt(self.attrs.into_iter())
             .await
-            .map(|records| records.into_iter().zip(descriptors.into_iter()).collect())
-            .map_err(SealError::from)
+            .map(|records| {
+                records
+                    .into_iter()
+                    // FIXME: We should change the decrypt method to return a plaintext and/or make a Plaintext::from_bytes method which consumes the bytes
+                    .map(|bytes| Plaintext::from_slice(&bytes).unwrap())
+                    .zip(descriptors.into_iter())
+                    .collect()
+            })
+            // FIXME: EncryptedRecord should return an error exposed in cipherstash_client
+            .map_err(|_| SealError::AssertionFailed("FIXME".to_string()))
     }
 
     /// Denormalize the encrypted records into a TableAttributes.
@@ -59,7 +63,7 @@ impl FlattenedEncryptedAttributes {
             .into_iter()
             .map(|record| {
                 record
-                    .to_vec()
+                    .to_mp_bytes()
                     .map(|data| (FlattenedAttrName::parse(&record.descriptor), data))
                     .map_err(|_| SealError::AssertionFailed("Decryption failed".to_string()))
             })

src/crypto/attrs/flattened_encrypted_attributes.rs

@@ -2,10 +2,13 @@ use super::{
     flattened_encrypted_attributes::FlattenedEncryptedAttributes,
     normalized_protected_attributes::NormalizedKey,
 };
-use crate::{crypto::SealError, encrypted_table::AttributeName};
+use crate::{
+    crypto::SealError,
+    encrypted_table::{AttributeName, ScopedZeroKmsCipher},
+};
 use cipherstash_client::{
-    credentials::{service_credentials::ServiceToken, Credentials},
-    encryption::{BytesWithDescriptor, Encryption, Plaintext},
+    encryption::{BytesWithDescriptor, Plaintext},
+    zerokms::EncryptPayload,
 };
 use itertools::Itertools;
 
@@ -32,13 +35,14 @@ impl FlattenedProtectedAttributes {
     /// The output is a vec of `chunk_into` [FlattenedEncryptedAttributes] objects.
     pub(crate) async fn encrypt_all(
         self,
-        cipher: &Encryption<impl Credentials<Token = ServiceToken>>,
+        cipher: &ScopedZeroKmsCipher,
         chunk_into: usize,
     ) -> Result<Vec<FlattenedEncryptedAttributes>, SealError> {
         let chunk_size = self.0.len() / chunk_into;
+        let payloads: Vec<BytesWithDescriptor> = self.0.into_iter().map(Into::into).collect();
 
         cipher
-            .encrypt(self.0.into_iter())
+            .encrypt(payloads.iter().map(EncryptPayload::from))
             .await?
             .into_iter()
             .chunks(chunk_size)

src/crypto/attrs/flattened_protected_attributes.rs

@@ -8,12 +8,8 @@ use crate::{
     Identifiable, IndexType, PrimaryKey,
 };
 use cipherstash_client::{
-    credentials::{service_credentials::ServiceToken, Credentials},
-    encryption::{
-        compound_indexer::{CompoundIndex, ExactIndex},
-        Encryption, EncryptionError, Plaintext, TypeParseError,
-    },
-    zerokms::Error as ZeroKmsError,
+    encryption::{EncryptionError, TypeParseError},
+    zerokms,
 };
 use miette::Diagnostic;
 use std::borrow::Cow;
@@ -47,15 +43,15 @@ pub enum SealError {
     #[error("Assertion failed: {0}")]
     AssertionFailed(String),
 
-    // Note that we don't expose the specific error type here
-    // so as to avoid leaking any information
     #[error(transparent)]
-    EncryptionError(#[from] EncryptionError),
+    CryptoError(#[from] zerokms::Error),
 
+    /// Error resulting from Indexing in `cipherstash_client::encryption::compound_indexer`
     #[error(transparent)]
-    ZeroKmsError(#[from] ZeroKmsError),
+    IndexError(#[from] EncryptionError),
 }
 
+// TODO: Possibly remove this
 #[derive(Error, Debug)]
 pub enum CryptoError {
     #[error("EncryptionError: {0}")]
@@ -94,32 +90,6 @@ pub(crate) fn all_index_keys<'a>(
         .collect()
 }
 
-/// Use a CipherStash [`ExactIndex`] to take the HMAC of a string with a provided salt
-///
-/// This value is used for term index keys and "encrypted" partition / sort keys
-pub fn hmac(
-    value: &str,
-    salt: Option<&str>,
-    cipher: &Encryption<impl Credentials<Token = ServiceToken>>,
-) -> Result<Vec<u8>, EncryptionError> {
-    let plaintext = Plaintext::Utf8Str(Some(value.to_string()));
-    let index = CompoundIndex::new(ExactIndex::new(vec![]));
-
-    cipher
-        .compound_index(
-            &index,
-            plaintext,
-            // passing None here results in no terms so pass an empty string
-            Some(salt.unwrap_or("")),
-            32,
-        )?
-        .as_binary()
-        .ok_or(EncryptionError::IndexingError(
-            "Invalid term type".to_string(),
-        ))
-}
-
-// Contains all the necessary information to encrypt the primary key pair
 #[derive(Clone)]
 pub struct PreparedPrimaryKey {
     pub primary_key_parts: PrimaryKeyParts,

src/crypto/mod.rs

@@ -1,14 +1,10 @@
 use crate::{
     crypto::attrs::FlattenedEncryptedAttributes,
-    encrypted_table::TableEntry,
+    encrypted_table::{TableEntry, ZeroKmsCipher},
     traits::{ReadConversionError, WriteConversionError},
     Decryptable, Identifiable,
 };
 use aws_sdk_dynamodb::{primitives::Blob, types::AttributeValue};
-use cipherstash_client::{
-    credentials::{service_credentials::ServiceToken, Credentials},
-    encryption::Encryption,
-};
 use itertools::Itertools;
 use std::{borrow::Cow, collections::HashMap};
 
@@ -72,7 +68,7 @@ impl SealedTableEntry {
     pub(crate) async fn unseal_all(
         items: Vec<Self>,
         spec: UnsealSpec<'_>,
-        cipher: &Encryption<impl Credentials<Token = ServiceToken>>,
+        cipher: &ZeroKmsCipher,
     ) -> Result<Vec<Unsealed>, SealError> {
         let UnsealSpec {
             protected_attributes,
@@ -134,7 +130,7 @@ impl SealedTableEntry {
     pub(crate) async fn unseal(
         self,
         spec: UnsealSpec<'_>,
-        cipher: &Encryption<impl Credentials<Token = ServiceToken>>,
+        cipher: &ZeroKmsCipher,
     ) -> Result<Unsealed, SealError> {
         let mut vec = Self::unseal_all(vec![self], spec, cipher).await?;
 
@@ -207,35 +203,36 @@ impl TryFrom<SealedTableEntry> for HashMap<String, AttributeValue> {
 
 #[cfg(test)]
 mod tests {
+    use crate::encrypted_table::ZeroKmsCipher;
+
     use super::SealedTableEntry;
     use cipherstash_client::{
-        credentials::{auto_refresh::AutoRefresh, service_credentials::ServiceCredentials},
-        encryption::Encryption,
-        ConsoleConfig, ZeroKMS, ZeroKMSConfig,
+        credentials::auto_refresh::AutoRefresh, ConsoleConfig, ZeroKMS, ZeroKMSConfig,
     };
     use miette::IntoDiagnostic;
-    use std::borrow::Cow;
-
-    type Cipher = Encryption<AutoRefresh<ServiceCredentials>>;
+    use std::{borrow::Cow, sync::Arc};
 
     // 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()?;
+    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)
             .with_env()
             .console_config(&console_config)
-            .build_with_client_key()?;
+            .build_with_client_key()
+            .into_diagnostic()?;
 
-        let zero_kms_client = ZeroKMS::new_with_client_key(
+        let cipher = 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))
+        Ok(Arc::new(cipher))
     }
 
     #[tokio::test]
@@ -248,6 +245,7 @@ mod tests {
         let results = SealedTableEntry::unseal_all(vec![], spec, &cipher)
             .await
             .into_diagnostic()?;
+
         assert!(results.is_empty());
 
         Ok(())