From 20165a6e84d66cd2182d3d3927daf75ac638bb0a Mon Sep 17 00:00:00 2001 From: Noor Date: Fri, 11 Apr 2025 01:58:00 +0530 Subject: [PATCH] refactor: DTMP grpc connction and methods updated Dtpm client connection remove mr_signer from connetion reuse dtpm client connection pure function for grpc methods some error handling --- src/sgx/cli_handler.rs | 6 +-- src/sgx/grpc_dtpm.rs | 104 +++++++++++++++++++---------------------- src/sgx/utils.rs | 27 +++++++---- 3 files changed, 71 insertions(+), 66 deletions(-) diff --git a/src/sgx/cli_handler.rs b/src/sgx/cli_handler.rs index d23bbfc..be133d3 100644 --- a/src/sgx/cli_handler.rs +++ b/src/sgx/cli_handler.rs @@ -2,7 +2,7 @@ use crate::config::Config; use crate::name_generator::random_app_name; use crate::sgx::config::{validate_yaml, DeteeCliExt}; use crate::sgx::grpc_brain::{delete_app, list_contracts}; -use crate::sgx::grpc_dtpm::{attest_and_send_config, get_config_from_enclave}; +use crate::sgx::grpc_dtpm::{get_config, update_config}; use crate::sgx::packaging::package_enclave; use crate::sgx::utils::{ deploy_new_app_and_update_config, fetch_config, override_envs_and_args_launch_config, @@ -225,7 +225,7 @@ fn handle_config_sub_update( (update_matche.get_one::("config"), update_matche.get_one::("uuid")) { let loaded_config = validate_yaml(file_path).unwrap(); - match block_on(attest_and_send_config(loaded_config, uuid)) { + match block_on(update_config(uuid, loaded_config)) { Ok(_) => Ok(SimpleOutput::from("App launch config updated successfully")), Err(e) => Err(Box::new(std::io::Error::other(format!( "Could not attest and update app launch config due to error: {e}" @@ -242,7 +242,7 @@ fn handle_config_sub_get( if let (Some(file_path_to_save), Some(uuid)) = (get_matche.get_one::("path"), get_matche.get_one::("uuid")) { - match block_on(get_config_from_enclave(uuid)) { + match block_on(get_config(uuid)) { Ok(config) => { let config_yaml = serde_yaml::to_string(&config).unwrap(); std::fs::write(file_path_to_save, config_yaml).unwrap(); diff --git a/src/sgx/grpc_dtpm.rs b/src/sgx/grpc_dtpm.rs index d8d45ff..04980cf 100644 --- a/src/sgx/grpc_dtpm.rs +++ b/src/sgx/grpc_dtpm.rs @@ -1,5 +1,8 @@ use detee_sgx::{prelude::*, HRaTlsConfigBuilder}; -use detee_shared::common_proto::Empty; +use detee_shared::{ + common_proto::Empty, + sgx::{pb::dtpm_proto::DtpmGetConfigRes, types::dtpm::FileEntry}, +}; use hyper_rustls::HttpsConnectorBuilder; use rustls::ClientConfig; use std::sync::{Arc, RwLock}; @@ -12,8 +15,8 @@ use tonic::{ use detee_shared::sgx::{ pb::dtpm_proto::{ - dtpm_config_manager_client::DtpmConfigManagerClient, DtpmConfigData, DtpmSetConfigReq, - FileEntry, + dtpm_config_manager_client::DtpmConfigManagerClient, DtpmSetConfigReq, + FileEntry as FileEntryPb, }, types::dtpm::DtpmConfig, }; @@ -33,28 +36,22 @@ pub enum Error { DiskAccess(#[from] std::io::Error), #[error("HRatls: {0}")] SgxHRatls(#[from] detee_sgx::error::SgxError), - #[error("DtpmConfig: {0}")] - DtpmConfig(String), + #[error("Dtpm: {0}")] + Dtpm(String), #[error(transparent)] ConfigError(#[from] crate::config::Error), } type Result = std::result::Result; -pub async fn connect_dtpm_grpc_client( - hratls_uri: String, - package_mr_enclave: Option<[u8; 32]>, -) -> Result> { +pub async fn connect_app_dtpm_client(app_uuid: &str) -> Result> { let private_key_pem = Config::get_hratls_private_key()?; - let mut mr_signer = [0u8; 32]; - hex::decode_to_slice(Config::get_mrsigner()?, &mut mr_signer)?; - let mr_signers = vec![mr_signer]; - let hratls_config = Arc::new(RwLock::new( - HRaTlsConfig::new() - .allow_instance_measurement(InstanceMeasurement::new().with_mrsigners(mr_signers)) - .with_hratls_private_key_pem(private_key_pem), - )); + let (hratls_uri, package_mr_enclave) = hratls_url_and_mr_enclave_from_app_id(app_uuid).await?; + log::info!("hratls uri: {}\nmr_enclave: {:?}", &hratls_uri, &package_mr_enclave); + + let hratls_config = + Arc::new(RwLock::new(HRaTlsConfig::new().with_hratls_private_key_pem(private_key_pem))); if let Some(mr_enclave) = package_mr_enclave { hratls_config.write().unwrap().allow_more_instance_measurement( @@ -74,60 +71,57 @@ pub async fn connect_dtpm_grpc_client( Ok(DtpmConfigManagerClient::new(channel).send_compressed(CompressionEncoding::Zstd)) } -pub async fn attest_and_send_config(loaded_config: DtpmConfig, uuid: &str) -> Result<()> { - let files = loaded_config.filesystems.clone(); +pub async fn update_config(app_uuid: &str, config: DtpmConfig) -> Result<()> { + let dtpm_client = connect_app_dtpm_client(app_uuid).await?; - let config_data = Some(DtpmConfigData::from(loaded_config)); - let req_data = DtpmSetConfigReq { config_data, ..Default::default() }; + upload_files_pb(config.filesystems.clone(), &dtpm_client).await?; + let req = DtpmSetConfigReq { config_data: Some(config.into()), ..Default::default() }; + set_config_pb(req, &dtpm_client).await?; + Ok(()) +} - log::trace!("Decoded the configuration... {:?}", req_data); - - let (hratls_uri, mr_enclave) = hratls_url_and_mr_enclave_from_app_id(uuid).await; - log::info!("hratls uri: {}\nmr_enclave: {:?}", &hratls_uri, &mr_enclave); - - let client = connect_dtpm_grpc_client(hratls_uri, mr_enclave).await?; +pub async fn get_config(app_uuid: &str) -> Result { + let dtpm_client = connect_app_dtpm_client(app_uuid).await?; + let config_res = get_config_pb(&dtpm_client).await?; + let config: DtpmConfig = + config_res.config_data.ok_or(Error::Dtpm("config data not found".to_string()))?.into(); + Ok(config) +} +pub async fn upload_files_pb( + reqs: Vec, + client: &DtpmConfigManagerClient, +) -> Result<()> { let (tx, rx) = mpsc::channel(6); tokio::spawn(async move { - for file in files { - let file_pb: FileEntry = file.into(); + for file in reqs { + let file_pb: FileEntryPb = file.into(); let _ = tx.send(file_pb).await; } }); let fs_stream = ReceiverStream::new(rx); - let _ = client - .clone() - .max_decoding_message_size(10240000) - .upload_files(tonic::Request::new(fs_stream)) - .await?; + let _ = client.clone().upload_files(tonic::Request::new(fs_stream)).await?; - let response = client - .max_decoding_message_size(10240000) - .set_config(tonic::Request::new(req_data)) - .await?; + Ok(()) +} + +pub(crate) async fn set_config_pb( + req: DtpmSetConfigReq, + client: &DtpmConfigManagerClient, +) -> Result<()> { + let response = client.clone().set_config(tonic::Request::new(req)).await?; log::trace!("Received respose from the server...{:?}", response.into_inner()); Ok(()) } -pub async fn get_config_from_enclave(uuid: &str) -> Result { - let (hratls_uri, mr_enclave) = hratls_url_and_mr_enclave_from_app_id(uuid).await; - log::info!("hratls uri: {}\nmr_enclave: {:?}", &hratls_uri, &mr_enclave); - - let client = connect_dtpm_grpc_client(hratls_uri, None).await?; - - let mgr_config_pb = client - .max_decoding_message_size(10240000) - .get_config(tonic::Request::new(Empty {})) - .await? - .into_inner(); - - let config: DtpmConfig = mgr_config_pb - .config_data - .ok_or(Error::DtpmConfig("config data not found".to_string()))? - .into(); - Ok(config) +pub(crate) async fn get_config_pb( + client: &DtpmConfigManagerClient, +) -> Result { + let mgr_config_pb = + client.clone().get_config(tonic::Request::new(Empty {})).await?.into_inner(); + Ok(mgr_config_pb) } diff --git a/src/sgx/utils.rs b/src/sgx/utils.rs index 422c3ee..eb2b7b7 100644 --- a/src/sgx/utils.rs +++ b/src/sgx/utils.rs @@ -1,9 +1,12 @@ use crate::constants::HRATLS_APP_PORT; use crate::sgx::get_one_contract; use crate::sgx::grpc_brain::new_app; -use crate::sgx::grpc_dtpm::attest_and_send_config; +use crate::sgx::grpc_dtpm::connect_app_dtpm_client; +use crate::sgx::grpc_dtpm::set_config_pb; +use crate::sgx::grpc_dtpm::upload_files_pb; use crate::sgx::package_entry_from_name; use detee_shared::app_proto::NewAppRes; +use detee_shared::sgx::pb::dtpm_proto::DtpmSetConfigReq; use detee_shared::sgx::types::brain::AppDeployConfig; use detee_shared::sgx::types::dtpm::DtpmConfig; use detee_shared::sgx::types::dtpm::EnvironmentEntry; @@ -17,7 +20,7 @@ pub enum Error { #[error(transparent)] Serde(#[from] serde_yaml::Error), #[error("{0}")] - PublicPackage(std::string::String), + Package(std::string::String), #[error("{0}")] Brain(#[from] crate::sgx::grpc_brain::Error), #[error("{0}")] @@ -26,7 +29,9 @@ pub enum Error { Deployment(String), } -pub async fn hratls_url_and_mr_enclave_from_app_id(app_id: &str) -> (String, Option<[u8; 32]>) { +pub async fn hratls_url_and_mr_enclave_from_app_id( + app_id: &str, +) -> Result<(String, Option<[u8; 32]>), crate::sgx::grpc_dtpm::Error> { let app_contract = get_one_contract(app_id).await; if app_contract.is_err() { eprintln!("Could not find App contract with ID: {}", app_id); @@ -44,15 +49,15 @@ pub async fn hratls_url_and_mr_enclave_from_app_id(app_id: &str) -> (String, Opt .mapped_ports .iter() .find(|port| port.app_port == HRATLS_APP_PORT) - .unwrap() + .ok_or(crate::sgx::grpc_dtpm::Error::Dtpm("Could not find DTMP port".to_string()))? .host_port; - (format!("https://{public_ip}:{dtpm_port}"), mr_enclave) + Ok((format!("https://{public_ip}:{dtpm_port}"), mr_enclave)) } pub async fn fetch_config(package_name: &str) -> Result { let index_package_entry = package_entry_from_name(package_name) - .ok_or(Error::PublicPackage("package not found for ".to_string() + package_name))?; + .ok_or(Error::Package("package not found for ".to_string() + package_name))?; let launch_config_url = index_package_entry.launch_config_url.clone(); @@ -128,11 +133,17 @@ pub async fn deploy_new_app_and_update_config( if let Some(launch_config) = launch_config { eprintln!("Deploying..."); tokio::time::sleep(tokio::time::Duration::from_millis(2500)).await; - Retry::spawn(FixedInterval::from_millis(1000).take(10), || { + let dtpm_client = Retry::spawn(FixedInterval::from_millis(1000).take(10), || { log::debug!("retrying attestation and launch config update"); - attest_and_send_config(launch_config.clone(), &new_app_res.uuid) + connect_app_dtpm_client(&new_app_res.uuid) }) .await?; + upload_files_pb(launch_config.filesystems.clone(), &dtpm_client).await?; + + let config_data = Some(launch_config.into()); + log::trace!("Decoded the configuration... {:?}", config_data); + let req = DtpmSetConfigReq { config_data, ..Default::default() }; + set_config_pb(req, &dtpm_client).await?; Ok(new_app_res) } else { Ok(new_app_res)