KBS/perf: promote the concurrency performance of KBS

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 brings some optimization to promote the performance and
stability.

1. Abondon the global Mutex of the attestation service and change the API definition of
Attest to non-mut. This would let the developers to handle the
concurrency safe inside the concrete attestation-service inside. In this
way, we prevent to lock the whole process logic.
2. Bring in a gRPC client pool to grpc-coco-as mode. This will help to
avoid errors that caused by runing out all the temporaty ports provided
by OS.
3. Replace the Mutex of session map with a concurrency-safe HashMap to
avoid bottle neck.

Fixes #256

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

kbs/config/kbs-config-grpc.toml

@@ -3,3 +3,4 @@ insecure_api = true
 
 [grpc_config]
 as_addr = "http://127.0.0.1:50004"
+pool_size = 200

kbs/src/api/Cargo.toml

@@ -15,7 +15,7 @@ opa = ["policy"]
 coco-as = ["as"]
 coco-as-builtin = ["coco-as", "attestation-service/default"]
 coco-as-builtin-no-verifier = ["coco-as", "attestation-service/rvps-builtin"]
-coco-as-grpc = ["coco-as", "tonic", "tonic-build", "prost"]
+coco-as-grpc = ["coco-as", "mobc", "tonic", "tonic-build", "prost"]
 intel-trust-authority-as = ["as", "reqwest", "jsonwebtoken"]
 rustls = ["actix-web/rustls", "dep:rustls", "dep:rustls-pemfile"]
 openssl = ["actix-web/openssl", "dep:openssl"]
@@ -37,12 +37,14 @@ jwt-simple = "0.11.6"
 kbs-types.workspace = true
 lazy_static = "1.4.0"
 log.workspace = true
+mobc = { version = "0.8.3", optional = true }
 prost = { version = "0.11", optional = true }
 rand = "0.8.5"
 reqwest = { version = "0.11", features = ["json"], optional = true }
 rsa = { version = "0.9.2", optional = true, features = ["sha2"] }
 rustls = { version = "0.20.8", optional = true }
 rustls-pemfile = { version = "1.0.4", optional = true }
+scc = "2"
 semver = "1.0.16"
 serde = { version = "1.0", features = ["derive"] }
 serde_json.workspace = true

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

@@ -11,26 +11,29 @@ use attestation_service::{
 };
 use kbs_types::{Attestation, Tee};
 use serde_json::json;
+use tokio::sync::RwLock;
 
-pub struct Native {
-    inner: AttestationService,
+pub struct BuiltInCoCoAs {
+    inner: RwLock<AttestationService>,
 }
 
 #[async_trait]
-impl Attest for Native {
-    async fn set_policy(&mut self, input: &[u8]) -> Result<()> {
+impl Attest for BuiltInCoCoAs {
+    async fn set_policy(&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> {
+    async fn verify(&self, tee: Tee, nonce: &str, attestation: &str) -> Result<String> {
         let attestation: Attestation = serde_json::from_str(attestation)?;
 
         // TODO: align with the guest-components/kbs-protocol side.
         let runtime_data_plaintext = json!({"tee-pubkey": attestation.tee_pubkey, "nonce": nonce});
 
         self.inner
+            .read()
+            .await
             .evaluate(
                 attestation.tee_evidence.into_bytes(),
                 tee,
@@ -44,10 +47,9 @@ impl Attest for Native {
     }
 }
 
-impl Native {
-    pub async fn new(config: &AsConfig) -> Result<Self> {
-        Ok(Self {
-            inner: AttestationService::new(config.clone()).await?,
-        })
+impl BuiltInCoCoAs {
+    pub async fn new(config: AsConfig) -> Result<Self> {
+        let inner = RwLock::new(AttestationService::new(config).await?);
+        Ok(Self { inner })
     }
 }