Add auxilliary data to KMS operations

4 files changed, 33 insertions, 23 deletions

diff --git a/src/kms/awskms.rs b/src/kms/awskms.rs
index 14f0804..4a244db 100644
--- a/src/kms/awskms.rs
+++ b/src/kms/awskms.rs
@@ -20,6 +20,7 @@ pub mod inner {
     extern crate rusoto_core;
     extern crate rusoto_kms;
 
+    use std::collections::HashMap;
     use std::default::Default;
     use std::error::Error;
     use std::fmt;
@@ -28,16 +29,18 @@ pub mod inner {
 
     use self::rusoto_core::Region;
     use self::rusoto_kms::{DecryptRequest, EncryptRequest, Kms, KmsClient};
-    use kms::{EncryptedDEK, KmsError, KmsProvider, PlaintextDEK, DEK_SIZE_BYTES};
+    use kms::{EncryptedDEK, KmsError, KmsProvider, PlaintextDEK, AD, DEK_SIZE_BYTES};
 
-    /// Amazon Key Management Service
+    /// Amazon Web Services Key Management Service
+    /// https://aws.amazon.com/kms/
     pub struct AwsKms {
         kms_client: KmsClient,
         key_id: String,
     }
 
     impl AwsKms {
-        /// Create a new instance from the ARN of a AWS KMS key.
+        /// Create a new instance from the full ARN of a AWS KMS key. The ARN is expected
+        /// to be of the form `arn:aws:kms:some-aws-region:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab`
         pub fn from_arn(arn: &str) -> Result<Self, KmsError> {
             let parts: Vec<&str> = arn.split(':').collect();
 
@@ -74,6 +77,10 @@ pub mod inner {
             encrypt_req.key_id = self.key_id.clone();
             encrypt_req.plaintext = plaintext_dek.clone();
 
+            let mut enc_context = HashMap::new();
+            enc_context.insert("AD".to_string(), AD.to_string());
+            encrypt_req.encryption_context = Some(enc_context);
+
             match self.kms_client.encrypt(encrypt_req).sync() {
                 Ok(result) => {
                     if let Some(ciphertext) = result.ciphertext_blob {
diff --git a/src/kms/envelope.rs b/src/kms/envelope.rs
index da75961..2dc7262 100644
--- a/src/kms/envelope.rs
+++ b/src/kms/envelope.rs
@@ -15,15 +15,14 @@
 extern crate hex;
 
 use std::io::{Cursor, Read, Write};
+use std::str::FromStr;
 
 use ring::aead::{open_in_place, seal_in_place, OpeningKey, SealingKey, AES_256_GCM};
-use ring::rand;
-use ring::rand::SecureRandom;
+use ring::rand::{SecureRandom, SystemRandom};
 
 use super::super::MIN_SEED_LENGTH;
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
-use kms::{KmsError, KmsProvider, DEK_SIZE_BYTES, NONCE_SIZE_BYTES, TAG_SIZE_BYTES};
-use std::string::ToString;
+use kms::{KmsError, KmsProvider, AD, DEK_SIZE_BYTES, NONCE_SIZE_BYTES, TAG_SIZE_BYTES};
 
 const DEK_LEN_FIELD: usize = 2;
 const NONCE_LEN_FIELD: usize = 2;
@@ -43,19 +42,11 @@ const MIN_PAYLOAD_SIZE: usize = DEK_LEN_FIELD
 // No input prefix to skip, consume entire buffer
 const IN_PREFIX_LEN: usize = 0;
 
-// Trivial domain separation to guard against KMS key reuse
-static AD: &[u8; 11] = b"roughenough";
-
 // Convenience function to create zero-filled Vec of given size
 fn vec_zero_filled(len: usize) -> Vec<u8> {
-    let mut v = Vec::with_capacity(len);
-    for _ in 0..len {
-        v.push(0);
-    }
-    return v;
+    (0..len).into_iter().map(|_| 0).collect()
 }
 
-///
 /// Envelope encryption of the long-term key seed value.
 ///
 /// The seed is encrypted using AES-GCM-256 with:
@@ -86,15 +77,22 @@ impl EnvelopeEncryption {
         let mut tmp = Cursor::new(ciphertext_blob);
 
         // Read the lengths of the wrapped DEK and of the nonce
-        let dek_len = tmp.read_u16::<LittleEndian>()?;
-        let nonce_len = tmp.read_u16::<LittleEndian>()?;
+        let dek_len = tmp.read_u16::<LittleEndian>()? as usize;
+        let nonce_len = tmp.read_u16::<LittleEndian>()? as usize;
+
+        if dek_len != DEK_SIZE_BYTES || nonce_len != NONCE_SIZE_BYTES {
+            return Err(KmsError::InvalidData(format!(
+                "invalid DEK ({}) or nonce ({}) length",
+                dek_len, nonce_len
+            )));
+        }
 
         // Consume the wrapped DEK
-        let mut encrypted_dek = vec_zero_filled(dek_len as usize);
+        let mut encrypted_dek = vec_zero_filled(dek_len);
         tmp.read_exact(&mut encrypted_dek)?;
 
         // Consume the nonce
-        let mut nonce = vec_zero_filled(nonce_len as usize);
+        let mut nonce = vec_zero_filled(nonce_len);
         tmp.read_exact(&mut nonce)?;
 
         // Consume the encrypted seed + tag
@@ -109,7 +107,7 @@ impl EnvelopeEncryption {
         match open_in_place(
             &dek_open_key,
             &nonce,
-            AD,
+            AD.as_bytes(),
             IN_PREFIX_LEN,
             &mut encrypted_seed,
         ) {
@@ -148,7 +146,7 @@ impl EnvelopeEncryption {
         let encrypted_seed = match seal_in_place(
             &dek_seal_key,
             &nonce,
-            AD,
+            AD.as_bytes(),
             &mut plaintext_buf,
             TAG_SIZE_BYTES,
         ) {
diff --git a/src/kms/gcpkms.rs b/src/kms/gcpkms.rs
index 7c79f13..d2590f5 100644
--- a/src/kms/gcpkms.rs
+++ b/src/kms/gcpkms.rs
@@ -82,6 +82,7 @@ pub mod inner {
         fn encrypt_dek(&self, plaintext_dek: &PlaintextDEK) -> Result<EncryptedDEK, KmsError> {
             let mut request = EncryptRequest::default();
             request.plaintext = Some(base64::encode(plaintext_dek));
+            request.additional_authenticated_data = Some(base64::encode(AD));
 
             let hub = self.new_hub();
             let result = hub
@@ -99,13 +100,14 @@ pub mod inner {
                         Err(self.pretty_http_error(&http_resp))
                     }
                 }
-                Err(e) => Err(KmsError::OperationFailed(format!("encrypt_dek() {:?}", e)))
+                Err(e) => Err(KmsError::OperationFailed(format!("encrypt_dek() {:?}", e))),
             }
         }
 
         fn decrypt_dek(&self, encrypted_dek: &EncryptedDEK) -> Result<PlaintextDEK, KmsError> {
             let mut request = DecryptRequest::default();
             request.ciphertext = Some(base64::encode(encrypted_dek));
+            request.additional_authenticated_data = Some(base64::encode(AD));
 
             let hub = self.new_hub();
             let result = hub
diff --git a/src/kms/mod.rs b/src/kms/mod.rs
index 56e7631..aa609b5 100644
--- a/src/kms/mod.rs
+++ b/src/kms/mod.rs
@@ -98,6 +98,9 @@ const TAG_SIZE_BYTES: usize = 16;
 // Size of the 256-bit Data Encryption Key (DEK) in bytes.
 const DEK_SIZE_BYTES: usize = 32;
 
+// Trivial domain separation to guard against KMS key reuse
+const AD: &str = "roughenough";
+
 /// An unencrypted (plaintext) 256-bit Data Encryption Key (DEK).
 pub type PlaintextDEK = Vec<u8>;