From 98e03cc37363b61311e75108531ae7025e1e2901 Mon Sep 17 00:00:00 2001 From: Xynnn007 Date: Mon, 11 Dec 2023 17:53:14 +0800 Subject: [PATCH] kbs: attestation service client Currently we are using a global Mutex to protect the only one attestation service client from unsafe thread sync & send. This brings performance bottle neck. This commit will create a temporary client object to handle the current attestation request, which can break the mutex. Fixes #256 Signed-off-by: Xynnn007 --- kbs/src/api/src/attestation/coco/builtin.rs | 9 ++- kbs/src/api/src/attestation/mod.rs | 66 +++++++++++++++------ kbs/src/api/src/http/attest.rs | 3 - kbs/src/api/src/http/config.rs | 3 - kbs/src/kbs/src/main.rs | 6 +- 5 files changed, 57 insertions(+), 30 deletions(-) diff --git a/kbs/src/api/src/attestation/coco/builtin.rs b/kbs/src/api/src/attestation/coco/builtin.rs index 7dab65787d..81197f324d 100644 --- a/kbs/src/api/src/attestation/coco/builtin.rs +++ b/kbs/src/api/src/attestation/coco/builtin.rs @@ -11,9 +11,10 @@ use attestation_service::{ }; use kbs_types::{Attestation, Tee}; use serde_json::json; +use tokio::sync::RwLock; pub struct Native { - inner: AttestationService, + inner: RwLock, } #[async_trait] @@ -21,7 +22,7 @@ impl Attest for Native { async fn set_policy(&mut self, input: &[u8]) -> Result<()> { let request: SetPolicyInput = serde_json::from_slice(input).context("parse SetPolicyInput")?; - self.inner.set_policy(request).await + self.inner.write().await.set_policy(request).await } async fn verify(&mut self, tee: Tee, nonce: &str, attestation: &str) -> Result { @@ -31,6 +32,8 @@ impl Attest for Native { let runtime_data_plaintext = json!({"tee-pubkey": attestation.tee_pubkey, "nonce": nonce}); self.inner + .read() + .await .evaluate( attestation.tee_evidence.into_bytes(), tee, @@ -47,7 +50,7 @@ impl Attest for Native { impl Native { pub async fn new(config: &AsConfig) -> Result { Ok(Self { - inner: AttestationService::new(config.clone()).await?, + inner: RwLock::new(AttestationService::new(config.clone()).await?), }) } } diff --git a/kbs/src/api/src/attestation/mod.rs b/kbs/src/api/src/attestation/mod.rs index 53aa3f035e..c804a04fe5 100644 --- a/kbs/src/api/src/attestation/mod.rs +++ b/kbs/src/api/src/attestation/mod.rs @@ -11,8 +11,6 @@ use coco::grpc::GrpcConfig; #[cfg(feature = "intel-trust-authority-as")] use intel_trust_authority::IntelTrustAuthorityConfig; use kbs_types::Tee; -use std::sync::Arc; -use tokio::sync::Mutex; #[cfg(feature = "coco-as")] #[allow(missing_docs)] @@ -38,34 +36,66 @@ pub trait Attest: Send + Sync { /// Attestation Service #[derive(Clone)] -pub struct AttestationService(pub Arc>); +pub enum AttestationService { + #[cfg(any(feature = "coco-as-builtin", feature = "coco-as-builtin-no-verifier"))] + CoCoASBuiltIn(AsConfig), + + #[cfg(feature = "coco-as-grpc")] + CoCoASgRPC(GrpcConfig), + + #[cfg(feature = "intel-trust-authority-as")] + IntelTA(IntelTrustAuthorityConfig), +} impl AttestationService { /// Create and initialize AttestationService. #[cfg(any(feature = "coco-as-builtin", feature = "coco-as-builtin-no-verifier"))] - pub async fn new(config: &AsConfig) -> Result { - let attestation_service: Arc> = - Arc::new(Mutex::new(coco::builtin::Native::new(config).await?)); - - Ok(Self(attestation_service)) + pub fn new(config: AsConfig) -> Self { + Self::CoCoASBuiltIn(config) } /// Create and initialize AttestationService. #[cfg(feature = "coco-as-grpc")] - pub async fn new(config: &GrpcConfig) -> Result { - let attestation_service: Arc> = - Arc::new(Mutex::new(coco::grpc::Grpc::new(config).await?)); - - Ok(Self(attestation_service)) + pub fn new(config: GrpcConfig) -> Self { + Self::CoCoASgRPC(config) } /// Create and initialize AttestationService. #[cfg(feature = "intel-trust-authority-as")] - pub fn new(config: &IntelTrustAuthorityConfig) -> Result { - let attestation_service: Arc> = Arc::new(Mutex::new( - intel_trust_authority::IntelTrustAuthority::new(config)?, - )); + pub fn new(config: IntelTrustAuthorityConfig) -> Self { + Self::IntelTA(config) + } + + pub async fn create_client(&self) -> Result> { + match self { + #[cfg(any(feature = "coco-as-builtin", feature = "coco-as-builtin-no-verifier"))] + AttestationService::CoCoASBuiltIn(config) => { + Ok(Box::new(coco::builtin::Native::new(config).await?)) + } + #[cfg(feature = "coco-as-grpc")] + AttestationService::CoCoASgRPC(config) => { + Ok(Box::new(coco::grpc::Grpc::new(config).await?)) + } + #[cfg(feature = "intel-trust-authority-as")] + AttestationService::IntelTA(config) => Ok(Box::new( + intel_trust_authority::IntelTrustAuthority::new(config)?, + )), + } + } + + pub async fn verify(&self, tee: Tee, nonce: &str, attestation: &str) -> Result { + let mut client = self + .create_client() + .await + .context("attestation service client initialization failed.")?; + client.verify(tee, nonce, attestation).await + } - Ok(Self(attestation_service)) + pub async fn set_policy(&self, input: &[u8]) -> Result<()> { + let mut client = self + .create_client() + .await + .context("attestation service client initialization failed.")?; + client.set_policy(input).await } } diff --git a/kbs/src/api/src/http/attest.rs b/kbs/src/api/src/http/attest.rs index 9d0c1477db..2b351f1afc 100644 --- a/kbs/src/api/src/http/attest.rs +++ b/kbs/src/api/src/http/attest.rs @@ -56,9 +56,6 @@ pub(crate) async fn attest( } let token = attestation_service - .0 - .lock() - .await .verify( session.tee(), session.nonce(), diff --git a/kbs/src/api/src/http/config.rs b/kbs/src/api/src/http/config.rs index 7f8ac53f42..b85af476f5 100644 --- a/kbs/src/api/src/http/config.rs +++ b/kbs/src/api/src/http/config.rs @@ -25,9 +25,6 @@ pub(crate) async fn attestation_policy( } attestation_service - .0 - .lock() - .await .set_policy(&input) .await .map_err(|e| Error::PolicyEndpoint(format!("Set policy error {e}")))?; diff --git a/kbs/src/kbs/src/main.rs b/kbs/src/kbs/src/main.rs index 021844d623..12060280a8 100644 --- a/kbs/src/kbs/src/main.rs +++ b/kbs/src/kbs/src/main.rs @@ -43,11 +43,11 @@ async fn main() -> Result<()> { let attestation_service = { cfg_if::cfg_if! { if #[cfg(any(feature = "coco-as-builtin", feature = "coco-as-builtin-no-verifier"))] { - AttestationService::new(&kbs_config.as_config.unwrap_or_default()).await? + AttestationService::new(kbs_config.as_config.unwrap_or_default()) } else if #[cfg(feature = "coco-as-grpc")] { - AttestationService::new(&kbs_config.grpc_config.unwrap_or_default()).await? + AttestationService::new(kbs_config.grpc_config.unwrap_or_default()) } else if #[cfg(feature = "intel-trust-authority-as")] { - AttestationService::new(&kbs_config.intel_trust_authority_config)? + AttestationService::new(kbs_config.intel_trust_authority_config) } else { compile_error!("Please enable at least one of the following features: `coco-as-builtin`, `coco-as-builtin-no-verifier`, `coco-as-grpc` or `intel-trust-authority-as` to continue."); }