AA: Support CoCo-AS Attestation Token

Cargo.toml

@@ -7,6 +7,7 @@ members = [
     "attestation-agent/kbc",
     "attestation-agent/kbs_protocol",
     "attestation-agent/attester",
+    "attestation-agent/token",
     "attestation-agent/deps/resource_uri",
     "attestation-agent/deps/crypto",
     "attestation-agent/deps/sev",

attestation-agent/app/src/grpc.rs

@@ -13,8 +13,10 @@ const DEFAULT_GETRESOURCE_ADDR: &str = "127.0.0.1:50001";
 const DEFAULT_ATTESTATION_AGENT_ADDR: &str = "127.0.0.1:50002";
 
 lazy_static! {
-    pub static ref ASYNC_ATTESTATION_AGENT: Arc<tokio::sync::Mutex<AttestationAgent>> =
-        Arc::new(tokio::sync::Mutex::new(AttestationAgent::new()));
+    pub static ref ASYNC_ATTESTATION_AGENT: Arc<tokio::sync::Mutex<AttestationAgent<'static>>> =
+        Arc::new(tokio::sync::Mutex::new(AttestationAgent::<'static>::new(
+            None
+        )));
 }
 
 #[derive(Debug, Parser)]

attestation-agent/app/src/rpc/mod.rs

@@ -21,7 +21,7 @@ const PROTOCOL: &str = "grpc";
 
 lazy_static! {
     pub static ref ABOUT: String = {
-        let aa_about = AttestationAgent::new().about();
+        let aa_about = AttestationAgent::new(None).about();
         format!("Protocol: {PROTOCOL}\n{aa_about}")
     };
 }

attestation-agent/app/src/ttrpc.rs

@@ -32,8 +32,8 @@ const DEFAULT_ATTESTATION_SOCKET_ADDR: &str = concatcp!(
 );
 
 lazy_static! {
-    pub static ref ASYNC_ATTESTATION_AGENT: Arc<Mutex<AttestationAgent>> =
-        Arc::new(Mutex::new(AttestationAgent::new()));
+    pub static ref ASYNC_ATTESTATION_AGENT: Arc<Mutex<AttestationAgent<'static>>> =
+        Arc::new(Mutex::new(AttestationAgent::<'static>::new(None)));
 }
 
 #[derive(Debug, Parser)]

attestation-agent/lib/Cargo.toml

@@ -16,6 +16,7 @@ resource_uri.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 strum.workspace = true
+token = { path = "../token", optional = true }
 tokio = { workspace = true, features = ["fs"] }
 toml.workspace = true
 tonic = { workspace = true, optional = true }
@@ -24,15 +25,19 @@ tonic = { workspace = true, optional = true }
 tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
 
 [features]
-default = ["sample_kbc", "rust-crypto"]
-
-cc_kbc = ["kbc/cc_kbc", "kbs_protocol/background_check"]
-all-attesters = ["kbc/all-attesters", "kbs_protocol?/all-attesters", "attester/all-attesters"]
-tdx-attester = ["kbc/tdx-attester", "kbs_protocol/tdx-attester", "attester/tdx-attester"]
-sgx-attester = ["kbc/sgx-attester", "kbs_protocol/sgx-attester", "attester/sgx-attester"]
-az-snp-vtpm-attester = ["kbc/az-snp-vtpm-attester", "kbs_protocol/az-snp-vtpm-attester", "attester/az-snp-vtpm-attester"]
-az-tdx-vtpm-attester = ["kbc/az-tdx-vtpm-attester", "kbs_protocol/az-tdx-vtpm-attester", "attester/az-tdx-vtpm-attester"]
-snp-attester = ["kbc/snp-attester", "kbs_protocol/snp-attester", "attester/snp-attester"]
+default = ["sample_kbc", "rust-crypto", "coco_as"]
+
+kbs_as = ["token/kbs"]
+
+coco_as = ["token/coco_as"]
+
+cc_kbc = ["kbc/cc_kbc", "kbs_as", "coco_as"]
+all-attesters = ["kbc/all-attesters", "kbs_protocol?/all-attesters", "attester/all-attesters", "token?/all-attesters"]
+tdx-attester = ["kbc/tdx-attester", "kbs_protocol/tdx-attester", "attester/tdx-attester", "token?/tdx"]
+sgx-attester = ["kbc/sgx-attester", "kbs_protocol/sgx-attester", "attester/sgx-attester", "token?/sgx"]
+az-snp-vtpm-attester = ["kbc/az-snp-vtpm-attester", "kbs_protocol/az-snp-vtpm-attester", "attester/az-snp-vtpm-attester", "token?/az-snp-vtpm"]
+az-tdx-vtpm-attester = ["kbc/az-tdx-vtpm-attester", "kbs_protocol/az-tdx-vtpm-attester", "attester/az-tdx-vtpm-attester", "token?/az-tdx-vtpm"]
+snp-attester = ["kbc/snp-attester", "kbs_protocol/snp-attester", "attester/snp-attester", "token?/snp"]
 
 sample_kbc = ["kbc/sample_kbc"]
 eaa_kbc = ["kbc/eaa_kbc"]

attestation-agent/lib/src/token.rs → attestation-agent/lib/src/config.rs

@@ -4,60 +4,50 @@
 //
 
 use anyhow::{anyhow, Context, Result};
-use kbs_protocol::{evidence_provider::NativeEvidenceProvider, KbsClientBuilder};
 use log::debug;
-use serde::{Deserialize, Serialize};
+use serde::Deserialize;
 use std::env;
+use std::fs::File;
 use std::path::Path;
 use std::sync::OnceLock;
 use tokio::fs;
 
 const PEER_POD_CONFIG_PATH: &str = "/run/peerpod/daemon.json";
+pub const DEFAULT_AA_CONFIG_PATH: &str = "/etc/attestation.toml";
 
-#[derive(Serialize)]
-struct Message {
-    token: String,
-    tee_keypair: String,
+static KATA_AGENT_CONFIG_PATH: OnceLock<String> = OnceLock::new();
+
+#[derive(Clone, Debug, Deserialize)]
+#[allow(dead_code)]
+pub struct Config {
+    /// URL Address of Attestation Service
+    pub as_url: String,
 }
 
-static KATA_AGENT_CONFIG_PATH: OnceLock<String> = OnceLock::new();
+impl TryFrom<&Path> for Config {
+    type Error = anyhow::Error;
+    fn try_from(config_path: &Path) -> Result<Self, Self::Error> {
+        let file = File::open(config_path)
+            .map_err(|e| anyhow!("failed to open AA config file {}", e.to_string()))?;
 
-pub(crate) async fn get_kbs_token() -> Result<Vec<u8>> {
-    let evidence_provider = Box::new(NativeEvidenceProvider::new()?);
+        serde_json::from_reader::<File, Config>(file)
+            .map_err(|e| anyhow!("failed to parse AA config file {}", e.to_string()))
+    }
+}
 
+#[allow(dead_code)]
+pub async fn get_host_url() -> Result<String> {
     // Check for /run/peerpod/daemon.json to see if we are in a peer pod
     // If so we need to read from the agent-config file, not /proc/cmdline
     let kbc_params = match Path::new(PEER_POD_CONFIG_PATH).exists() {
-        true => get_kbc_params_from_config_file().await?,
-        false => get_kbc_params_from_cmdline().await?,
-    };
-
-    let kbs_host_url = extract_kbs_host_url(&kbc_params)?;
-
-    let mut client =
-        KbsClientBuilder::with_evidence_provider(evidence_provider, &kbs_host_url).build()?;
-
-    let (token, tee_keypair) = client.get_token().await?;
-    let message = Message {
-        token: token.content,
-        tee_keypair: tee_keypair.to_pkcs1_pem()?.to_string(),
+        true => get_aa_params_from_kata_agent_config_file().await?,
+        false => get_aa_params_from_kernel_cmdline().await?,
     };
 
-    let res = serde_json::to_vec(&message)?;
-    Ok(res)
+    extract_host_url(&kbc_params)
 }
 
-fn extract_kbs_host_url(kbc_params: &str) -> Result<String> {
-    let kbs_host = kbc_params
-        .split("::")
-        .last()
-        .ok_or(anyhow!("illegal input `agent.aa_kbc_params` format",))?
-        .to_string();
-
-    Ok(kbs_host)
-}
-
-pub(crate) async fn get_kbc_params_from_cmdline() -> Result<String> {
+async fn get_aa_params_from_kernel_cmdline() -> Result<String> {
     let cmdline = fs::read_to_string("/proc/cmdline").await?;
     let kbc_params = cmdline
         .split_ascii_whitespace()
@@ -71,7 +61,7 @@ pub(crate) async fn get_kbc_params_from_cmdline() -> Result<String> {
     Ok(kbc_params)
 }
 
-pub(crate) async fn get_kbc_params_from_config_file() -> Result<String> {
+async fn get_aa_params_from_kata_agent_config_file() -> Result<String> {
     // We only care about the aa_kbc_params value at the moment
     #[derive(Debug, Deserialize)]
     struct AgentConfig {
@@ -95,3 +85,13 @@ pub(crate) async fn get_kbc_params_from_config_file() -> Result<String> {
         .aa_kbc_params
         .ok_or(anyhow!("no `aa_kbc_params` found in {path}!"))
 }
+
+fn extract_host_url(kbc_params: &str) -> Result<String> {
+    let kbs_host = kbc_params
+        .split("::")
+        .last()
+        .ok_or(anyhow!("illegal input `agent.aa_kbc_params` format",))?
+        .to_string();
+
+    Ok(kbs_host)
+}

attestation-agent/lib/src/lib.rs

@@ -12,12 +12,12 @@ use async_trait::async_trait;
 use attester::{detect_tee_type, BoxedAttester};
 use kbc::{AnnotationPacket, KbcCheckInfo, KbcInstance, KbcModuleList};
 use resource_uri::ResourceUri;
-use std::collections::HashMap;
+use std::{collections::HashMap, path::Path};
 
-#[cfg(feature = "cc_kbc")]
-mod token;
-#[cfg(feature = "cc_kbc")]
-use token::get_kbs_token;
+mod config;
+
+#[cfg(any(feature = "coco_as", feature = "kbs_as"))]
+use token::GetToken;
 
 /// Attestation Agent (AA for short) is a rust library crate for attestation procedure
 /// in confidential containers. It provides kinds of service APIs that need to make
@@ -31,7 +31,7 @@ use token::get_kbs_token;
 /// use attestation_agent::AttestationAgent;
 /// use attestation_agent::AttestationAPIs;
 ///
-/// let mut aa = AttestationAgent::new();
+/// let mut aa = AttestationAgent::new(None);
 ///
 /// let key_result = aa.decrypt_image_layer_annotation(
 ///     "sample_kbc",
@@ -88,22 +88,26 @@ pub trait AttestationAPIs {
     ) -> Result<()>;
 }
 
+#[allow(dead_code)]
 /// Attestation agent to provide attestation service.
-pub struct AttestationAgent {
+pub struct AttestationAgent<'a> {
+    config_file_path: &'a Path,
     kbc_module_list: KbcModuleList,
     kbc_instance_map: HashMap<String, KbcInstance>,
 }
 
-impl Default for AttestationAgent {
+impl<'a> Default for AttestationAgent<'a> {
     fn default() -> Self {
-        Self::new()
+        Self::new(None)
     }
 }
 
-impl AttestationAgent {
+impl<'a> AttestationAgent<'a> {
     /// Create a new instance of [AttestationAgent].
-    pub fn new() -> Self {
+    pub fn new(config_path: Option<&'a str>) -> Self {
+        let config_path = config_path.unwrap_or(config::DEFAULT_AA_CONFIG_PATH);
         AttestationAgent {
+            config_file_path: &Path::new(config_path),
             kbc_module_list: KbcModuleList::new(),
             kbc_instance_map: HashMap::new(),
         }
@@ -135,7 +139,7 @@ impl AttestationAgent {
 }
 
 #[async_trait]
-impl AttestationAPIs for AttestationAgent {
+impl<'a> AttestationAPIs for AttestationAgent<'a> {
     async fn decrypt_image_layer_annotation(
         &mut self,
         kbc_name: &str,
@@ -174,23 +178,36 @@ impl AttestationAPIs for AttestationAgent {
             .await
     }
 
+    #[allow(unused_variables)]
+    #[allow(unreachable_code)]
     async fn get_token(&mut self, _token_type: &str) -> Result<Vec<u8>> {
-        #[cfg(feature = "cc_kbc")]
-        {
-            let token = match _token_type {
-                "kbs" => get_kbs_token().await?,
-                typ => bail!("Unsupported token type {typ}"),
-            };
-
-            Ok(token)
-        }
-
-        // TODO: remove the feature flags after refactoring AA. Currently, kbs_host_url
-        // is only set by user in aa_kbc_params when cc_kbc is enabled.
-        #[cfg(not(feature = "cc_kbc"))]
-        {
-            bail!("unimplemented!");
-        }
+        let token = match _token_type {
+            #[cfg(any(feature = "cc_kbc", feature = "kbs_as"))]
+            "kbs" => {
+                let kbs_host_url = config::get_host_url().await?;
+                let kbs_token = token::kbs::KbsTokenGetter::default()
+                    .get_token(kbs_host_url)
+                    .await?;
+                kbs_token
+            }
+            #[cfg(feature = "coco_as")]
+            "coco_as" => {
+                let as_url = match config::get_host_url().await {
+                    Ok(url) => url,
+                    Err(_) => {
+                        let config = config::Config::try_from(self.config_file_path)?;
+                        config.as_url.clone()
+                    }
+                };
+                let coco_as_token = token::coco_as::CoCoASTokenGetter::default()
+                    .get_token(as_url)
+                    .await?;
+                coco_as_token
+            }
+            typ => bail!("Unsupported token type {typ}"),
+        };
+
+        Ok(token)
     }
 
     /// Get TEE hardware signed evidence that includes the runtime data.

attestation-agent/token/Cargo.toml

@@ -0,0 +1,33 @@
+[package]
+name = "token"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+anyhow.workspace = true
+async-trait.workspace = true
+attester = { path = "../attester", optional = true, default-features = false }
+base64.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+kbs_protocol = { path = "../kbs_protocol", optional = true }
+reqwest = { version = "0.11", features = ["json"], optional = true }
+tokio = { workspace = true, features = ["fs"] }
+toml.workspace = true
+
+[features]
+default = ["kbs", "coco_as"]
+
+kbs = ["kbs_protocol/background_check"]
+coco_as = ["reqwest", "attester"]
+
+all-attesters = ["attester?/all-attesters"]
+tdx = ["attester?/tdx-attester"]
+sgx = ["attester?/sgx-attester"]
+az-snp-vtpm = ["attester?/az-snp-vtpm-attester"]
+az-tdx-vtpm = ["attester?/az-tdx-vtpm-attester"]
+snp = ["attester?/snp-attester"]
+csv = ["attester?/csv-attester"]
+cca = ["attester?/cca-attester"]