Feat: Implement TPM AK credential fetching

attestation-agent/attestation-agent/Cargo.toml

@@ -34,7 +34,9 @@ kbs-types.workspace = true
 log.workspace = true
 prost = { workspace = true, optional = true }
 protobuf = { workspace = true, optional = true }
-reqwest = { workspace = true, features = ["json"], optional = true }
+reqwest = { workspace = true, features = ["json", "native-tls"], optional = true }
+rsa = { version = "0.9", features = ["pem"], optional = true }
+pkcs8 = { version = "0.10", optional = true }
 serde.workspace = true
 serde_json.workspace = true
 sha2.workspace = true
@@ -45,6 +47,8 @@ tokio = { workspace = true, features = ["fs", "sync"] }
 toml.workspace = true
 tonic = { workspace = true, optional = true }
 ttrpc = { workspace = true, features = ["async"], optional = true }
+urlencoding = { version = "2.1", optional = true }
+jsonwebtoken = { version = "9.2", optional = true }
 
 [dev-dependencies]
 rstest.workspace = true
@@ -69,6 +73,8 @@ coco_as = ["reqwest", "token"]
 
 # AA Instance Info support
 instance_info = ["reqwest"]
+# TPM Credential Token
+tpm_credential = ["reqwest", "rsa", "pkcs8", "urlencoding", "token", "tpm-attester", "jsonwebtoken", "kbs"]
 
 all-attesters = [
     "tdx-attester",

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

@@ -189,6 +189,14 @@ impl AttestationAPIs for AttestationAgent {
                 .get_token()
                 .await
             }
+            #[cfg(feature = "tpm_credential")]
+            token::TokenType::TpmCredential => {
+                token::tpm_credential::TpmCredentialGetter::new(
+                    &self.config.read().await.token_configs.kbs,
+                )
+                .get_token()
+                .await
+            }
         }
     }
 

attestation-agent/attestation-agent/src/token/mod.rs

@@ -13,6 +13,9 @@ pub mod kbs;
 #[cfg(feature = "coco_as")]
 pub mod coco_as;
 
+#[cfg(feature = "tpm_credential")]
+pub mod tpm_credential;
+
 #[derive(EnumString, Clone, Copy)]
 pub enum TokenType {
     #[cfg(feature = "kbs")]
@@ -22,6 +25,10 @@ pub enum TokenType {
     #[cfg(feature = "coco_as")]
     #[strum(serialize = "coco_as")]
     CoCoAS,
+
+    #[cfg(feature = "tpm_credential")]
+    #[strum(serialize = "tpm_credential")]
+    TpmCredential,
 }
 
 #[async_trait]

attestation-agent/attestation-agent/src/token/tpm_credential.rs

@@ -0,0 +1,92 @@
+// Copyright (c) 2024 Alibaba Cloud
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+use super::GetToken;
+use crate::config::kbs::KbsConfig;
+use anyhow::*;
+use async_trait::async_trait;
+use attester::tpm::utils;
+use reqwest::Client;
+use rsa::pkcs8::EncodePublicKey;
+use serde::Serialize;
+use std::fs;
+use std::path::Path;
+
+const DEFAULT_AK_CERT_PATH: &str = "/opt/tpm-credential/ak.cert";
+
+#[derive(Serialize)]
+struct TpmCredential {
+    ak_cert_chain: String,
+}
+
+#[derive(Default)]
+pub struct TpmCredentialGetter {
+    kbs_host_url: String,
+    cert: Option<String>,
+}
+
+#[async_trait]
+impl GetToken for TpmCredentialGetter {
+    async fn get_token(&self) -> Result<Vec<u8>> {
+        // Generate AK
+        let _ = utils::generate_rsa_ak()?;
+
+        // Get AK public key
+        let ak_pub = utils::get_ak_pub()?;
+        let ak_pubkey = ak_pub.to_public_key_pem(rsa::pkcs8::LineEnding::LF)?;
+
+        // Try to get EK certificate
+        let ek_cert = utils::dump_ek_cert_pem().ok();
+
+        // Construct request URL and parameters
+        let url = format!("{}/kbs/v0/tpm_pca/ak_credential", self.kbs_host_url);
+        let mut query = format!("ak_pubkey={}", urlencoding::encode(&ak_pubkey));
+
+        // If EK certificate is successfully obtained, add it to query parameters
+        if let Some(ek_cert) = ek_cert {
+            query = format!("{}&ek_cert={}", query, urlencoding::encode(&ek_cert));
+        }
+
+        // Create HTTP client
+        let client = if let Some(cert) = &self.cert {
+            let cert = reqwest::Certificate::from_pem(cert.as_bytes())?;
+            Client::builder().add_root_certificate(cert).build()?
+        } else {
+            Client::new()
+        };
+
+        // Send request
+        let response = client.get(&url).query(&query).send().await?;
+
+        if !response.status().is_success() {
+            bail!("Failed to get AK credential: {}", response.status());
+        }
+
+        let ak_credential = response.text().await?;
+
+        // Save certificate to file
+        let cert_path = Path::new(DEFAULT_AK_CERT_PATH);
+        if let Some(parent) = cert_path.parent() {
+            fs::create_dir_all(parent)?;
+        }
+        fs::write(cert_path, &ak_credential)?;
+
+        let credential = TpmCredential {
+            ak_cert_chain: ak_credential,
+        };
+
+        let res = serde_json::to_vec(&credential)?;
+        Ok(res)
+    }
+}
+
+impl TpmCredentialGetter {
+    pub fn new(config: &KbsConfig) -> Self {
+        Self {
+            kbs_host_url: config.url.clone(),
+            cert: config.cert.clone(),
+        }
+    }
+}

attestation-agent/attester/Cargo.toml

@@ -20,7 +20,7 @@ eventlog-rs = { version = "0.1.7", optional = true }
 hex.workspace = true
 kbs-types.workspace = true
 log.workspace = true
-nix = { workspace = true, optional = true, default-features = false }
+nix = { workspace = true, optional = true }
 num-traits = { version = "0.2.19", optional = true }
 occlum_dcap = { git = "https://github.com/occlum/occlum", tag = "v0.29.7", optional = true }
 pnet = { version = "0.35.0", optional = true }
@@ -32,6 +32,7 @@ scroll = { version = "0.12.0", default-features = false, features = [
 serde.workspace = true
 serde_json.workspace = true
 serde_with.workspace = true
+serial_test = "0.9.0"
 sev = { version = "4.0.0", default-features = false, features = [
     "snp",
 ], optional = true }