aa_kbc_params: centralize handling in CDH and AA

attestation-agent/kbc/src/lib.rs

@@ -81,6 +81,7 @@ impl Default for KbcModuleList {
 impl KbcModuleList {
     /// Create a new [KbcModuleList] and register all known KBC modules.
     pub fn new() -> KbcModuleList {
+        #[allow(unused_mut)]
         let mut mod_list = HashMap::new();
 
         #[cfg(feature = "sample_kbc")]

attestation-agent/lib/Cargo.toml

@@ -16,6 +16,7 @@ resource_uri.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 strum.workspace = true
+thiserror.workspace = true
 tokio = { workspace = true, features = ["fs"] }
 toml.workspace = true
 tonic = { workspace = true, optional = true }

attestation-agent/lib/src/aa_kbc_params.rs

@@ -0,0 +1,105 @@
+use log::debug;
+use serde::Deserialize;
+use std::convert::TryFrom;
+use std::env;
+use std::path::Path;
+use std::sync::OnceLock;
+use thiserror::Error;
+use tokio::fs;
+
+const PEER_POD_CONFIG_PATH: &str = "/run/peerpod/daemon.json";
+static KATA_AGENT_CONFIG_PATH: OnceLock<String> = OnceLock::new();
+
+#[derive(Error, Debug)]
+pub enum ParamError {
+    #[error("illegal aa_kbc_params format: {0}")]
+    IllegalFormat(String),
+    #[error("unable to read `aa_kbc_params` entry from kata-agent config file")]
+    AgentConfigParsing(#[from] toml::de::Error),
+    #[error("io error")]
+    Io(#[from] std::io::Error),
+    #[error("no `agent.aa_kbc_params` provided in kernel commandline")]
+    MissingInCmdline,
+}
+
+pub struct AaKbcParams {
+    kbc: String,
+    uri: String,
+}
+
+impl AaKbcParams {
+    pub fn kbc(&self) -> &str {
+        &self.kbc
+    }
+
+    pub fn uri(&self) -> &str {
+        &self.uri
+    }
+}
+
+impl TryFrom<String> for AaKbcParams {
+    type Error = ParamError;
+
+    fn try_from(value: String) -> Result<Self, ParamError> {
+        let segments: Vec<&str> = value.split("::").collect();
+
+        if segments.len() != 2 {
+            return Err(ParamError::IllegalFormat(value));
+        }
+
+        let params = AaKbcParams {
+            kbc: segments[0].into(),
+            uri: segments[1].into(),
+        };
+
+        Ok(params)
+    }
+}
+
+async fn get_value() -> Result<String, ParamError> {
+    // first check whether we are in a peer pod
+    if Path::new(PEER_POD_CONFIG_PATH).exists() {
+        return from_config_file().await;
+    }
+    // finally use the kernel cmdline
+    from_cmdline().await
+}
+
+pub async fn get_params() -> Result<AaKbcParams, ParamError> {
+    let value = get_value().await?;
+    value.try_into()
+}
+
+// We only care about the aa_kbc_params value at the moment
+#[derive(Debug, Deserialize)]
+struct AgentConfig {
+    aa_kbc_params: String,
+}
+
+async fn from_config_file() -> Result<String, ParamError> {
+    debug!("get aa_kbc_params from file");
+
+    // check env for KATA_AGENT_CONFIG_PATH, fall back to default path
+    let path: &String = KATA_AGENT_CONFIG_PATH.get_or_init(|| {
+        env::var("KATA_AGENT_CONFIG_PATH").unwrap_or_else(|_| "/etc/agent-config.toml".into())
+    });
+
+    debug!("reading agent config from {}", path);
+    let agent_config_str = std::fs::read_to_string(path)?;
+
+    let agent_config: AgentConfig = toml::from_str(&agent_config_str)?;
+
+    Ok(agent_config.aa_kbc_params)
+}
+
+async fn from_cmdline() -> Result<String, ParamError> {
+    debug!("get aa_kbc_params from kernel cmdline");
+    let cmdline = fs::read_to_string("/proc/cmdline").await?;
+    let value = cmdline
+        .split_ascii_whitespace()
+        .find(|para| para.starts_with("agent.aa_kbc_params="))
+        .ok_or(ParamError::MissingInCmdline)?
+        .strip_prefix("agent.aa_kbc_params=")
+        .expect("must have a prefix");
+    Ok(value.into())
+}

attestation-agent/lib/src/lib.rs

@@ -19,6 +19,8 @@ mod token;
 #[cfg(feature = "cc_kbc")]
 use token::get_kbs_token;
 
+pub mod aa_kbc_params;
+
 /// 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
 /// requests to the Relying Party (Key Broker Service) in Confidential Containers,

attestation-agent/lib/src/token.rs

@@ -3,39 +3,25 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-use anyhow::{anyhow, Context, Result};
+use crate::aa_kbc_params;
+use anyhow::Result;
 use kbs_protocol::{evidence_provider::NativeEvidenceProvider, KbsClientBuilder};
-use log::debug;
-use serde::{Deserialize, Serialize};
-use std::env;
-use std::path::Path;
-use std::sync::OnceLock;
-use tokio::fs;
-
-const PEER_POD_CONFIG_PATH: &str = "/run/peerpod/daemon.json";
+use serde::Serialize;
 
 #[derive(Serialize)]
 struct Message {
     token: String,
     tee_keypair: String,
 }
 
-static KATA_AGENT_CONFIG_PATH: OnceLock<String> = OnceLock::new();
-
 pub(crate) async fn get_kbs_token() -> Result<Vec<u8>> {
     let evidence_provider = Box::new(NativeEvidenceProvider::new()?);
 
-    // 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 params = aa_kbc_params::get_params().await?;
+    let kbs_host_url = params.uri();
 
     let mut client =
-        KbsClientBuilder::with_evidence_provider(evidence_provider, &kbs_host_url).build()?;
+        KbsClientBuilder::with_evidence_provider(evidence_provider, kbs_host_url).build()?;
 
     let (token, tee_keypair) = client.get_token().await?;
     let message = Message {
@@ -46,52 +32,3 @@ pub(crate) async fn get_kbs_token() -> Result<Vec<u8>> {
     let res = serde_json::to_vec(&message)?;
     Ok(res)
 }
-
-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> {
-    let cmdline = fs::read_to_string("/proc/cmdline").await?;
-    let kbc_params = cmdline
-        .split_ascii_whitespace()
-        .find(|para| para.starts_with("agent.aa_kbc_params="))
-        .ok_or(anyhow!(
-            "no `agent.aa_kbc_params` provided in kernel commandline!",
-        ))?
-        .strip_prefix("agent.aa_kbc_params=")
-        .expect("must have one")
-        .to_string();
-    Ok(kbc_params)
-}
-
-pub(crate) async fn get_kbc_params_from_config_file() -> Result<String> {
-    // We only care about the aa_kbc_params value at the moment
-    #[derive(Debug, Deserialize)]
-    struct AgentConfig {
-        aa_kbc_params: Option<String>,
-    }
-
-    // check env for KATA_AGENT_CONFIG_PATH, fall back to default path
-    let path: &String = KATA_AGENT_CONFIG_PATH.get_or_init(|| {
-        env::var("KATA_AGENT_CONFIG_PATH").unwrap_or_else(|_| "/etc/agent-config.toml".into())
-    });
-
-    debug!("reading agent config from {}", path);
-    let agent_config_str = fs::read_to_string(path)
-        .await
-        .context(format!("Failed to read {path}"))?;
-
-    let agent_config: AgentConfig =
-        toml::from_str(&agent_config_str).context(format!("Failed to deserialize {path}"))?;
-
-    agent_config
-        .aa_kbc_params
-        .ok_or(anyhow!("no `aa_kbc_params` found in {path}!"))
-}

confidential-data-hub/kms/Cargo.toml

@@ -8,6 +8,7 @@ edition = "2021"
 [dependencies]
 anyhow.workspace = true
 async-trait.workspace = true
+attestation_agent = { path = "../../attestation-agent/lib", default-features = false }
 base64.workspace = true
 bincode = { workspace = true, optional = true }
 chrono = { workspace = true, optional = true }

confidential-data-hub/kms/src/error.rs

@@ -3,6 +3,7 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
+use attestation_agent::aa_kbc_params;
 use thiserror::Error;
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -22,4 +23,7 @@ pub enum Error {
 
     #[error("Unsupported provider: {0}")]
     UnsupportedProvider(String),
+
+    #[error("aa_kbc_params error")]
+    AaKbcParamsError(#[from] aa_kbc_params::ParamError),
 }

confidential-data-hub/kms/src/plugins/kbs/mod.rs

@@ -16,21 +16,13 @@ mod offline_fs;
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use attestation_agent::aa_kbc_params;
 use lazy_static::lazy_static;
-use log::debug;
 pub use resource_uri::ResourceUri;
-use serde::Deserialize;
-use std::path::Path;
-use std::sync::OnceLock;
-use std::{env, fs};
 use tokio::sync::Mutex;
 
 use crate::{Annotations, Error, Getter, Result};
 
-const PEER_POD_CONFIG_PATH: &str = "/run/peerpod/daemon.json";
-
-static KATA_AGENT_CONFIG_PATH: OnceLock<String> = OnceLock::new();
-
 enum RealClient {
     #[cfg(feature = "kbs")]
     Cc(cc_kbc::CcKbc),
@@ -41,18 +33,13 @@ enum RealClient {
 
 impl RealClient {
     async fn new() -> Result<Self> {
-        // 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, _kbs_host) = match Path::new(PEER_POD_CONFIG_PATH).exists() {
-            true => get_aa_params_from_config_file().await?,
-            false => get_aa_params_from_cmdline().await?,
-        };
+        let params = aa_kbc_params::get_params().await?;
 
-        let c = match &kbc[..] {
+        let c = match params.kbc() {
             #[cfg(feature = "kbs")]
-            "cc_kbc" => RealClient::Cc(cc_kbc::CcKbc::new(&_kbs_host).await?),
+            "cc_kbc" => RealClient::Cc(cc_kbc::CcKbc::new(params.uri()).await?),
             #[cfg(feature = "sev")]
-            "online_sev_kbc" => RealClient::Sev(sev::OnlineSevKbc::new(&_kbs_host).await?),
+            "online_sev_kbc" => RealClient::Sev(sev::OnlineSevKbc::new(params.uri()).await?),
             "offline_fs_kbc" => RealClient::OfflineFs(offline_fs::OfflineFsKbc::new().await?),
             others => return Err(Error::KbsClientError(format!("unknown kbc name {others}, only support `cc_kbc`(feature `kbs`), `online_sev_kbc` (feature `sev`) and `offline_fs_kbc`."))),
         };
@@ -116,69 +103,3 @@ impl KbcClient {
         Ok(KbcClient {})
     }
 }
-
-async fn get_aa_params_from_cmdline() -> Result<(String, String)> {
-    use tokio::fs;
-    debug!("get aa_kbc_params from kernel cmdline");
-    let cmdline = fs::read_to_string("/proc/cmdline")
-        .await
-        .map_err(|e| Error::KbsClientError(format!("read kernel cmdline failed: {e}")))?;
-    let aa_kbc_params = cmdline
-        .split_ascii_whitespace()
-        .find(|para| para.starts_with("agent.aa_kbc_params="))
-        .ok_or(Error::KbsClientError(
-            "no `agent.aa_kbc_params` provided in kernel commandline!".into(),
-        ))?
-        .strip_prefix("agent.aa_kbc_params=")
-        .expect("must have a prefix")
-        .split("::")
-        .collect::<Vec<&str>>();
-
-    if aa_kbc_params.len() != 2 {
-        return Err(Error::KbsClientError(
-            "Illegal `agent.aa_kbc_params` format provided in kernel commandline.".to_string(),
-        ));
-    }
-
-    Ok((aa_kbc_params[0].to_string(), aa_kbc_params[1].to_string()))
-}
-
-async fn get_aa_params_from_config_file() -> Result<(String, String)> {
-    debug!("get aa_kbc_params from file");
-    // We only care about the aa_kbc_params value at the moment
-    #[derive(Debug, Deserialize)]
-    struct AgentConfig {
-        aa_kbc_params: Option<String>,
-    }
-
-    // check env for KATA_AGENT_CONFIG_PATH, fall back to default path
-    let path: &String = KATA_AGENT_CONFIG_PATH.get_or_init(|| {
-        env::var("KATA_AGENT_CONFIG_PATH").unwrap_or_else(|_| "/etc/agent-config.toml".into())
-    });
-
-    debug!("reading agent config from {}", path);
-    let agent_config_str = fs::read_to_string(path)
-        .map_err(|e| Error::KbsClientError(format!("Failed to read {path} file: {e}")))?;
-
-    let agent_config: AgentConfig = toml::from_str(&agent_config_str)
-        .map_err(|e| Error::KbsClientError(format!("Failed to deserialize {path}: {e}")))?;
-
-    let aa_kbc_params = agent_config
-        .aa_kbc_params
-        .ok_or(Error::KbsClientError(format!(
-            "no `aa_kbc_params` found in {path}"
-        )))?;
-
-    let aa_kbc_params_vec = aa_kbc_params.split("::").collect::<Vec<&str>>();
-
-    if aa_kbc_params_vec.len() != 2 {
-        return Err(Error::KbsClientError(format!(
-            "Illegal `aa_kbc_params` format provided in {path}."
-        )));
-    }
-
-    Ok((
-        aa_kbc_params_vec[0].to_string(),
-        aa_kbc_params_vec[1].to_string(),
-    ))
-}