KMS: add API definitions and crate for KMS integration

This commit include the following:
- Added a crate `deps/kms`
- Added Datakey API abstraction inside `deps/kms`
- Added KMS module API abstraction inside `deps/kms`
- Implemented an example for KMS

Signed-off-by: Xynnn007 <[email protected]>

Cargo.toml

@@ -8,6 +8,7 @@ members = [
     "deps/resource_uri",
     "deps/crypto",
     "deps/sev",
+    "deps/kms",
     "coco_keyprovider"
 ]
 

deps/kms/Cargo.toml

@@ -0,0 +1,26 @@
+[package]
+name = "kms"
+version = "0.1.0"
+authors = ["The Attestation Agent Authors"]
+publish = false
+edition = "2021"
+
+[dependencies]
+aes-gcm = "0.10.1"
+anyhow.workspace = true
+async-trait.workspace = true
+crypto = { path = "../crypto", default-features = false, optional = true}
+rand = { version = "0.8.4", optional = true }
+strum.workspace = true
+tokio = "1.0"
+zeroize = { version = "1.6.0", optional = true }
+
+[dev-dependencies]
+tokio = { version = "1.0", features = ["rt", "macros" ] }
+
+[features]
+default = [ "sample-rust" ]
+
+sample = [ "zeroize", "rand" ]
+sample-openssl = [ "crypto/openssl", "sample" ]
+sample-rust = [ "crypto/rust-crypto", "sample" ]

deps/kms/src/api.rs

@@ -0,0 +1,36 @@
+// Copyright (c) 2023 Alibaba Cloud
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+use anyhow::Result;
+use async_trait::async_trait;
+
+#[async_trait]
+pub trait KMS {
+    /// Generate a [`Datakey`] object related to this kind of KMS.
+    async fn generate_data_key(&mut self) -> Result<Box<dyn Datakey>>;
+}
+
+#[async_trait]
+pub trait Datakey {
+    /// Use this key to encrypt the given data slice, and return the
+    /// ciphertext blob that can be decrypted by calling `decrypt` API.
+    async fn encrypt(&mut self, plaintext: &[u8]) -> Result<Vec<u8>>;
+
+    /// Use this datakey to decrypt the given data blob, and return the
+    /// plaintext.
+    async fn decrypt(&mut self, ciphertext: &[u8]) -> Result<Vec<u8>>;
+
+    /// Export the key blob that can be used to reconstruct the Datakey
+    /// except the plaintext of the Datakey.
+    ///
+    /// Typical information of key blob is the CMK id, uuid of the key,
+    /// etc. This keyblob will be stored inside KBS.
+    async fn export_keyblob(&mut self) -> Result<Vec<u8>>;
+
+    /// Construct the Datakey from the given blob.
+    async fn from_keyblob(blob: &[u8]) -> Result<Self>
+    where
+        Self: Sized;
+}

deps/kms/src/lib.rs

@@ -0,0 +1,31 @@
+// Copyright (c) 2023 Alibaba Cloud
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+//! # KMS
+//!
+//! This lib defines the API of the KMSes which will be used to implement
+//! image encryption and decryption.
+//!
+//! ## Trait Definitions
+//!
+//! ### Datakey
+//!
+//! [`Datakey`] is an abstraction of a key with the following attributes
+//! in different KMSes:
+//! - Can be used to encrypt/decrypt large pieces of data
+//! - Plaintext of the key can be exported outside the KMS
+//! - KMS can help to decrypt the ciphertext of the key to get the plaintext
+//! of the key.
+//!
+//! Typical examples of [`Datakey`]s are
+//! - [Datakey](https://www.alibabacloud.com/help/en/key-management-service/latest/kms-generatedatakey-v2#doc-api-Kms-GenerateDataKey) in Alibaba Cloud
+
+pub mod api;
+pub use api::Datakey;
+pub use api::KMS;
+
+pub mod plugins;
+
+pub mod types;

deps/kms/src/plugins/mod.rs

@@ -0,0 +1,6 @@
+// Copyright (c) 2023 Alibaba Cloud
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+pub mod sample;

deps/kms/src/plugins/sample.rs

@@ -0,0 +1,127 @@
+// Copyright (c) 2023 Alibaba Cloud
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+use aes_gcm::{aead::Aead, Aes256Gcm, Key, KeyInit, Nonce};
+use anyhow::{anyhow, Result};
+use async_trait::async_trait;
+use crypto::WrapType;
+use rand::RngCore;
+use zeroize::Zeroizing;
+
+use crate::{Datakey, KMS};
+
+const IV: &[u8] = b"unique nonce";
+
+/// Only for sample
+pub const EXAMPLE_CMK: &[u8] = &[
+    217, 155, 119, 5, 176, 186, 122, 22, 130, 149, 179, 163, 54, 114, 112, 176, 221, 155, 55, 27,
+    245, 20, 202, 139, 155, 167, 240, 163, 55, 17, 218, 234,
+];
+
+/// A fake a KMS implementation
+pub struct SampleKms;
+
+#[async_trait]
+impl KMS for SampleKms {
+    async fn generate_data_key(&mut self) -> Result<Box<dyn Datakey>> {
+        Ok(Box::<SampleDatakey>::default())
+    }
+}
+
+/// Fake Datakey implementation
+pub struct SampleDatakey {
+    /// The `Zeroizing` here will help to flush the memory when the
+    /// [`Datakey`] object is dropped
+    key: Zeroizing<Vec<u8>>,
+}
+
+impl Default for SampleDatakey {
+    /// Here we simulate behavior of KMS. A random datakey will be generated
+    fn default() -> Self {
+        let mut key = Zeroizing::new(Vec::new());
+        key.resize(32, b' ');
+        rand::thread_rng().fill_bytes(&mut key);
+        Self { key }
+    }
+}
+
+#[async_trait]
+impl Datakey for SampleDatakey {
+    async fn encrypt(&mut self, plaintext: &[u8]) -> Result<Vec<u8>> {
+        let encryption_key = Key::<Aes256Gcm>::from_slice(&self.key);
+        let cipher = Aes256Gcm::new(encryption_key);
+        let nonce = Nonce::from_slice(IV);
+        cipher
+            .encrypt(nonce, plaintext)
+            .map_err(|e| anyhow!("Decrypt failed: {:?}", e))
+    }
+
+    /// We use this datakey to encrypt
+    async fn decrypt(&mut self, ciphertext: &[u8]) -> Result<Vec<u8>> {
+        crypto::decrypt(
+            self.key.clone(),
+            ciphertext.to_vec(),
+            IV.to_vec(),
+            WrapType::Aes256Gcm.as_ref(),
+        )
+    }
+
+    /// Here we simulate the operation of a real KMS datakey. An encrypted
+    /// version of Datakey will be exported as the key blob.
+    async fn export_keyblob(&mut self) -> Result<Vec<u8>> {
+        let encryption_key = Key::<Aes256Gcm>::from_slice(EXAMPLE_CMK);
+        let cipher = Aes256Gcm::new(encryption_key);
+        let nonce = Nonce::from_slice(IV);
+        cipher
+            .encrypt(nonce, &**self.key)
+            .map_err(|e| anyhow!("Decrypt failed: {:?}", e))
+    }
+
+    /// Here we simulate the operation of a real KMS datakey. An encrypted
+    /// version of Datakey will be imported, and then be decrypted by a
+    /// given CMK. Here the CMK is hardcoded. In real scenarios of KMS,
+    /// a CMK id should be included the key blob.
+    async fn from_keyblob(blob: &[u8]) -> Result<Self>
+    where
+        Self: Sized,
+    {
+        let key = crypto::decrypt(
+            EXAMPLE_CMK.to_vec().into(),
+            blob.to_vec(),
+            IV.to_vec(),
+            WrapType::Aes256Gcm.as_ref(),
+        )?;
+
+        Ok(Self { key: key.into() })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        plugins::sample::{SampleDatakey, SampleKms},
+        Datakey, KMS,
+    };
+
+    const TEST_TEXT: &[u8] = b"this is for example";
+
+    #[tokio::test]
+    async fn datakey_lifetime() {
+        let mut kms = SampleKms {};
+        let mut datakey = kms.generate_data_key().await.expect("generate datakey");
+        let cipher_text = datakey.encrypt(TEST_TEXT).await.expect("encryption");
+
+        // export key blob
+        let key_blob = datakey.export_keyblob().await.expect("export blob");
+
+        // import key blob and decrypt
+        let mut datakey2 = SampleDatakey::from_keyblob(&key_blob)
+            .await
+            .expect("create key from blob");
+
+        let plaintext = datakey2.decrypt(&cipher_text).await.expect("decryption");
+        assert_eq!(plaintext, TEST_TEXT);
+    }
+}

deps/kms/src/types.rs

@@ -0,0 +1,25 @@
+// Copyright (c) 2023 Alibaba Cloud
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+use anyhow::Result;
+use strum::{AsRefStr, EnumString};
+
+use crate::KMS;
+
+#[derive(EnumString, AsRefStr)]
+pub enum KMSTypes {
+    #[cfg(feature = "sample")]
+    #[strum(serialize = "sample")]
+    Sample,
+}
+
+impl KMSTypes {
+    pub fn to_kms(&self) -> Result<Box<dyn KMS>> {
+        match self {
+            #[cfg(feature = "sample")]
+            KMSTypes::Sample => Ok(Box::new(crate::plugins::sample::SampleKms {})),
+        }
+    }
+}