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>
confidential-containers · Jul 18, 2023 · e26aa2c · e26aa2c
1 parent f3d9264
commit e26aa2c
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 7 deletions.
diff --git 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
@@ -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))
 }