kbs/as: use global kv-storage registry

Every key value storage instance should be alloced via
`register_namespace` function. The global registry will help
1. provide an easy way to get a key value storage
2. manage instances

Signed-off-by: Xynnn007 <xynnn@linux.alibaba.com>

Author: Xynnn007

attestation-service/src/lib.rs

@@ -126,22 +126,17 @@ pub struct AttestationService {
 impl AttestationService {
     /// Create a new Attestation Service instance.
     pub async fn new(config: Config) -> Result<Self, ServiceError> {
-        let rvps = rvps::initialize_rvps_client(
-            &config.rvps_config,
-            config.storage_backend.storage_type,
-            &config.storage_backend.backends,
+        let rvps =
+            rvps::initialize_rvps_client(&config.rvps_config, config.storage_backend.clone())
+                .await
+                .map_err(ServiceError::Rvps)?;
+
+        let policy_storage = key_value_storage::register_namespace(
+            AS_POLICY_STORAGE_NAMESPACE,
+            &config.storage_backend,
         )
         .await
-        .map_err(ServiceError::Rvps)?;
-
-        let policy_storage = config
-            .storage_backend
-            .backends
-            .to_client_with_namespace(
-                config.storage_backend.storage_type,
-                AS_POLICY_STORAGE_NAMESPACE,
-            )
-            .await?;
+        .context("initialize policy storage")?;
         let token_broker =
             EarAttestationTokenBroker::new(config.attestation_token_broker.clone(), policy_storage)
                 .await?;

attestation-service/src/rvps/builtin.rs

@@ -2,7 +2,7 @@ use super::{Result, RvpsApi};
 use anyhow::Context;
 use async_trait::async_trait;
 use core::result::Result::Ok;
-use key_value_storage::{KeyValueStorageStructConfig, KeyValueStorageType};
+use key_value_storage::StorageBackendConfig;
 use reference_value_provider_service::{
     extractors::ExtractorsConfig, Rvps, REFERENCE_VALUE_STORAGE_NAMESPACE,
 };
@@ -15,13 +15,14 @@ pub struct BuiltinRvps {
 impl BuiltinRvps {
     pub async fn new(
         config: Option<ExtractorsConfig>,
-        storage_type: KeyValueStorageType,
-        storage_config: &KeyValueStorageStructConfig,
+        storage_backend_config: &StorageBackendConfig,
     ) -> Result<Self> {
-        let storage = storage_config
-            .to_client_with_namespace(storage_type, REFERENCE_VALUE_STORAGE_NAMESPACE)
-            .await
-            .context("initialize RVPS storage")?;
+        let storage = key_value_storage::register_namespace(
+            REFERENCE_VALUE_STORAGE_NAMESPACE,
+            storage_backend_config,
+        )
+        .await
+        .context("initialize RVPS storage")?;
         let rvps = Rvps::new_with_storage(config, storage)
             .await
             .context("initialize RVPS")?;

attestation-service/src/rvps/mod.rs

@@ -5,7 +5,7 @@
 
 use std::sync::Arc;
 
-use key_value_storage::{KeyValueStorageStructConfig, KeyValueStorageType};
+use key_value_storage::{KeyValueStorageType, StorageBackendConfig};
 pub use reference_value_provider_service::config::Config as RvpsCrateConfig;
 use reference_value_provider_service::extractors::ExtractorsConfig;
 use serde::Deserialize;
@@ -85,8 +85,7 @@ pub type RvpsClient = Arc<Mutex<dyn RvpsApi + Send + Sync>>;
 #[instrument(skip_all, name = "Initialize RVPS")]
 pub async fn initialize_rvps_client(
     config: &RvpsConfig,
-    unified_storage_type: KeyValueStorageType,
-    storage_config: &KeyValueStorageStructConfig,
+    mut storage_backend_config: StorageBackendConfig,
 ) -> Result<RvpsClient> {
     match config {
         RvpsConfig::BuiltIn {
@@ -95,12 +94,12 @@ pub async fn initialize_rvps_client(
         } => {
             info!("launch a built-in RVPS.");
 
-            // Use RVPS-specific storage if provided, otherwise fall back to defaults
-            let actual_storage_type = storage_type.unwrap_or(unified_storage_type);
+            if let Some(storage_type) = storage_type {
+                storage_backend_config.storage_type = *storage_type;
+            }
 
             Ok(Arc::new(Mutex::new(
-                builtin::BuiltinRvps::new(extractors.clone(), actual_storage_type, storage_config)
-                    .await?,
+                builtin::BuiltinRvps::new(extractors.clone(), &storage_backend_config).await?,
             )) as RvpsClient)
         }
         #[cfg(feature = "rvps-grpc")]

kbs/src/api_server.rs

@@ -90,12 +90,11 @@ impl ApiServer {
             "Using storage backend type: {:?}",
             config.storage_backend.storage_type
         );
-        let policy_storage_backend = config
-            .storage_backend
-            .backends
-            .to_client_with_namespace(config.storage_backend.storage_type, KBS_STORAGE_NAMESPACE)
-            .await
-            .map_err(|e| Error::StorageBackendInitialization { source: e })?;
+        let policy_storage_backend =
+            key_value_storage::register_namespace(KBS_STORAGE_NAMESPACE, &config.storage_backend)
+                .await
+                .map_err(|e| Error::StorageBackendInitialization { source: e })?;
+
         let policy_engine = PolicyEngine::new(policy_storage_backend);
 
         policy_engine

kbs/src/plugins/implementations/resource/backend.rs

@@ -111,13 +111,12 @@ impl ResourceStorage {
     ) -> Result<Self> {
         match value {
             RepositoryConfig::KvStorage => {
-                let storage = storage_backend_config
-                    .backends
-                    .to_client_with_namespace(
-                        storage_backend_config.storage_type,
-                        RESOURCE_STORAGE_NAMESPACE,
-                    )
-                    .await?;
+                let storage = key_value_storage::register_namespace(
+                    RESOURCE_STORAGE_NAMESPACE,
+                    storage_backend_config,
+                )
+                .await?;
+
                 let backend = kv_storage::KvStorage::new(storage);
                 Ok(Self {
                     backend: Arc::new(backend),