From 184fff8328abb792eaca1a8db309382c63a30592 Mon Sep 17 00:00:00 2001
From: Jiale Zhang <zhangjiale@linux.alibaba.com>
Date: Tue, 11 Jul 2023 14:33:11 +0800
Subject: [PATCH] AA: Allow setting custom KBS certificates to enable TLS

Add parameters for setting customized KBS Root certificate when creating HTTP client,
so that TLS communication can be enabled in various deployment scenarios.

Signed-off-by: Jiale Zhang <zhangjiale@linux.alibaba.com>
---
 attestation-agent/kbc/src/cc_kbc/mod.rs   |  2 +-
 attestation-agent/kbs_protocol/src/lib.rs | 19 +++++++++++++------
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/attestation-agent/kbc/src/cc_kbc/mod.rs b/attestation-agent/kbc/src/cc_kbc/mod.rs
index 26199fd1d..99b88a53a 100644
--- a/attestation-agent/kbc/src/cc_kbc/mod.rs
+++ b/attestation-agent/kbc/src/cc_kbc/mod.rs
@@ -61,7 +61,7 @@ impl Kbc {
         Ok(Kbc {
             kbs_uri: url,
             token: None,
-            kbs_protocol_wrapper: KbsProtocolWrapper::new().unwrap(),
+            kbs_protocol_wrapper: KbsProtocolWrapper::new(vec![]).unwrap(),
         })
     }
 
diff --git a/attestation-agent/kbs_protocol/src/lib.rs b/attestation-agent/kbs_protocol/src/lib.rs
index bf3c70fca..9df9f09b9 100644
--- a/attestation-agent/kbs_protocol/src/lib.rs
+++ b/attestation-agent/kbs_protocol/src/lib.rs
@@ -37,7 +37,7 @@ pub struct KbsProtocolWrapper {
 }
 
 impl KbsProtocolWrapper {
-    pub fn new() -> Result<KbsProtocolWrapper> {
+    pub fn new(kbs_root_certs_pem: Vec<String>) -> Result<KbsProtocolWrapper> {
         // Detect TEE type of the current platform.
         let tee_type = detect_tee_type();
         // Create attester instance.
@@ -48,7 +48,7 @@ impl KbsProtocolWrapper {
             tee_key: None,
             nonce: String::default(),
             attester,
-            http_client: build_http_client().unwrap(),
+            http_client: build_http_client(kbs_root_certs_pem).unwrap(),
             authenticated: false,
         })
     }
@@ -87,7 +87,7 @@ impl KbsProtocolWrapper {
 
     async fn attestation(
         &mut self,
-        kbs_root_url: String
+        kbs_root_url: String,
         tee_pubkey_pem: Option<String>,
     ) -> Result<String> {
         let challenge = self
@@ -191,14 +191,21 @@ impl KbsRequest for KbsProtocolWrapper {
     }
 }
 
-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))
 }