Refactor kbs client

Cargo.toml

@@ -30,21 +30,26 @@ const_format = "0.2.30"
 ctr = "0.9.2"
 env_logger = "0.10.0"
 hmac = "0.12.1"
+jwt-simple = "0.11"
 kbs-types = "0.4"
 lazy_static = "1.4.0"
 log = "0.4.14"
 openssl = "0.10"
 prost = "0.11"
 protobuf = "3.2.0"
 rand = "0.8.5"
-reqwest = "0.11.18"
+reqwest = { version = "0.11.18", default-features = false }
 resource_uri = { path = "attestation-agent/deps/resource_uri" }
 ring = "0.16.20"
 rsa = "0.9.2"
+rstest = "0.17"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
+serial_test = "1"
 sha2 = "0.10.7"
 strum = { version = "0.25", features = ["derive"] }
+tempfile = "3.2"
+testcontainers = "0.14"
 thiserror = "1.0"
 tokio = "1.0"
 tonic = "0.9"

attestation-agent/app/Cargo.toml

@@ -20,7 +20,7 @@ prost = { workspace = true, optional = true }
 protobuf = { workspace = true, optional = true }
 serde.workspace = true
 serde_json.workspace = true
-tokio = { workspace = true, default-features = false, features = ["rt-multi-thread", "rt", "sync", "signal"]}
+tokio = { workspace = true, features = ["rt-multi-thread", "rt", "sync", "signal"]}
 tonic = { workspace = true, optional = true }
 ttrpc = { workspace = true, features = ["async"], optional = true }
 

attestation-agent/attester/Cargo.toml

@@ -10,6 +10,7 @@ anyhow.workspace = true
 async-trait.workspace = true
 az-snp-vtpm = { git = "https://github.com/kinvolk/azure-cvm-tooling", rev = "2c2e411", default-features = false, features = ["attester"], optional = true }
 base64.workspace = true
+kbs-types.workspace = true
 log.workspace = true
 occlum_dcap = { git = "https://github.com/occlum/occlum", tag = "v0.29.6", optional = true }
 serde.workspace = true

attestation-agent/attester/src/lib.rs

@@ -3,10 +3,8 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-#[macro_use]
-extern crate strum;
-
 use anyhow::*;
+use kbs_types::Tee;
 
 pub mod sample;
 
@@ -22,38 +20,26 @@ pub mod sgx_dcap;
 #[cfg(feature = "snp-attester")]
 pub mod snp;
 
-/// The supported TEE types:
-/// - Tdx: TDX TEE.
-/// - Sgx: SGX TEE with a LibOS.
-/// - AzSnpVtpm: SEV-SNP TEE for Azure CVMs.
-/// - Snp: SEV-SNP TEE.
-/// - Sample: A dummy TEE that used to test/demo the KBC functionalities.
-#[derive(Debug, EnumString, Display, Clone, Copy)]
-#[strum(ascii_case_insensitive, serialize_all = "lowercase")]
-pub enum Tee {
-    Tdx,
-    #[strum(serialize = "sgx")]
-    Sgx,
-    AzSnpVtpm,
-    Snp,
-    Sample,
-    Unknown,
-}
+pub type BoxedAttester = Box<dyn Attester + Send + Sync>;
+
+impl TryFrom<Tee> for BoxedAttester {
+    type Error = anyhow::Error;
 
-impl Tee {
-    pub fn to_attester(&self) -> Result<Box<dyn Attester + Send + Sync>> {
-        match self {
-            Tee::Sample => Ok(Box::<sample::SampleAttester>::default()),
+    fn try_from(value: Tee) -> Result<Self> {
+        let attester: Box<dyn Attester + Send + Sync> = match value {
+            Tee::Sample => Box::<sample::SampleAttester>::default(),
             #[cfg(feature = "tdx-attester")]
-            Tee::Tdx => Ok(Box::<tdx::TdxAttester>::default()),
+            Tee::Tdx => Box::<tdx::TdxAttester>::default(),
             #[cfg(feature = "sgx-attester")]
-            Tee::Sgx => Ok(Box::<sgx_dcap::SgxDcapAttester>::default()),
+            Tee::Sgx => Box::<sgx_dcap::SgxDcapAttester>::default(),
             #[cfg(feature = "az-snp-vtpm-attester")]
-            Tee::AzSnpVtpm => Ok(Box::<az_snp_vtpm::AzSnpVtpmAttester>::default()),
+            Tee::AzSnpVtpm => Box::<az_snp_vtpm::AzSnpVtpmAttester>::default(),
             #[cfg(feature = "snp-attester")]
-            Tee::Snp => Ok(Box::<snp::SnpAttester>::default()),
+            Tee::Snp => Box::<snp::SnpAttester>::default(),
             _ => bail!("TEE is not supported!"),
-        }
+        };
+
+        Ok(attester)
     }
 }
 
@@ -66,29 +52,30 @@ pub trait Attester {
 }
 
 // Detect which TEE platform the KBC running environment is.
-pub fn detect_tee_type() -> Tee {
+pub fn detect_tee_type() -> Option<Tee> {
     if sample::detect_platform() {
-        return Tee::Sample;
+        return Some(Tee::Sample);
     }
 
     #[cfg(feature = "tdx-attester")]
     if tdx::detect_platform() {
-        return Tee::Tdx;
+        return Some(Tee::Tdx);
     }
 
     #[cfg(feature = "sgx-attester")]
     if sgx_dcap::detect_platform() {
-        return Tee::Sgx;
+        return Some(Tee::Sgx);
     }
 
     #[cfg(feature = "az-snp-vtpm-attester")]
     if az_snp_vtpm::detect_platform() {
-        return Tee::AzSnpVtpm;
+        return Some(Tee::AzSnpVtpm);
     }
 
     #[cfg(feature = "snp-attester")]
     if snp::detect_platform() {
-        return Tee::Snp;
+        return Some(Tee::Snp);
     }
-    Tee::Unknown
+
+    None
 }

attestation-agent/coco_keyprovider/Cargo.toml

@@ -30,6 +30,6 @@ shadow-rs = "0.5.25"
 tonic-build.workspace = true
 
 [dev-dependencies]
-rstest = "0.17.0"
+rstest.workspace = true
 
 [features]

attestation-agent/deps/crypto/Cargo.toml

@@ -13,17 +13,17 @@ ctr = { workspace = true, optional = true }
 kbs-types.workspace = true
 openssl = { workspace = true, features = ["vendored"], optional = true}
 rand.workspace = true
-rsa.workspace = true
+rsa = { workspace = true, optional = true }
 serde.workspace = true
 serde_json.workspace = true
 sha2.workspace = true
 strum.workspace = true
 zeroize.workspace = true
 
 [dev-dependencies]
-rstest = "0.17.0"
+rstest.workspace = true
 
 [features]
 default = ["rust-crypto"]
-rust-crypto = ["dep:aes-gcm", "ctr"]
+rust-crypto = ["dep:aes-gcm", "ctr", "rsa"]
 openssl = ["dep:openssl"]

attestation-agent/deps/crypto/src/asymmetric.rs

@@ -0,0 +1,28 @@
+// Copyright (c) 2023 Alibaba Cloud
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+pub mod rsa {
+    #[cfg(feature = "openssl")]
+    pub use crate::native::rsa::*;
+
+    #[cfg(all(feature = "rust-crypto", not(feature = "openssl")))]
+    pub use crate::rust::rsa::*;
+
+    /// Definations of different Padding mode for encryption. Refer to
+    /// <https://datatracker.ietf.org/doc/html/rfc7518#section-4.1> for
+    /// more information.
+    #[derive(EnumString, AsRefStr)]
+    pub enum PaddingMode {
+        #[strum(serialize = "RSA-OAEP")]
+        OAEP,
+
+        #[strum(serialize = "RSA1_5")]
+        PKCS1v15,
+    }
+
+    pub const RSA_PUBKEY_LENGTH: usize = 2048;
+
+    pub const RSA_KTY: &str = "RSA";
+}

attestation-agent/deps/crypto/src/lib.rs

@@ -29,5 +29,5 @@ mod rust;
 mod symmetric;
 pub use symmetric::*;
 
-mod teekey;
-pub use teekey::*;
+mod asymmetric;
+pub use asymmetric::*;

attestation-agent/deps/crypto/src/native/mod.rs

@@ -10,3 +10,6 @@ pub mod aes256gcm;
 
 pub use aes256ctr::*;
 pub use aes256gcm::*;
+
+pub mod rsa;
+pub use self::rsa::*;

attestation-agent/deps/crypto/src/native/rsa.rs

@@ -0,0 +1,61 @@
+// Copyright (c) 2023 Alibaba Cloud
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+use anyhow::*;
+use openssl::{
+    pkey::Private,
+    rsa::{Padding, Rsa},
+};
+use zeroize::Zeroizing;
+
+use crate::rsa::{PaddingMode, RSA_PUBKEY_LENGTH};
+
+#[derive(Debug, Clone)]
+pub struct RSAKeyPair {
+    private_key: Rsa<Private>,
+}
+
+impl RSAKeyPair {
+    pub fn new() -> Result<RSAKeyPair> {
+        let private_key = Rsa::generate(RSA_PUBKEY_LENGTH as u32)?;
+        Ok(Self { private_key })
+    }
+
+    pub fn decrypt(&self, mode: PaddingMode, cipher_text: Vec<u8>) -> Result<Vec<u8>> {
+        let mut plaintext = [0; RSA_PUBKEY_LENGTH];
+        let decrypted_size = match mode {
+            PaddingMode::OAEP => self
+                .private_key
+                .private_decrypt(&cipher_text, &mut plaintext, Padding::PKCS1_OAEP)
+                .map_err(|e| anyhow!("RSA key decrypt OAEP failed: {:?}", e))?,
+            PaddingMode::PKCS1v15 => self
+                .private_key
+                .private_decrypt(&cipher_text, &mut plaintext, Padding::PKCS1)
+                .map_err(|e| anyhow!("RSA key pkcs1v15 decrypt failed: {:?}", e))?,
+        };
+
+        Ok(plaintext[..decrypted_size].to_vec())
+    }
+
+    pub fn n(&self) -> Vec<u8> {
+        self.private_key.n().to_vec()
+    }
+
+    pub fn e(&self) -> Vec<u8> {
+        self.private_key.e().to_vec()
+    }
+
+    pub fn to_pkcs1_pem(&self) -> Result<Zeroizing<String>> {
+        let res = self.private_key.private_key_to_pem()?;
+        let pem = String::from_utf8(res)?;
+        Ok(Zeroizing::new(pem))
+    }
+
+    pub fn from_pkcs1_pem(pem: &str) -> Result<Self> {
+        let private_key = Rsa::<Private>::private_key_from_pem(pem.as_bytes())?;
+
+        Ok(Self { private_key })
+    }
+}

attestation-agent/deps/crypto/src/rust/mod.rs

@@ -10,3 +10,6 @@ pub mod aes256gcm;
 
 pub use aes256ctr::*;
 pub use aes256gcm::*;
+
+pub mod rsa;
+pub use self::rsa::*;

attestation-agent/deps/crypto/src/rust/rsa.rs

@@ -0,0 +1,73 @@
+// Copyright (c) 2023 Alibaba Cloud
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+//! Implementations of the TeeKey
+
+use anyhow::*;
+use rsa::{
+    pkcs1::{DecodeRsaPrivateKey, EncodeRsaPrivateKey},
+    pkcs8::LineEnding,
+    traits::PublicKeyParts,
+    Oaep, Pkcs1v15Encrypt, RsaPrivateKey, RsaPublicKey,
+};
+use zeroize::Zeroizing;
+
+use crate::rsa::{PaddingMode, RSA_PUBKEY_LENGTH};
+
+#[derive(Debug, Clone)]
+pub struct RSAKeyPair {
+    private_key: RsaPrivateKey,
+    public_key: RsaPublicKey,
+}
+
+impl RSAKeyPair {
+    pub fn new() -> Result<RSAKeyPair> {
+        let mut rng = rand::thread_rng();
+
+        let private_key = RsaPrivateKey::new(&mut rng, RSA_PUBKEY_LENGTH)?;
+        let public_key = RsaPublicKey::from(&private_key);
+
+        Ok(RSAKeyPair {
+            private_key,
+            public_key,
+        })
+    }
+
+    pub fn decrypt(&self, mode: PaddingMode, cipher_text: Vec<u8>) -> Result<Vec<u8>> {
+        match mode {
+            PaddingMode::OAEP => self
+                .private_key
+                .decrypt(Oaep::new::<sha2::Sha256>(), &cipher_text)
+                .map_err(|e| anyhow!("RSA key decrypt OAEP failed: {:?}", e)),
+            PaddingMode::PKCS1v15 => self
+                .private_key
+                .decrypt(Pkcs1v15Encrypt, &cipher_text)
+                .map_err(|e| anyhow!("RSA key pkcs1v15 decrypt failed: {:?}", e)),
+        }
+    }
+
+    pub fn n(&self) -> Vec<u8> {
+        self.public_key.n().to_bytes_be()
+    }
+
+    pub fn e(&self) -> Vec<u8> {
+        self.private_key.e().to_bytes_be()
+    }
+
+    pub fn to_pkcs1_pem(&self) -> Result<Zeroizing<String>> {
+        let res = self.private_key.to_pkcs1_pem(LineEnding::default())?;
+        Ok(res)
+    }
+
+    pub fn from_pkcs1_pem(pem: &str) -> Result<Self> {
+        let private_key = RsaPrivateKey::from_pkcs1_pem(pem)?;
+        let public_key = RsaPublicKey::from(&private_key);
+
+        Ok(Self {
+            private_key,
+            public_key,
+        })
+    }
+}