CDH: fix ttrpc cdh memory occupation

Fixes confidential-containers#688 This commit will copy only the pointer of CDH to implement different APIs. In old version, each CDH instance will serve for one API, thus occupies multiple times of memory. Signed-off-by: Xynnn007 <[email protected]>
Xynnn007 · Sep 27, 2024 · ee65f67 · ee65f67
1 parent 49476b5
commit ee65f67
Show file tree

Hide file tree

Showing 9 changed files with 53 additions and 86 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -65,7 +65,7 @@ tokio = "1.39"
 toml = "0.8.19"
 tonic = "0.9"
 tonic-build = "0.9"
-ttrpc = "0.8.0"
+ttrpc = "0.8.2"
 ttrpc-codegen = "0.4.2"
 url = "2.5.2"
 uuid = "1"

diff --git a/attestation-agent/kbs_protocol/src/token_provider/aa/attestation_agent_ttrpc.rs b/attestation-agent/kbs_protocol/src/token_provider/aa/attestation_agent_ttrpc.rs
@@ -1,4 +1,4 @@
-// This file is generated by ttrpc-compiler 0.6.2. Do not edit
+// This file is generated by ttrpc-compiler 0.6.3. Do not edit
 // @generated
 
 #![cfg_attr(rustfmt, rustfmt_skip)]
@@ -63,7 +63,7 @@ impl AttestationAgentServiceClient {
 }
 
 struct GetEvidenceMethod {
-    service: Arc<Box<dyn AttestationAgentService + Send + Sync>>,
+    service: Arc<dyn AttestationAgentService + Send + Sync>,
 }
 
 #[async_trait]
@@ -74,7 +74,7 @@ impl ::ttrpc::r#async::MethodHandler for GetEvidenceMethod {
 }
 
 struct GetTokenMethod {
-    service: Arc<Box<dyn AttestationAgentService + Send + Sync>>,
+    service: Arc<dyn AttestationAgentService + Send + Sync>,
 }
 
 #[async_trait]
@@ -85,7 +85,7 @@ impl ::ttrpc::r#async::MethodHandler for GetTokenMethod {
 }
 
 struct ExtendRuntimeMeasurementMethod {
-    service: Arc<Box<dyn AttestationAgentService + Send + Sync>>,
+    service: Arc<dyn AttestationAgentService + Send + Sync>,
 }
 
 #[async_trait]
@@ -96,7 +96,7 @@ impl ::ttrpc::r#async::MethodHandler for ExtendRuntimeMeasurementMethod {
 }
 
 struct CheckInitDataMethod {
-    service: Arc<Box<dyn AttestationAgentService + Send + Sync>>,
+    service: Arc<dyn AttestationAgentService + Send + Sync>,
 }
 
 #[async_trait]
@@ -107,7 +107,7 @@ impl ::ttrpc::r#async::MethodHandler for CheckInitDataMethod {
 }
 
 struct UpdateConfigurationMethod {
-    service: Arc<Box<dyn AttestationAgentService + Send + Sync>>,
+    service: Arc<dyn AttestationAgentService + Send + Sync>,
 }
 
 #[async_trait]
@@ -118,7 +118,7 @@ impl ::ttrpc::r#async::MethodHandler for UpdateConfigurationMethod {
 }
 
 struct GetTeeTypeMethod {
-    service: Arc<Box<dyn AttestationAgentService + Send + Sync>>,
+    service: Arc<dyn AttestationAgentService + Send + Sync>,
 }
 
 #[async_trait]
@@ -150,7 +150,7 @@ pub trait AttestationAgentService: Sync {
     }
 }
 
-pub fn create_attestation_agent_service(service: Arc<Box<dyn AttestationAgentService + Send + Sync>>) -> HashMap<String, ::ttrpc::r#async::Service> {
+pub fn create_attestation_agent_service(service: Arc<dyn AttestationAgentService + Send + Sync>) -> HashMap<String, ::ttrpc::r#async::Service> {
     let mut ret = HashMap::new();
     let mut methods = HashMap::new();
     let streams = HashMap::new();

diff --git a/confidential-data-hub/hub/src/bin/protos/api_ttrpc.rs b/confidential-data-hub/hub/src/bin/protos/api_ttrpc.rs
@@ -1,4 +1,4 @@
-// This file is generated by ttrpc-compiler 0.6.2. Do not edit
+// This file is generated by ttrpc-compiler 0.6.3. Do not edit
 // @generated
 
 #![cfg_attr(rustfmt, rustfmt_skip)]
@@ -38,7 +38,7 @@ impl SealedSecretServiceClient {
 }
 
 struct UnsealSecretMethod {
-    service: Arc<Box<dyn SealedSecretService + Send + Sync>>,
+    service: Arc<dyn SealedSecretService + Send + Sync>,
 }
 
 #[async_trait]
@@ -55,7 +55,7 @@ pub trait SealedSecretService: Sync {
     }
 }
 
-pub fn create_sealed_secret_service(service: Arc<Box<dyn SealedSecretService + Send + Sync>>) -> HashMap<String, ::ttrpc::r#async::Service> {
+pub fn create_sealed_secret_service(service: Arc<dyn SealedSecretService + Send + Sync>) -> HashMap<String, ::ttrpc::r#async::Service> {
     let mut ret = HashMap::new();
     let mut methods = HashMap::new();
     let streams = HashMap::new();
@@ -86,7 +86,7 @@ impl GetResourceServiceClient {
 }
 
 struct GetResourceMethod {
-    service: Arc<Box<dyn GetResourceService + Send + Sync>>,
+    service: Arc<dyn GetResourceService + Send + Sync>,
 }
 
 #[async_trait]
@@ -103,7 +103,7 @@ pub trait GetResourceService: Sync {
     }
 }
 
-pub fn create_get_resource_service(service: Arc<Box<dyn GetResourceService + Send + Sync>>) -> HashMap<String, ::ttrpc::r#async::Service> {
+pub fn create_get_resource_service(service: Arc<dyn GetResourceService + Send + Sync>) -> HashMap<String, ::ttrpc::r#async::Service> {
     let mut ret = HashMap::new();
     let mut methods = HashMap::new();
     let streams = HashMap::new();
@@ -134,7 +134,7 @@ impl SecureMountServiceClient {
 }
 
 struct SecureMountMethod {
-    service: Arc<Box<dyn SecureMountService + Send + Sync>>,
+    service: Arc<dyn SecureMountService + Send + Sync>,
 }
 
 #[async_trait]
@@ -151,7 +151,7 @@ pub trait SecureMountService: Sync {
     }
 }
 
-pub fn create_secure_mount_service(service: Arc<Box<dyn SecureMountService + Send + Sync>>) -> HashMap<String, ::ttrpc::r#async::Service> {
+pub fn create_secure_mount_service(service: Arc<dyn SecureMountService + Send + Sync>) -> HashMap<String, ::ttrpc::r#async::Service> {
     let mut ret = HashMap::new();
     let mut methods = HashMap::new();
     let streams = HashMap::new();
@@ -182,7 +182,7 @@ impl ImagePullServiceClient {
 }
 
 struct PullImageMethod {
-    service: Arc<Box<dyn ImagePullService + Send + Sync>>,
+    service: Arc<dyn ImagePullService + Send + Sync>,
 }
 
 #[async_trait]
@@ -199,7 +199,7 @@ pub trait ImagePullService: Sync {
     }
 }
 
-pub fn create_image_pull_service(service: Arc<Box<dyn ImagePullService + Send + Sync>>) -> HashMap<String, ::ttrpc::r#async::Service> {
+pub fn create_image_pull_service(service: Arc<dyn ImagePullService + Send + Sync>) -> HashMap<String, ::ttrpc::r#async::Service> {
     let mut ret = HashMap::new();
     let mut methods = HashMap::new();
     let streams = HashMap::new();

diff --git a/confidential-data-hub/hub/src/bin/protos/keyprovider_ttrpc.rs b/confidential-data-hub/hub/src/bin/protos/keyprovider_ttrpc.rs
@@ -1,4 +1,4 @@
-// This file is generated by ttrpc-compiler 0.6.2. Do not edit
+// This file is generated by ttrpc-compiler 0.6.3. Do not edit
 // @generated
 
 #![cfg_attr(rustfmt, rustfmt_skip)]
@@ -43,7 +43,7 @@ impl KeyProviderServiceClient {
 }
 
 struct WrapKeyMethod {
-    service: Arc<Box<dyn KeyProviderService + Send + Sync>>,
+    service: Arc<dyn KeyProviderService + Send + Sync>,
 }
 
 #[async_trait]
@@ -54,7 +54,7 @@ impl ::ttrpc::r#async::MethodHandler for WrapKeyMethod {
 }
 
 struct UnWrapKeyMethod {
-    service: Arc<Box<dyn KeyProviderService + Send + Sync>>,
+    service: Arc<dyn KeyProviderService + Send + Sync>,
 }
 
 #[async_trait]
@@ -74,7 +74,7 @@ pub trait KeyProviderService: Sync {
     }
 }
 
-pub fn create_key_provider_service(service: Arc<Box<dyn KeyProviderService + Send + Sync>>) -> HashMap<String, ::ttrpc::r#async::Service> {
+pub fn create_key_provider_service(service: Arc<dyn KeyProviderService + Send + Sync>) -> HashMap<String, ::ttrpc::r#async::Service> {
     let mut ret = HashMap::new();
     let mut methods = HashMap::new();
     let streams = HashMap::new();

diff --git a/confidential-data-hub/hub/src/bin/ttrpc-cdh.rs b/confidential-data-hub/hub/src/bin/ttrpc-cdh.rs
@@ -41,14 +41,6 @@ struct Cli {
     config: Option<String>,
 }
 
-macro_rules! ttrpc_service {
-    ($func: expr, $conf: expr) => {{
-        let server = Server::new($conf).await?;
-        let server = Arc::new(Box::new(server) as _);
-        $func(server)
-    }};
-}
-
 #[tokio::main]
 async fn main() -> Result<()> {
     env_logger::init_from_env(env_logger::Env::new().default_filter_or("info"));
@@ -64,20 +56,17 @@ async fn main() -> Result<()> {
     create_socket_parent_directory(unix_socket_path).await?;
     clean_previous_sock_file(unix_socket_path).await?;
 
-    let sealed_secret_service = ttrpc_service!(create_sealed_secret_service, &config);
-    let get_resource_service = ttrpc_service!(create_get_resource_service, &config);
-    let key_provider_service = ttrpc_service!(create_key_provider_service, &config);
-    let secure_mount_service = ttrpc_service!(create_secure_mount_service, &config);
-    let image_pull_service = ttrpc_service!(create_image_pull_service, &config);
+    let server = Server::new(&config).await.context("create CDH instance")?;
+    let server = Arc::new(server);
 
     let mut server = TtrpcServer::new()
         .bind(&config.socket)
         .context("cannot bind cdh ttrpc service")?
-        .register_service(sealed_secret_service)
-        .register_service(get_resource_service)
-        .register_service(secure_mount_service)
-        .register_service(key_provider_service)
-        .register_service(image_pull_service);
+        .register_service(create_sealed_secret_service(server.clone() as _))
+        .register_service(create_get_resource_service(server.clone() as _))
+        .register_service(create_key_provider_service(server.clone() as _))
+        .register_service(create_secure_mount_service(server.clone() as _))
+        .register_service(create_image_pull_service(server.clone() as _));
 
     info!(
         "[ttRPC] Confidential Data Hub starts to listen to request: {}",

diff --git a/confidential-data-hub/hub/src/bin/ttrpc_server/mod.rs b/confidential-data-hub/hub/src/bin/ttrpc_server/mod.rs
@@ -3,15 +3,13 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-use std::{error::Error as _, sync::Arc};
+use std::error::Error as _;
 
 use anyhow::Result;
 use async_trait::async_trait;
 use confidential_data_hub::{hub::Hub, CdhConfig, DataHub};
-use lazy_static::lazy_static;
 use log::{debug, error};
 use storage::volume_type::Storage;
-use tokio::sync::RwLock;
 use ttrpc::{asynchronous::TtrpcContext, Code, Error, Status};
 
 use crate::{
@@ -30,26 +28,15 @@ use crate::{
     },
 };
 
-lazy_static! {
-    static ref HUB: Arc<RwLock<Option<Hub>>> = Arc::new(RwLock::new(None));
+pub struct Server {
+    hub: Hub,
 }
 
-pub struct Server;
-
 impl Server {
-    async fn init(config: &CdhConfig) -> Result<()> {
-        let mut writer = HUB.write().await;
-        if writer.is_none() {
-            let hub = Hub::new(config.clone()).await?;
-            *writer = Some(hub);
-        }
-
-        Ok(())
-    }
-
     pub async fn new(config: &CdhConfig) -> Result<Self> {
-        Self::init(config).await?;
-        Ok(Self)
+        let hub = Hub::new(config.clone()).await?;
+
+        Ok(Self { hub })
     }
 }
 
@@ -61,9 +48,7 @@ impl SealedSecretService for Server {
         input: UnsealSecretInput,
     ) -> ::ttrpc::Result<UnsealSecretOutput> {
         debug!("[ttRPC CDH] get new UnsealSecret request");
-        let reader = HUB.read().await;
-        let reader = reader.as_ref().expect("must be initialized");
-        let plaintext = reader.unseal_secret(input.secret).await.map_err(|e| {
+        let plaintext = self.hub.unseal_secret(input.secret).await.map_err(|e| {
             let detailed_error = format_error!(e);
             error!("[ttRPC CDH] UnsealSecret :\n{detailed_error}");
             let mut status = Status::new();
@@ -87,9 +72,7 @@ impl GetResourceService for Server {
         req: GetResourceRequest,
     ) -> ::ttrpc::Result<GetResourceResponse> {
         debug!("[ttRPC CDH] get new GetResource request");
-        let reader = HUB.read().await;
-        let reader = reader.as_ref().expect("must be initialized");
-        let resource = reader.get_resource(req.ResourcePath).await.map_err(|e| {
+        let resource = self.hub.get_resource(req.ResourcePath).await.map_err(|e| {
             let detailed_error = format_error!(e);
             error!("[ttRPC CDH] GetResource :\n{detailed_error}");
             let mut status = Status::new();
@@ -113,8 +96,6 @@ impl KeyProviderService for Server {
         req: KeyProviderKeyWrapProtocolInput,
     ) -> ::ttrpc::Result<KeyProviderKeyWrapProtocolOutput> {
         debug!("[ttRPC CDH] get new UnWrapKey request");
-        let reader = HUB.read().await;
-        let reader = reader.as_ref().expect("must be initialized");
         let key_provider_input: KeyProviderInput =
             serde_json::from_slice(&req.KeyProviderKeyWrapProtocolInput[..]).map_err(|e| {
                 error!("[ttRPC CDH] UnwrapKey parse KeyProviderInput failed : {e:?}");
@@ -133,7 +114,7 @@ impl KeyProviderService for Server {
         })?;
 
         debug!("[ttRPC CDH] Call CDH to Unwrap Key...");
-        let decrypted_optsdata = reader.unwrap_key(&annotation_packet).await.map_err(|e| {
+        let decrypted_optsdata = self.hub.unwrap_key(&annotation_packet).await.map_err(|e| {
             let detailed_error = format_error!(e);
             error!("[ttRPC CDH] UnWrapKey :\n{detailed_error}");
             let mut status = Status::new();
@@ -173,15 +154,13 @@ impl SecureMountService for Server {
         req: SecureMountRequest,
     ) -> ::ttrpc::Result<SecureMountResponse> {
         debug!("[ttRPC CDH] get new secure mount request");
-        let reader = HUB.read().await;
-        let reader = reader.as_ref().expect("must be initialized");
         let storage = Storage {
             volume_type: req.volume_type,
             options: req.options,
             flags: req.flags,
             mount_point: req.mount_point,
         };
-        let resource = reader.secure_mount(storage).await.map_err(|e| {
+        let resource = self.hub.secure_mount(storage).await.map_err(|e| {
             let detailed_error = format_error!(e);
             error!("[ttRPC CDH] Secure Mount :\n{detailed_error}");
             let mut status = Status::new();
@@ -205,9 +184,8 @@ impl ImagePullService for Server {
         req: ImagePullRequest,
     ) -> ::ttrpc::Result<ImagePullResponse> {
         debug!("[ttRPC CDH] get new image pull request");
-        let reader = HUB.read().await;
-        let reader = reader.as_ref().expect("must be initialized");
-        let manifest_digest = reader
+        let manifest_digest = self
+            .hub
             .pull_image(&req.image_url, &req.bundle_path)
             .await
             .map_err(|e| {