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 <[email protected]>

kbs/src/api/src/attestation/coco/builtin.rs

@@ -11,17 +11,18 @@ use attestation_service::{
 };
 use kbs_types::{Attestation, Tee};
 use serde_json::json;
+use tokio::sync::RwLock;
 
 pub struct Native {
-    inner: AttestationService,
+    inner: RwLock<AttestationService>,
 }
 
 #[async_trait]
 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<String> {
@@ -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<Self> {
         Ok(Self {
-            inner: AttestationService::new(config.clone()).await?,
+            inner: RwLock::new(AttestationService::new(config.clone()).await?),
         })
     }
 }

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<Mutex<dyn Attest>>);
+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<Self> {
-        let attestation_service: Arc<Mutex<dyn Attest>> =
-            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<Self> {
-        let attestation_service: Arc<Mutex<dyn Attest>> =
-            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<Self> {
-        let attestation_service: Arc<Mutex<dyn Attest>> = 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<Box<dyn Attest>> {
+        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<String> {
+        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
     }
 }

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(),

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}")))?;

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.");
             }