From d3c1ed5a2f8b37562394ce1d191693104142cb79 Mon Sep 17 00:00:00 2001 From: Xynnn007 Date: Wed, 27 Dec 2023 17:15:05 +0800 Subject: [PATCH] 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 Signed-off-by: Biao Lu --- Cargo.lock | 115 ++++++++++++++++++ kbs/config/kbs-config-grpc.toml | 1 + kbs/src/api/Cargo.toml | 4 +- kbs/src/api/src/attestation/coco/builtin.rs | 24 ++-- kbs/src/api/src/attestation/coco/grpc.rs | 76 ++++++++---- .../attestation/intel_trust_authority/mod.rs | 4 +- kbs/src/api/src/attestation/mod.rs | 67 ++++++---- kbs/src/api/src/http/attest.rs | 43 +++---- kbs/src/api/src/http/config.rs | 5 +- kbs/src/api/src/http/resource.rs | 9 +- kbs/src/api/src/lib.rs | 8 +- kbs/src/api/src/session.rs | 7 +- kbs/src/kbs/src/main.rs | 8 +- 13 files changed, 268 insertions(+), 103 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 114c20861..cc4be4d5f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -417,6 +417,7 @@ dependencies = [ "kbs-types", "lazy_static", "log", + "mobc", "openssl", "prost", "rand", @@ -425,6 +426,7 @@ dependencies = [ "rstest", "rustls 0.20.9", "rustls-pemfile", + "scc", "semver 1.0.20", "serde", "serde_json", @@ -2617,6 +2619,27 @@ dependencies = [ "autocfg", ] +[[package]] +name = "metrics" +version = "0.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e52eb6380b6d2a10eb3434aec0885374490f5b82c8aaf5cd487a183c98be834" +dependencies = [ + "ahash 0.7.7", + "metrics-macros", +] + +[[package]] +name = "metrics-macros" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49e30813093f757be5cf21e50389a24dc7dbb22c49f23b7e8f51d69b508a5ffa" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "mime" version = "0.3.17" @@ -2650,6 +2673,25 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "mobc" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90eb49dc5d193287ff80e72a86f34cfb27aae562299d22fea215e06ea1059dd3" +dependencies = [ + "async-trait", + "futures-channel", + "futures-core", + "futures-timer", + "futures-util", + "log", + "metrics", + "thiserror", + "tokio", + "tracing", + "tracing-subscriber", +] + [[package]] name = "multimap" version = "0.8.3" @@ -2697,6 +2739,16 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "nu-ansi-term" +version = "0.46.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84" +dependencies = [ + "overload", + "winapi", +] + [[package]] name = "num-bigint" version = "0.4.4" @@ -2911,6 +2963,12 @@ version = "6.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e2355d85b9a3786f481747ced0e0ff2ba35213a1f9bd406ed906554d7af805a1" +[[package]] +name = "overload" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" + [[package]] name = "p256" version = "0.13.2" @@ -3893,6 +3951,12 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "scc" +version = "2.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60da9a72c824ff528dbae0c744d24b9f039dcde49cca9dd2f34438d5b0a1578c" + [[package]] name = "schannel" version = "0.1.22" @@ -4209,6 +4273,15 @@ dependencies = [ "tzdb", ] +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + [[package]] name = "shlex" version = "1.2.0" @@ -4559,6 +4632,16 @@ dependencies = [ "syn 2.0.39", ] +[[package]] +name = "thread_local" +version = "1.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3fdd6f064ccff2d6567adcb3873ca630700f00b5ad3f060c25b5dcfd9a4ce152" +dependencies = [ + "cfg-if", + "once_cell", +] + [[package]] name = "time" version = "0.3.30" @@ -4857,6 +4940,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c06d3da6113f116aaee68e4d601191614c9053067f9ab7f6edbcb161237daa54" dependencies = [ "once_cell", + "valuable", ] [[package]] @@ -4869,6 +4953,31 @@ dependencies = [ "tracing", ] +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad0f048c97dbd9faa9b7df56362b8ebcaa52adb06b498c050d2f4e32f90a7a8b" +dependencies = [ + "nu-ansi-term", + "sharded-slab", + "smallvec", + "thread_local", + "tracing-core", + "tracing-log", +] + [[package]] name = "try-lock" version = "0.2.4" @@ -5038,6 +5147,12 @@ dependencies = [ "serde", ] +[[package]] +name = "valuable" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d" + [[package]] name = "vcpkg" version = "0.2.15" diff --git a/kbs/config/kbs-config-grpc.toml b/kbs/config/kbs-config-grpc.toml index 1ef4b83ac..7200f365d 100644 --- a/kbs/config/kbs-config-grpc.toml +++ b/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 \ No newline at end of file diff --git a/kbs/src/api/Cargo.toml b/kbs/src/api/Cargo.toml index 2ab93028e..a6a4d02bd 100644 --- a/kbs/src/api/Cargo.toml +++ b/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 diff --git a/kbs/src/api/src/attestation/coco/builtin.rs b/kbs/src/api/src/attestation/coco/builtin.rs index 7dab65787..3f8eb4bac 100644 --- a/kbs/src/api/src/attestation/coco/builtin.rs +++ b/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, } #[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 { + async fn verify(&self, tee: Tee, nonce: &str, attestation: &str) -> Result { 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 { - Ok(Self { - inner: AttestationService::new(config.clone()).await?, - }) +impl BuiltInCoCoAs { + pub async fn new(config: AsConfig) -> Result { + let inner = RwLock::new(AttestationService::new(config).await?); + Ok(Self { inner }) } } diff --git a/kbs/src/api/src/attestation/coco/grpc.rs b/kbs/src/api/src/attestation/coco/grpc.rs index e708d0ddb..d37b215e7 100644 --- a/kbs/src/api/src/attestation/coco/grpc.rs +++ b/kbs/src/api/src/attestation/coco/grpc.rs @@ -8,8 +8,10 @@ use async_trait::async_trait; use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine}; use kbs_types::{Attestation, Tee}; use log::info; +use mobc::{Manager, Pool}; use serde::Deserialize; use serde_json::json; +use tokio::sync::Mutex; use tonic::transport::Channel; use self::attestation::{ @@ -18,12 +20,11 @@ use self::attestation::{ }; mod attestation { - #![allow(unknown_lints)] - #![allow(clippy::derive_partial_eq_without_eq)] tonic::include_proto!("attestation"); } pub const DEFAULT_AS_ADDR: &str = "http://127.0.0.1:50004"; +pub const DEFAULT_POOL_SIZE: u64 = 100; pub const COCO_AS_HASH_ALGORITHM: &str = "sha384"; @@ -44,44 +45,51 @@ fn to_grpc_tee(tee: Tee) -> GrpcTee { #[derive(Clone, Debug, Deserialize)] pub struct GrpcConfig { as_addr: Option, + pool_size: Option, } impl Default for GrpcConfig { fn default() -> Self { Self { as_addr: Some(DEFAULT_AS_ADDR.to_string()), + pool_size: Some(DEFAULT_POOL_SIZE), } } } -pub struct Grpc { - inner: AttestationServiceClient, +pub struct GrpcClientPool { + pool: Mutex>, } -impl Grpc { - pub async fn new(config: &GrpcConfig) -> Result { - let as_addr = match &config.as_addr { - Some(addr) => addr.clone(), - None => { - log::info!("Default remote AS address (127.0.0.1:50004) is used"); - DEFAULT_AS_ADDR.to_string() - } - }; - - info!("connect to remote AS [{as_addr}]"); - let inner = AttestationServiceClient::connect(as_addr).await?; - Ok(Self { inner }) +impl GrpcClientPool { + pub async fn new(config: GrpcConfig) -> Result { + let as_addr = config.as_addr.unwrap_or_else(|| { + log::info!("Default remote AS address ({DEFAULT_AS_ADDR}) is used"); + DEFAULT_AS_ADDR.to_string() + }); + + let pool_size = config.pool_size.unwrap_or_else(|| { + log::info!("Default AS connection pool size ({DEFAULT_POOL_SIZE}) is used"); + DEFAULT_POOL_SIZE + }); + + info!("connect to remote AS [{as_addr}] with pool size {pool_size}"); + let manager = GrpcManager { as_addr }; + let pool = Mutex::new(Pool::builder().max_open(pool_size).build(manager)); + + Ok(Self { pool }) } } #[async_trait] -impl Attest for Grpc { - async fn set_policy(&mut self, input: &[u8]) -> Result<()> { +impl Attest for GrpcClientPool { + async fn set_policy(&self, input: &[u8]) -> Result<()> { let input = String::from_utf8(input.to_vec()).context("parse SetPolicyInput")?; let req = tonic::Request::new(SetPolicyRequest { input }); - let _ = self - .inner + let mut client = { self.pool.lock().await.get().await? }; + + client .set_attestation_policy(req) .await .map_err(|e| anyhow!("Set Policy Failed: {:?}", e))?; @@ -89,7 +97,7 @@ impl Attest for Grpc { Ok(()) } - async fn verify(&mut self, tee: Tee, nonce: &str, attestation: &str) -> Result { + async fn verify(&self, tee: Tee, nonce: &str, attestation: &str) -> Result { let attestation: Attestation = serde_json::from_str(attestation)?; // TODO: align with the guest-components/kbs-protocol side. @@ -107,8 +115,9 @@ impl Attest for Grpc { policy_ids: vec!["default".to_string()], }); - let token = self - .inner + let mut client = { self.pool.lock().await.get().await? }; + + let token = client .attestation_evaluate(req) .await? .into_inner() @@ -117,3 +126,22 @@ impl Attest for Grpc { Ok(token) } } + +pub struct GrpcManager { + as_addr: String, +} + +#[async_trait] +impl Manager for GrpcManager { + type Connection = AttestationServiceClient; + type Error = tonic::transport::Error; + + async fn connect(&self) -> Result { + let connection = AttestationServiceClient::connect(self.as_addr.clone()).await?; + std::result::Result::Ok(connection) + } + + async fn check(&self, conn: Self::Connection) -> Result { + std::result::Result::Ok(conn) + } +} diff --git a/kbs/src/api/src/attestation/intel_trust_authority/mod.rs b/kbs/src/api/src/attestation/intel_trust_authority/mod.rs index 1648ff154..2c987acb4 100644 --- a/kbs/src/api/src/attestation/intel_trust_authority/mod.rs +++ b/kbs/src/api/src/attestation/intel_trust_authority/mod.rs @@ -49,7 +49,7 @@ pub struct IntelTrustAuthority { #[async_trait] impl Attest for IntelTrustAuthority { - async fn verify(&mut self, tee: Tee, _nonce: &str, attestation: &str) -> Result { + async fn verify(&self, tee: Tee, _nonce: &str, attestation: &str) -> Result { if tee != Tee::Tdx && tee != Tee::Sgx { bail!("Intel Trust Authority: TEE {tee:?} is not supported."); } @@ -119,7 +119,7 @@ impl Attest for IntelTrustAuthority { } impl IntelTrustAuthority { - pub fn new(config: &IntelTrustAuthorityConfig) -> Result { + pub fn new(config: IntelTrustAuthorityConfig) -> Result { let file = File::open(&config.certs_file) .map_err(|e| anyhow!("Open certs file failed: {:?}", e))?; let reader = BufReader::new(file); diff --git a/kbs/src/api/src/attestation/mod.rs b/kbs/src/api/src/attestation/mod.rs index 53aa3f035..db7fc4926 100644 --- a/kbs/src/api/src/attestation/mod.rs +++ b/kbs/src/api/src/attestation/mod.rs @@ -7,12 +7,10 @@ use async_trait::async_trait; #[cfg(any(feature = "coco-as-builtin", feature = "coco-as-builtin-no-verifier"))] use attestation_service::config::Config as AsConfig; #[cfg(feature = "coco-as-grpc")] -use coco::grpc::GrpcConfig; +use coco::grpc::*; #[cfg(feature = "intel-trust-authority-as")] -use intel_trust_authority::IntelTrustAuthorityConfig; +use intel_trust_authority::*; use kbs_types::Tee; -use std::sync::Arc; -use tokio::sync::Mutex; #[cfg(feature = "coco-as")] #[allow(missing_docs)] @@ -27,45 +25,68 @@ pub mod intel_trust_authority; #[async_trait] pub trait Attest: Send + Sync { /// Set Attestation Policy - async fn set_policy(&mut self, _input: &[u8]) -> Result<()> { + async fn set_policy(&self, _input: &[u8]) -> Result<()> { Err(anyhow!("Set Policy API is unimplemented")) } /// Verify Attestation Evidence /// Return Attestation Results Token - async fn verify(&mut self, tee: Tee, nonce: &str, attestation: &str) -> Result; + async fn verify(&self, tee: Tee, nonce: &str, attestation: &str) -> Result; } /// 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(coco::builtin::BuiltInCoCoAs), + + #[cfg(feature = "coco-as-grpc")] + CoCoASgRPC(GrpcClientPool), + + #[cfg(feature = "intel-trust-authority-as")] + IntelTA(IntelTrustAuthority), +} 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 async fn new(config: AsConfig) -> Result { + let built_in_as = coco::builtin::BuiltInCoCoAs::new(config).await?; + Ok(Self::CoCoASBuiltIn(built_in_as)) } /// 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 async fn new(config: GrpcConfig) -> Result { + let pool = GrpcClientPool::new(config).await?; + Ok(Self::CoCoASgRPC(pool)) } /// 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) -> Result { + let ta_client = intel_trust_authority::IntelTrustAuthority::new(config)?; + Ok(Self::IntelTA(ta_client)) + } + + pub async fn verify(&self, tee: Tee, nonce: &str, attestation: &str) -> Result { + match self { + #[cfg(feature = "coco-as-grpc")] + AttestationService::CoCoASgRPC(inner) => inner.verify(tee, nonce, attestation).await, + #[cfg(any(feature = "coco-as-builtin", feature = "coco-as-builtin-no-verifier"))] + AttestationService::CoCoASBuiltIn(inner) => inner.verify(tee, nonce, attestation).await, + #[cfg(feature = "intel-trust-authority-as")] + AttestationService::IntelTA(inner) => inner.verify(tee, nonce, attestation).await, + } + } - Ok(Self(attestation_service)) + pub async fn set_policy(&self, input: &[u8]) -> Result<()> { + match self { + #[cfg(feature = "coco-as-grpc")] + AttestationService::CoCoASgRPC(inner) => inner.set_policy(input).await, + #[cfg(any(feature = "coco-as-builtin", feature = "coco-as-builtin-no-verifier"))] + AttestationService::CoCoASBuiltIn(inner) => inner.set_policy(input).await, + #[cfg(feature = "intel-trust-authority-as")] + AttestationService::IntelTA(inner) => inner.set_policy(input).await, + } } } diff --git a/kbs/src/api/src/http/attest.rs b/kbs/src/api/src/http/attest.rs index 9d0c1477d..189b1b1a4 100644 --- a/kbs/src/api/src/http/attest.rs +++ b/kbs/src/api/src/http/attest.rs @@ -27,10 +27,7 @@ pub(crate) async fn auth( extra_params: "".to_string(), }); - map.sessions - .write() - .await - .insert(session.id().to_string(), Arc::new(Mutex::new(session))); + let _ = map.sessions.insert(session.id().to_string(), session); Ok(response) } @@ -40,30 +37,28 @@ pub(crate) async fn attest( attestation: web::Json, request: HttpRequest, map: web::Data>, - attestation_service: web::Data, + attestation_service: web::Data>, ) -> Result { let cookie = request.cookie(KBS_SESSION_ID).ok_or(Error::MissingCookie)?; - let sessions = map.sessions.read().await; - let locked_session = sessions.get(cookie.value()).ok_or(Error::InvalidCookie)?; - - let mut session = locked_session.lock().await; + let (tee, nonce) = { + let session = map + .sessions + .get_async(cookie.value()) + .await + .ok_or(Error::InvalidCookie)?; + let session = session.get(); - info!("Cookie {} attestation {:?}", session.id(), attestation); + info!("Cookie {} attestation {:?}", session.id(), attestation); - if session.is_expired() { - raise_error!(Error::ExpiredCookie); - } + if session.is_expired() { + raise_error!(Error::ExpiredCookie); + } + (session.tee(), session.nonce().to_string()) + }; let token = attestation_service - .0 - .lock() - .await - .verify( - session.tee(), - session.nonce(), - &serde_json::to_string(&attestation).unwrap(), - ) + .verify(tee, &nonce, &serde_json::to_string(&attestation).unwrap()) .await .map_err(|e| Error::AttestationFailed(e.to_string()))?; @@ -78,6 +73,12 @@ pub(crate) async fn attest( ) .map_err(|e| Error::TokenIssueFailed(format!("Illegal token base64 claims: {e}")))?; + let mut session = map + .sessions + .get_async(cookie.value()) + .await + .ok_or(Error::InvalidCookie)?; + let session = session.get_mut(); session.set_tee_public_key(attestation.tee_pubkey.clone()); session.set_authenticated(); session.set_attestation_claims(claims); diff --git a/kbs/src/api/src/http/config.rs b/kbs/src/api/src/http/config.rs index 7f8ac53f4..dcf46a136 100644 --- a/kbs/src/api/src/http/config.rs +++ b/kbs/src/api/src/http/config.rs @@ -11,7 +11,7 @@ pub(crate) async fn attestation_policy( input: web::Bytes, user_pub_key: web::Data>, insecure: web::Data, - attestation_service: web::Data, + attestation_service: web::Data>, ) -> Result { if !insecure.get_ref() { let user_pub_key = user_pub_key @@ -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/api/src/http/resource.rs b/kbs/src/api/src/http/resource.rs index dbb74b60a..071e3259c 100644 --- a/kbs/src/api/src/http/resource.rs +++ b/kbs/src/api/src/http/resource.rs @@ -138,12 +138,13 @@ async fn get_attest_claims_from_session( .cookie(KBS_SESSION_ID) .ok_or(Error::UnAuthenticatedCookie)?; - let session_map = map.sessions.read().await; - let locked_session = session_map - .get(cookie.value()) + let session = map + .sessions + .get_async(cookie.value()) + .await .ok_or(Error::UnAuthenticatedCookie)?; - let session = locked_session.lock().await; + let session = session.get(); info!("Cookie {} request to get resource", session.id()); diff --git a/kbs/src/api/src/lib.rs b/kbs/src/api/src/lib.rs index a9d0180af..029993deb 100644 --- a/kbs/src/api/src/lib.rs +++ b/kbs/src/api/src/lib.rs @@ -25,8 +25,8 @@ use jwt_simple::prelude::Ed25519PublicKey; #[cfg(feature = "resource")] use resource::RepositoryConfig; use semver::{BuildMetadata, Prerelease, Version, VersionReq}; -use std::net::SocketAddr; use std::path::PathBuf; +use std::{net::SocketAddr, sync::Arc}; #[cfg(feature = "resource")] use token::AttestationTokenVerifierType; @@ -105,7 +105,7 @@ pub struct ApiServer { insecure: bool, #[cfg(feature = "as")] - attestation_service: AttestationService, + attestation_service: Arc, http_timeout: i64, insecure_api: bool, @@ -126,7 +126,7 @@ impl ApiServer { certificate: Option, insecure: bool, - #[cfg(feature = "as")] attestation_service: &AttestationService, + #[cfg(feature = "as")] attestation_service: AttestationService, http_timeout: i64, insecure_api: bool, @@ -152,7 +152,7 @@ impl ApiServer { insecure, #[cfg(feature = "as")] - attestation_service: attestation_service.clone(), + attestation_service: Arc::new(attestation_service), http_timeout, insecure_api, diff --git a/kbs/src/api/src/session.rs b/kbs/src/api/src/session.rs index 0d4bcc0f7..ec6d4d269 100644 --- a/kbs/src/api/src/session.rs +++ b/kbs/src/api/src/session.rs @@ -12,9 +12,6 @@ use base64::Engine; use kbs_types::{Request, Tee, TeePubKey}; use rand::{thread_rng, Rng}; use semver::Version; -use std::collections::HashMap; -use std::sync::Arc; -use tokio::sync::{Mutex, RwLock}; use uuid::Uuid; pub(crate) static KBS_SESSION_ID: &str = "kbs-session-id"; @@ -123,13 +120,13 @@ impl<'a> Session<'a> { } pub(crate) struct SessionMap<'a> { - pub sessions: RwLock>>>>, + pub sessions: scc::HashMap>, } impl<'a> SessionMap<'a> { pub fn new() -> Self { SessionMap { - sessions: RwLock::new(HashMap::new()), + sessions: scc::HashMap::new(), } } } diff --git a/kbs/src/kbs/src/main.rs b/kbs/src/kbs/src/main.rs index 021844d62..c03b669aa 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()).await? } 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()).await? } 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."); } @@ -61,7 +61,7 @@ async fn main() -> Result<()> { kbs_config.certificate, kbs_config.insecure_http, #[cfg(feature = "as")] - &attestation_service, + attestation_service, kbs_config.timeout, kbs_config.insecure_api, #[cfg(feature = "resource")]