From ee65f675b6b1d647bc20637a22a0ae79224384f3 Mon Sep 17 00:00:00 2001 From: Xynnn007 Date: Fri, 27 Sep 2024 16:44:17 +0800 Subject: [PATCH] CDH: fix ttrpc cdh memory occupation Fixes #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 --- Cargo.lock | 10 ++-- Cargo.toml | 2 +- .../aa/attestation_agent_ttrpc.rs | 16 +++---- .../hub/src/bin/protos/api_ttrpc.rs | 18 ++++---- .../hub/src/bin/protos/keyprovider_ttrpc.rs | 8 ++-- .../hub/src/bin/ttrpc-cdh.rs | 25 +++------- .../hub/src/bin/ttrpc_server/mod.rs | 46 +++++-------------- .../kbs/ttrpc_proto/getresource_ttrpc.rs | 6 +-- .../src/utils/ttrpc/keyprovider_ttrpc.rs | 8 ++-- 9 files changed, 53 insertions(+), 86 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 18e50bdb9..6e78dcad6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3347,7 +3347,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c2a198fb6b0eada2a8df47933734e6d35d350665a33a3593d7164fa52c75c19" dependencies = [ "cfg-if", - "windows-targets 0.52.6", + "windows-targets 0.48.5", ] [[package]] @@ -6571,9 +6571,9 @@ dependencies = [ [[package]] name = "ttrpc" -version = "0.8.1" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "55ea338db445bee75c596cf8a478fbfcefad5a943c9e92a7e1c805c65ed39551" +checksum = "e376927d4422245ae3e0a0d7df0e805f99652536999b5c671144de9fe4120d8c" dependencies = [ "async-trait", "byteorder", @@ -6604,9 +6604,9 @@ dependencies = [ [[package]] name = "ttrpc-compiler" -version = "0.6.2" +version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0672eb06e5663ad190c7b93b2973f5d730259859b62e4e3381301a12a7441107" +checksum = "04c4c51f20209ae3ec2579b947b54cf52685825238002bc2e5af8e1e075d4813" dependencies = [ "derive-new", "prost 0.8.0", diff --git a/Cargo.toml b/Cargo.toml index b575169de..fef535079 100644 --- 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 index a2ff57e23..85fe45a8f 100644 --- 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>, + service: Arc, } #[async_trait] @@ -74,7 +74,7 @@ impl ::ttrpc::r#async::MethodHandler for GetEvidenceMethod { } struct GetTokenMethod { - service: Arc>, + service: Arc, } #[async_trait] @@ -85,7 +85,7 @@ impl ::ttrpc::r#async::MethodHandler for GetTokenMethod { } struct ExtendRuntimeMeasurementMethod { - service: Arc>, + service: Arc, } #[async_trait] @@ -96,7 +96,7 @@ impl ::ttrpc::r#async::MethodHandler for ExtendRuntimeMeasurementMethod { } struct CheckInitDataMethod { - service: Arc>, + service: Arc, } #[async_trait] @@ -107,7 +107,7 @@ impl ::ttrpc::r#async::MethodHandler for CheckInitDataMethod { } struct UpdateConfigurationMethod { - service: Arc>, + service: Arc, } #[async_trait] @@ -118,7 +118,7 @@ impl ::ttrpc::r#async::MethodHandler for UpdateConfigurationMethod { } struct GetTeeTypeMethod { - service: Arc>, + service: Arc, } #[async_trait] @@ -150,7 +150,7 @@ pub trait AttestationAgentService: Sync { } } -pub fn create_attestation_agent_service(service: Arc>) -> HashMap { +pub fn create_attestation_agent_service(service: Arc) -> HashMap { 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 index 7142428e3..35d2c704d 100644 --- 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>, + service: Arc, } #[async_trait] @@ -55,7 +55,7 @@ pub trait SealedSecretService: Sync { } } -pub fn create_sealed_secret_service(service: Arc>) -> HashMap { +pub fn create_sealed_secret_service(service: Arc) -> HashMap { let mut ret = HashMap::new(); let mut methods = HashMap::new(); let streams = HashMap::new(); @@ -86,7 +86,7 @@ impl GetResourceServiceClient { } struct GetResourceMethod { - service: Arc>, + service: Arc, } #[async_trait] @@ -103,7 +103,7 @@ pub trait GetResourceService: Sync { } } -pub fn create_get_resource_service(service: Arc>) -> HashMap { +pub fn create_get_resource_service(service: Arc) -> HashMap { let mut ret = HashMap::new(); let mut methods = HashMap::new(); let streams = HashMap::new(); @@ -134,7 +134,7 @@ impl SecureMountServiceClient { } struct SecureMountMethod { - service: Arc>, + service: Arc, } #[async_trait] @@ -151,7 +151,7 @@ pub trait SecureMountService: Sync { } } -pub fn create_secure_mount_service(service: Arc>) -> HashMap { +pub fn create_secure_mount_service(service: Arc) -> HashMap { let mut ret = HashMap::new(); let mut methods = HashMap::new(); let streams = HashMap::new(); @@ -182,7 +182,7 @@ impl ImagePullServiceClient { } struct PullImageMethod { - service: Arc>, + service: Arc, } #[async_trait] @@ -199,7 +199,7 @@ pub trait ImagePullService: Sync { } } -pub fn create_image_pull_service(service: Arc>) -> HashMap { +pub fn create_image_pull_service(service: Arc) -> HashMap { 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 index 059b79f72..9d4364cb9 100644 --- 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>, + service: Arc, } #[async_trait] @@ -54,7 +54,7 @@ impl ::ttrpc::r#async::MethodHandler for WrapKeyMethod { } struct UnWrapKeyMethod { - service: Arc>, + service: Arc, } #[async_trait] @@ -74,7 +74,7 @@ pub trait KeyProviderService: Sync { } } -pub fn create_key_provider_service(service: Arc>) -> HashMap { +pub fn create_key_provider_service(service: Arc) -> HashMap { 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 index 9138da32f..869619ee3 100644 --- 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, } -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 index e7a71e5b4..448c49323 100644 --- 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>> = 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::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 { 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 { 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 { 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 { 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 { 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| { diff --git a/image-rs/src/resource/kbs/ttrpc_proto/getresource_ttrpc.rs b/image-rs/src/resource/kbs/ttrpc_proto/getresource_ttrpc.rs index 5e8aebecf..b565ca784 100644 --- a/image-rs/src/resource/kbs/ttrpc_proto/getresource_ttrpc.rs +++ b/image-rs/src/resource/kbs/ttrpc_proto/getresource_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 GetResourceServiceClient { } struct GetResourceMethod { - service: Arc>, + service: Arc, } #[async_trait] @@ -55,7 +55,7 @@ pub trait GetResourceService: Sync { } } -pub fn create_get_resource_service(service: Arc>) -> HashMap { +pub fn create_get_resource_service(service: Arc) -> HashMap { let mut ret = HashMap::new(); let mut methods = HashMap::new(); let streams = HashMap::new(); diff --git a/ocicrypt-rs/src/utils/ttrpc/keyprovider_ttrpc.rs b/ocicrypt-rs/src/utils/ttrpc/keyprovider_ttrpc.rs index 9a68b073e..7f18d5a51 100644 --- a/ocicrypt-rs/src/utils/ttrpc/keyprovider_ttrpc.rs +++ b/ocicrypt-rs/src/utils/ttrpc/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)] @@ -44,7 +44,7 @@ impl KeyProviderServiceClient { } struct WrapKeyMethod { - service: Arc>, + service: Arc, } impl ::ttrpc::MethodHandler for WrapKeyMethod { @@ -55,7 +55,7 @@ impl ::ttrpc::MethodHandler for WrapKeyMethod { } struct UnWrapKeyMethod { - service: Arc>, + service: Arc, } impl ::ttrpc::MethodHandler for UnWrapKeyMethod { @@ -74,7 +74,7 @@ pub trait KeyProviderService { } } -pub fn create_key_provider_service(service: Arc>) -> HashMap> { +pub fn create_key_provider_service(service: Arc) -> HashMap> { let mut methods = HashMap::new(); methods.insert("/keyprovider.KeyProviderService/WrapKey".to_string(),