Using ZeroKMS and ScopedCipher directly, rough working version - tests pass

Author: coderdan

Cargo.lock

@@ -13,7 +13,7 @@ keywords = ["cryptography", "security", "databases", "encryption", "dynamodb"]
 categories = ["cryptography", "database"]
 
 [dependencies]
-cipherstash-client = { version = ">=0.12.4" }
+cipherstash-client = { version = ">=0.12.4", path = "../cipherstash-suite/packages/cipherstash-client/" }
 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,12 +1,10 @@
 use crate::{
     crypto::{attrs::flattened_protected_attributes::FlattenedAttrName, SealError},
-    encrypted_table::TableAttributes,
+    encrypted_table::{ScopedCipherWithCreds, TableAttributes},
     traits::TableAttribute,
 };
 use cipherstash_client::{
-    credentials::{service_credentials::ServiceToken, Credentials},
-    encryption::Encryption,
-    zerokms::EncryptedRecord,
+    encryption::Plaintext, zerokms::EncryptedRecord
 };
 use itertools::Itertools;
 
@@ -36,7 +34,7 @@ impl FlattenedEncryptedAttributes {
     /// Decrypt self, returning a [FlattenedProtectedAttributes].
     pub(crate) async fn decrypt_all(
         self,
-        cipher: &Encryption<impl Credentials<Token = ServiceToken>>,
+        cipher: &ScopedCipherWithCreds,
     ) -> Result<FlattenedProtectedAttributes, SealError> {
         let descriptors = self
             .attrs
@@ -47,8 +45,13 @@ 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.

src/crypto/attrs/flattened_encrypted_attributes.rs

@@ -2,10 +2,9 @@ use super::{
     flattened_encrypted_attributes::FlattenedEncryptedAttributes,
     normalized_protected_attributes::NormalizedKey,
 };
-use crate::{crypto::SealError, encrypted_table::AttributeName};
+use crate::{crypto::SealError, encrypted_table::{AttributeName, ScopedCipherWithCreds}};
 use cipherstash_client::{
-    credentials::{service_credentials::ServiceToken, Credentials},
-    encryption::{BytesWithDescriptor, Encryption, Plaintext},
+    encryption::{BytesWithDescriptor, Plaintext}, zerokms::EncryptPayload,
 };
 use itertools::Itertools;
 
@@ -32,13 +31,16 @@ 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: &ScopedCipherWithCreds,
         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

@@ -1,6 +1,7 @@
 use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
 
 /// Base64 encode the provided buffer using a URL safe scheme with no padding
+// TODO: Make this consume and zeroize
 pub fn b64_encode(x: impl AsRef<[u8]>) -> String {
     URL_SAFE_NO_PAD.encode(x)
 }

src/crypto/b64_encode.rs

@@ -4,16 +4,12 @@ mod sealed;
 mod sealer;
 mod unsealed;
 use crate::{
-    traits::{PrimaryKeyError, PrimaryKeyParts, ReadConversionError, WriteConversionError},
-    Identifiable, IndexType, PrimaryKey,
+    traits::{PrimaryKeyError, PrimaryKeyParts, ReadConversionError, WriteConversionError}, Identifiable, IndexType, PrimaryKey
 };
 use cipherstash_client::{
-    credentials::{service_credentials::ServiceToken, Credentials},
     encryption::{
-        compound_indexer::{CompoundIndex, ExactIndex},
-        Encryption, EncryptionError, Plaintext, TypeParseError,
-    },
-    zerokms::Error as ZeroKmsError,
+        EncryptionError, TypeParseError
+    },zerokms
 };
 use miette::Diagnostic;
 use std::borrow::Cow;
@@ -47,15 +43,16 @@ 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),
+    //#[diagnostic(transparent)] // TODO
+    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,31 +91,64 @@ pub(crate) fn all_index_keys<'a>(
         .collect()
 }
 
-/// Use a CipherStash [`ExactIndex`] to take the HMAC of a string with a provided salt
+/* /// 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(
+pub fn prf(
     value: &str,
     salt: Option<&str>,
-    cipher: &Encryption<impl Credentials<Token = ServiceToken>>,
+    cipher: &Cipher,
+    // TODO: Pass a DatasetWithRootKey (use a Protected)
+    root_key: [u8; 32],
 ) -> 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,
-        )?
+    // passing None here results in no terms so pass an empty string
+    let salt = salt.unwrap_or("");
+    let accumulator = Accumulator::from_salt(salt);
+
+    index
+        .compose_index(root_key, plaintext.into(), accumulator)?
+        // TODO: Use a constant for the 32
+        .truncate(32)
+        .map(IndexTerm::from)?
         .as_binary()
         .ok_or(EncryptionError::IndexingError(
             "Invalid term type".to_string(),
         ))
+} */
+
+/*// FIXME: Don't use the root key here
+pub fn query_compound_prf<I>(index: I, plaintext: ComposablePlaintext, info: String, root_key: [u8; 32]) -> Result<IndexTerm, SealError> where I: ComposableIndex + Send {
+    let index = CompoundIndex::new(index);
+    let accumulator = Accumulator::from_salt(info);
+
+    index
+        .compose_query(root_key, plaintext, accumulator)?
+        .exactly_one()
+        // FIXME: Don't use a magic number
+        .and_then(|term| term.truncate(12))
+        .and_then(|term| IndexTerm::try_from(term))
+        .map_err(EncryptionError::from)
+        .map_err(SealError::from)
 }
 
+// FIXME: Don't use the root key here
+pub fn compound_prf<I>(index: I, plaintext: ComposablePlaintext, info: String, root_key: [u8; 32]) -> Result<IndexTerm, SealError> where I: ComposableIndex + Send {
+    let index = CompoundIndex::new(index);
+    let accumulator = Accumulator::from_salt(info);
+
+    let term = index
+        .compose_index(root_key, plaintext, accumulator)?
+        // FIXME: Don't use a magic number
+        .truncate(12)
+        .map_err(EncryptionError::from)?;
+
+    // Saftey: This conversion is Infallible
+    Ok(IndexTerm::try_from(term).unwrap())
+} */
+
 // Contains all the necessary information to encrypt the primary key pair
 #[derive(Clone)]
 pub struct PreparedPrimaryKey {

src/crypto/mod.rs

@@ -1,14 +1,10 @@
 use crate::{
     crypto::attrs::FlattenedEncryptedAttributes,
-    encrypted_table::TableEntry,
+    encrypted_table::{ScopedCipherWithCreds, TableEntry},
     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: &ScopedCipherWithCreds,
     ) -> 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: &ScopedCipherWithCreds,
     ) -> Result<Unsealed, SealError> {
         let mut vec = Self::unseal_all(vec![self], spec, cipher).await?;
 
@@ -207,35 +203,34 @@ impl TryFrom<SealedTableEntry> for HashMap<String, AttributeValue> {
 
 #[cfg(test)]
 mod tests {
+    use crate::encrypted_table::{Cipher, ScopedCipherWithCreds};
+
     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 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<Cipher, Box<dyn std::error::Error>> {
-        let console_config = ConsoleConfig::builder().with_env().build()?;
+    async fn get_cipher() -> Result<Arc<Cipher>, 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]
@@ -245,9 +240,14 @@ mod tests {
             sort_key_prefix: "test".to_string(),
         };
         let cipher = get_cipher().await?;
-        let results = SealedTableEntry::unseal_all(vec![], spec, &cipher)
+        // TODO: Temporary obvs
+        let dataset_id = Uuid::parse_str("93e10481-2692-4d65-a619-37e36a496e64").unwrap();
+        let scoped_cipher = ScopedCipherWithCreds::init(cipher, dataset_id).await;
+
+        let results = SealedTableEntry::unseal_all(vec![], spec, &scoped_cipher)
             .await
             .into_diagnostic()?;
+
         assert!(results.is_empty());
 
         Ok(())