From 974d6b0736c4e758a2be2b6ab7dd521cdea82aff Mon Sep 17 00:00:00 2001
From: ChengyuZhu6 <chengyu.zhu@intel.com>
Date: Mon, 30 Sep 2024 12:51:31 +0800
Subject: [PATCH] agent:cdh: initialize cdhclient with the input cdh socket uri

Refactor cdh code to initialize cdhclient with the input cdh socket uri.

Signed-off-by: ChengyuZhu6 <chengyu.zhu@intel.com>
---
 src/agent/src/cdh.rs  | 11 ++++++-----
 src/agent/src/main.rs |  2 +-
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/agent/src/cdh.rs b/src/agent/src/cdh.rs
index 551b92740b..5a1250b14a 100644
--- a/src/agent/src/cdh.rs
+++ b/src/agent/src/cdh.rs
@@ -8,7 +8,6 @@
 // https://github.com/confidential-containers/guest-components/tree/main/confidential-data-hub
 
 use crate::AGENT_CONFIG;
-use crate::CDH_SOCKET_URI;
 use anyhow::{bail, Context, Result};
 use derivative::Derivative;
 use protocols::{
@@ -43,8 +42,8 @@ pub struct CDHClient {
 }
 
 impl CDHClient {
-    pub fn new() -> Result<Self> {
-        let client = ttrpc::asynchronous::Client::connect(CDH_SOCKET_URI)?;
+    pub fn new(cdh_socket_uri: &str) -> Result<Self> {
+        let client = ttrpc::asynchronous::Client::connect(cdh_socket_uri)?;
         let sealed_secret_client =
             confidential_data_hub_ttrpc_async::SealedSecretServiceClient::new(client.clone());
         let secure_mount_client =
@@ -87,9 +86,11 @@ impl CDHClient {
     }
 }
 
-pub async fn init_cdh_client() -> Result<()> {
+pub async fn init_cdh_client(cdh_socket_uri: &str) -> Result<()> {
     CDH_CLIENT
-        .get_or_try_init(|| async { CDHClient::new().context("Failed to create CDH Client") })
+        .get_or_try_init(|| async {
+            CDHClient::new(cdh_socket_uri).context("Failed to create CDH Client")
+        })
         .await?;
     Ok(())
 }
diff --git a/src/agent/src/main.rs b/src/agent/src/main.rs
index 88a5ec01a5..8a057bb367 100644
--- a/src/agent/src/main.rs
+++ b/src/agent/src/main.rs
@@ -498,7 +498,7 @@ async fn init_attestation_components(logger: &Logger, config: &AgentConfig) -> R
     .map_err(|e| anyhow!("launch_process {} failed: {:?}", CDH_PATH, e))?;
 
     // initialize cdh client
-    cdh::init_cdh_client().await?;
+    cdh::init_cdh_client(CDH_SOCKET_URI).await?;
 
     // skip launch of api-server-rest
     if config.guest_components_procs == GuestComponentsProcs::ConfidentialDataHub {