Added two security enhancements to AA

attestation-agent/attester/src/lib.rs

@@ -28,7 +28,7 @@ pub mod snp;
 /// - 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)]
+#[derive(Debug, EnumString, Display, Clone, Copy)]
 #[strum(ascii_case_insensitive, serialize_all = "lowercase")]
 pub enum Tee {
     Tdx,

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

@@ -6,8 +6,10 @@
 //! Implementations of the TeeKey
 
 use anyhow::*;
+use base64::engine::general_purpose::STANDARD;
 use base64::Engine;
 use kbs_types::TeePubKey;
+use rsa::pkcs1::DecodeRsaPublicKey;
 use rsa::{traits::PublicKeyParts, Pkcs1v15Encrypt, RsaPrivateKey, RsaPublicKey};
 use sha2::{Digest, Sha384};
 
@@ -69,3 +71,17 @@ pub fn hash_chunks(chunks: Vec<Vec<u8>>) -> Vec<u8> {
 
     hasher.finalize().to_vec()
 }
+
+// Convert PKCS#1 PEM public key to TeePubKey
+pub fn pkcs1_pem_to_teepubkey(pem: String) -> Result<TeePubKey> {
+    let public_key = RsaPublicKey::from_pkcs1_pem(&pem)?;
+    let k_mod = STANDARD.encode(public_key.n().to_bytes_be());
+    let k_exp = STANDARD.encode(public_key.e().to_bytes_be());
+    let tee_pubkey = TeePubKey {
+        kty: RSA_KEY_TYPE.to_string(),
+        alg: RSA_ALGORITHM.to_string(),
+        k_mod,
+        k_exp,
+    };
+    Ok(tee_pubkey)
+}

attestation-agent/kbc/src/cc_kbc/mod.rs

@@ -62,7 +62,7 @@ impl Kbc {
         Ok(Kbc {
             kbs_uri: url,
             token: None,
-            kbs_protocol_wrapper: KbsProtocolWrapper::new().unwrap(),
+            kbs_protocol_wrapper: KbsProtocolWrapper::new(vec![]).unwrap(),
         })
     }
 

attestation-agent/kbs_protocol/src/lib.rs

@@ -7,8 +7,8 @@ use anyhow::*;
 use async_trait::async_trait;
 use attester::{detect_tee_type, Attester, Tee};
 use core::time::Duration;
-use crypto::{hash_chunks, TeeKey};
-use kbs_types::{Attestation, ErrorInformation};
+use crypto::{hash_chunks, pkcs1_pem_to_teepubkey, TeeKey};
+use kbs_types::{Attestation, ErrorInformation, TeePubKey};
 use std::convert::TryFrom;
 use types::*;
 use url::Url;
@@ -25,7 +25,7 @@ pub trait KbsRequest {
     /// Get confidential resource
     async fn http_get(&mut self, resource_url: String) -> Result<Vec<u8>>;
     /// Attestation and get attestation results token (Base64 endcoded)
-    async fn attest(&mut self, host_url: String) -> Result<String>;
+    async fn attest(&mut self, url: String, tee_pubkey_pem: Option<String>) -> Result<String>;
 }
 
 type BoxedAttester = Box<dyn Attester + Send + Sync>;
@@ -45,22 +45,32 @@ pub struct NoAttester;
 pub struct KbsProtocolWrapperBuilder<T, A> {
     tee: T,
     attester: A,
+    kbs_certs: Vec<String>,
 }
 type WrapperBuilder<T, A> = KbsProtocolWrapperBuilder<T, A>;
 
+impl<T, A> WrapperBuilder<T, A> {
+    pub fn add_kbs_cert(mut self, cert_pem: String) -> Self {
+        self.kbs_certs.push(cert_pem);
+        self
+    }
+}
+
 impl WrapperBuilder<NoTee, NoAttester> {
     #[allow(clippy::new_without_default)]
     pub fn new() -> Self {
         Self {
             tee: NoTee,
             attester: NoAttester,
+            kbs_certs: vec![],
         }
     }
 
     pub fn with_tee(self, tee: Tee) -> WrapperBuilder<Tee, NoAttester> {
         WrapperBuilder {
             tee,
             attester: self.attester,
+            kbs_certs: self.kbs_certs,
         }
     }
 }
@@ -70,6 +80,7 @@ impl WrapperBuilder<Tee, NoAttester> {
         WrapperBuilder {
             tee: self.tee,
             attester,
+            kbs_certs: self.kbs_certs,
         }
     }
 }
@@ -94,15 +105,15 @@ impl WrapperBuilder<Tee, BoxedAttester> {
 impl TryFrom<WrapperBuilder<Tee, BoxedAttester>> for KbsProtocolWrapper {
     type Error = Error;
     fn try_from(builder: WrapperBuilder<Tee, BoxedAttester>) -> Result<Self> {
-        let mut wrapper = Self::new()?;
+        let mut wrapper = Self::new(builder.kbs_certs)?;
         wrapper.tee = builder.tee.to_string();
         wrapper.attester = Some(builder.attester);
         Ok(wrapper)
     }
 }
 
 impl KbsProtocolWrapper {
-    pub fn new() -> Result<KbsProtocolWrapper> {
+    pub fn new(kbs_root_certs_pem: Vec<String>) -> Result<KbsProtocolWrapper> {
         let tee_key = TeeKey::new().map_err(|e| anyhow!("Generate TEE key failed: {e}"))?;
         // Detect TEE type of the current platform.
         let tee_type = detect_tee_type();
@@ -114,17 +125,12 @@ impl KbsProtocolWrapper {
             attester,
             tee_key,
             nonce: None,
-            http_client: build_http_client()?,
+            http_client: build_http_client(kbs_root_certs_pem)?,
             authenticated: false,
         })
     }
 
-    fn generate_evidence(&self) -> Result<Attestation> {
-        let tee_pubkey = self
-            .tee_key
-            .export_pubkey()
-            .map_err(|e| anyhow!("Export TEE pubkey failed: {:?}", e))?;
-
+    fn generate_evidence(&self, tee_pubkey: TeePubKey) -> Result<Attestation> {
         let nonce = self
             .nonce
             .to_owned()
@@ -161,7 +167,11 @@ impl KbsProtocolWrapper {
         &mut self.http_client
     }
 
-    async fn attestation(&mut self, kbs_root_url: String) -> Result<String> {
+    async fn attestation(
+        &mut self,
+        kbs_root_url: String,
+        tee_pubkey_pem: Option<String>,
+    ) -> Result<String> {
         let challenge = self
             .http_client()
             .post(format!("{kbs_root_url}{KBS_PREFIX}/auth"))
@@ -173,11 +183,23 @@ impl KbsProtocolWrapper {
             .await?;
         self.nonce = Some(challenge.nonce.clone());
 
+        let pubkey = match tee_pubkey_pem {
+            Some(k) => pkcs1_pem_to_teepubkey(k)?,
+            None => {
+                let key = TeeKey::new()?;
+                let key_exp = key
+                    .export_pubkey()
+                    .map_err(|e| anyhow!("Export TEE pubkey failed: {:?}", e))?;
+                self.tee_key = key;
+                key_exp
+            }
+        };
+
         let attest_response = self
             .http_client()
             .post(format!("{kbs_root_url}{KBS_PREFIX}/attest"))
             .header("Content-Type", "application/json")
-            .json(&self.generate_evidence()?)
+            .json(&self.generate_evidence(pubkey)?)
             .send()
             .await?;
 
@@ -213,7 +235,7 @@ impl KbsRequest for KbsProtocolWrapper {
             log::info!("CC-KBC: trying to request KBS, attempt {attempt}");
 
             if !self.authenticated {
-                self.attestation(root_url.clone()).await?;
+                self.attestation(root_url.clone(), None).await?;
             }
 
             let res = self.http_client().get(&url_string).send().await?;
@@ -243,19 +265,49 @@ impl KbsRequest for KbsProtocolWrapper {
         bail!("Request KBS: Attested but KBS still return Unauthorized")
     }
 
-    async fn attest(&mut self, host_url: String) -> Result<String> {
-        self.attestation(host_url).await
+    async fn attest(&mut self, url: String, tee_pubkey_pem: Option<String>) -> Result<String> {
+        self.attestation(url, tee_pubkey_pem).await
     }
 }
 
-fn build_http_client() -> Result<reqwest::Client> {
-    reqwest::Client::builder()
+fn build_http_client(kbs_root_certs_pem: Vec<String>) -> Result<reqwest::Client> {
+    let mut client_builder = reqwest::Client::builder()
         .cookie_store(true)
         .user_agent(format!(
             "attestation-agent-kbs-client/{}",
             env!("CARGO_PKG_VERSION")
         ))
-        .timeout(Duration::from_secs(KBS_REQ_TIMEOUT_SEC))
+        .timeout(Duration::from_secs(KBS_REQ_TIMEOUT_SEC));
+
+    for custom_root_cert in kbs_root_certs_pem.iter() {
+        let cert = reqwest::Certificate::from_pem(custom_root_cert.as_bytes())?;
+        client_builder = client_builder.add_root_certificate(cert);
+    }
+
+    client_builder
         .build()
         .map_err(|e| anyhow!("Build KBS http client failed: {:?}", e))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_build_http_client() {
+        let test_certs = vec!["-----BEGIN CERTIFICATE-----
+MIIBzTCCAX+gAwIBAgIUOGdGRmt/IDSVIem7iFwsuxnV62AwBQYDK2VwMGkxCzAJ
+BgNVBAYTAkNOMREwDwYDVQQIDAhTaGFuZ2hhaTERMA8GA1UEBwwIU2hhbmdoYWkx
+EDAOBgNVBAoMB0FsaWJhYmExDzANBgNVBAsMBkFsaXl1bjERMA8GA1UEAwwIS0JT
+LXJvb3QwHhcNMjMwNzE0MDYzMzA1WhcNMjMwODEzMDYzMzA1WjBpMQswCQYDVQQG
+EwJDTjERMA8GA1UECAwIU2hhbmdoYWkxETAPBgNVBAcMCFNoYW5naGFpMRAwDgYD
+VQQKDAdBbGliYWJhMQ8wDQYDVQQLDAZBbGl5dW4xETAPBgNVBAMMCEtCUy1yb290
+MCowBQYDK2VwAyEAOo8z6/Ul3XvNBf2Oa7qDevljyhGSKyGMjV+4qneVNr+jOTA3
+MAkGA1UdEwQCMAAwCwYDVR0PBAQDAgXgMB0GA1UdDgQWBBREKNLFRe7fCBKRffTv
+x13TMfDeczAFBgMrZXADQQBpP6ABBkzVj3mF55nWUtP5vxwq3t91wqQJ6NyC7WsT
+3Z29bFfJn7C280JfkCqiqeSZjYV/JjTepATH659kktcA
+-----END CERTIFICATE-----"
+            .to_string()];
+        assert!(build_http_client(test_certs).is_ok());
+    }
+}