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 confidential-containers#256

Signed-off-by: Xynnn007 <[email protected]>

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