AA: Add Config file mechanism

Signed-off-by: Jiale Zhang <zhangjiale@linux.alibaba.com>

attestation-agent/app/src/grpc.rs

@@ -14,7 +14,7 @@ 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()));
+        Arc::new(tokio::sync::Mutex::new(AttestationAgent::default()));
 }
 
 #[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::default().about();
         format!("Protocol: {PROTOCOL}\n{aa_about}")
     };
 }

attestation-agent/app/src/ttrpc.rs

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

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

@@ -0,0 +1,38 @@
+// Copyright (c) 2024 Alibaba Cloud
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+use anyhow::Result;
+use serde::Deserialize;
+use std::fs::File;
+use thiserror::Error;
+
+pub mod aa_kbc_params;
+
+pub const DEFAULT_AA_CONFIG_PATH: &str = "/etc/attestation-agent.toml";
+
+#[derive(Clone, Debug, Deserialize)]
+#[allow(dead_code)]
+pub struct Config {
+    /// URL Address of Attestation Service
+    pub as_uri: String,
+    // TODO: Add more fields that accessing AS needs.
+}
+
+#[derive(Error, Debug)]
+pub enum ConfigFileError {
+    #[error("failed to open")]
+    Io(#[from] std::io::Error),
+    #[error("failed to parse")]
+    Parse(#[from] serde_json::Error),
+}
+
+impl TryFrom<&str> for Config {
+    type Error = ConfigFileError;
+    fn try_from(config_path: &str) -> Result<Self, Self::Error> {
+        let file = File::open(config_path)?;
+        let cfg: Config = serde_json::from_reader(file)?;
+        Ok(cfg)
+    }
+}

attestation-agent/lib/src/lib.rs

@@ -14,12 +14,13 @@ use kbc::{AnnotationPacket, KbcCheckInfo, KbcInstance, KbcModuleList};
 use resource_uri::ResourceUri;
 use std::collections::HashMap;
 
+pub mod config;
 mod token;
 
 #[allow(unused_imports)]
 use token::{GetToken, TokenType};
 
-pub mod aa_kbc_params;
+use crate::config::{aa_kbc_params, Config};
 
 /// 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
@@ -33,7 +34,7 @@ pub mod aa_kbc_params;
 /// use attestation_agent::AttestationAgent;
 /// use attestation_agent::AttestationAPIs;
 ///
-/// let mut aa = AttestationAgent::new();
+/// let mut aa = AttestationAgent::default();
 ///
 /// let key_result = aa.decrypt_image_layer_annotation(
 ///     "sample_kbc",
@@ -94,20 +95,29 @@ pub trait AttestationAPIs {
 pub struct AttestationAgent {
     kbc_module_list: KbcModuleList,
     kbc_instance_map: HashMap<String, KbcInstance>,
+    config: Option<Config>,
 }
 
 impl Default for AttestationAgent {
     fn default() -> Self {
-        Self::new()
+        let config = Config::try_from(config::DEFAULT_AA_CONFIG_PATH).ok();
+        AttestationAgent {
+            kbc_module_list: KbcModuleList::new(),
+            kbc_instance_map: HashMap::new(),
+            config,
+        }
     }
 }
 
 impl AttestationAgent {
     /// Create a new instance of [AttestationAgent].
-    pub fn new() -> Self {
+    pub fn new(config_path: &str) -> Self {
+        let config = Config::try_from(config_path).ok();
+
         AttestationAgent {
             kbc_module_list: KbcModuleList::new(),
             kbc_instance_map: HashMap::new(),
+            config,
         }
     }
 
@@ -178,21 +188,29 @@ impl AttestationAPIs for AttestationAgent {
 
     #[allow(unreachable_code)]
     async fn get_token(&mut self, _token_type: &str) -> Result<Vec<u8>> {
-        let _params = aa_kbc_params::get_params().await?;
+        let _uri = match self.config.as_ref() {
+            Some(c) => c.as_uri.clone(),
+            None => {
+                let params = aa_kbc_params::get_params()
+                    .await
+                    .map_err(|_| anyhow!("Get AS URI failed"))?;
+                params.uri().to_string()
+            }
+        };
 
         let _token = match serde_json::from_str::<TokenType>(_token_type)
             .map_err(|e| anyhow!("Unsupported token type: {e}"))?
         {
             #[cfg(feature = "kbs")]
             TokenType::Kbs => {
                 token::kbs::KbsTokenGetter::default()
-                    .get_token(_params.uri().to_string())
+                    .get_token(_uri)
                     .await?
             }
             #[cfg(feature = "coco_as")]
             TokenType::CoCoAS => {
                 token::coco_as::CoCoASTokenGetter::default()
-                    .get_token(_params.uri().to_string())
+                    .get_token(_uri)
                     .await?
             }
         };

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

@@ -14,7 +14,7 @@ pub struct CoCoASTokenGetter {}
 
 #[async_trait]
 impl GetToken for CoCoASTokenGetter {
-    async fn get_token(&self, as_url: String) -> Result<Vec<u8>> {
+    async fn get_token(&self, as_uri: String) -> Result<Vec<u8>> {
         let tee_type = attester::detect_tee_type();
         let attester = attester::BoxedAttester::try_from(tee_type)?;
         let evidence = attester.get_evidence(vec![]).await?;
@@ -26,7 +26,7 @@ impl GetToken for CoCoASTokenGetter {
 
         let client = reqwest::Client::new();
         let res = client
-            .post(as_url)
+            .post(as_uri)
             .header("Content-Type", "application/json")
             .json(&request_body)
             .send()

confidential-data-hub/hub/src/auth/mod.rs

@@ -16,7 +16,7 @@ impl Hub {
         #[cfg(feature = "sev")]
         {
             use log::{info, warn};
-            match attestation_agent::aa_kbc_params::get_params().await {
+            match attestation_agent::config::aa_kbc_params::get_params().await {
                 Ok(aa_kbc_params) => {
                     if aa_kbc_params.kbc() == "online_sev_kbc" {
                         info!("online_sev_kbc used. Start to initialize sev.");

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

@@ -3,7 +3,7 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-use attestation_agent::aa_kbc_params;
+use attestation_agent::config::aa_kbc_params;
 use thiserror::Error;
 
 pub type Result<T> = std::result::Result<T, Error>;

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

@@ -16,7 +16,7 @@ mod offline_fs;
 use std::sync::Arc;
 
 use async_trait::async_trait;
-use attestation_agent::aa_kbc_params;
+use attestation_agent::config::aa_kbc_params;
 use lazy_static::lazy_static;
 pub use resource_uri::ResourceUri;
 use tokio::sync::Mutex;