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")]