diff --git a/ai-docs/ARCHITECTURE.md b/ai-docs/ARCHITECTURE.md index 88193484e5..ff17f80d04 100644 --- a/ai-docs/ARCHITECTURE.md +++ b/ai-docs/ARCHITECTURE.md @@ -351,7 +351,7 @@ store. Epoch IDs are compared as big-endian integers. ### Centralized -A single `RealCentralizedKms` instance holds all key material. No MPC; keys +A single `CentralizedKms` instance holds all key material. No MPC; keys live in the configured vault backend. Preprocessing / reshare RPCs are not applicable. diff --git a/core/service/src/bin/kms-server.rs b/core/service/src/bin/kms-server.rs index 0d99eb7ce9..8ccd0c3d15 100644 --- a/core/service/src/bin/kms-server.rs +++ b/core/service/src/bin/kms-server.rs @@ -20,7 +20,7 @@ use kms_lib::{ signatures::NodeSigningIdentity, }, engine::{ - backup_operator::boot_base_kms, centralized::central_kms::RealCentralizedKms, + backup_operator::boot_base_kms, centralized::central_kms::CentralizedKms, context::SoftwareVersion, context_manager::create_default_centralized_context_in_storage, migration::migrate_to_0_15_x, rng_source::RngSource, run_server, threshold::service::new_real_threshold_kms, @@ -680,7 +680,7 @@ async fn main_exec() -> anyhow::Result<()> { .await?; // A node without its signing key boots in recovery mode, as a threshold node does, // so that it can recover its keys from the custodians. - let (kms, (health, health_service)) = RealCentralizedKms::new_from_base_kms( + let (kms, (health, health_service)) = CentralizedKms::new_from_base_kms( core_config, public_vault, private_vault, diff --git a/core/service/src/client/test_tools.rs b/core/service/src/client/test_tools.rs index 3f1f6d1a9d..9f5e82eb8e 100644 --- a/core/service/src/client/test_tools.rs +++ b/core/service/src/client/test_tools.rs @@ -7,10 +7,11 @@ use crate::consts::{DEC_CAPACITY, DEFAULT_PROTOCOL, DEFAULT_URL, MAX_TRIES, MIN_ use crate::cryptography::signatures::PublicSigKey; use crate::engine::backup_operator::boot_base_kms; use crate::engine::base::BaseKmsStruct; -use crate::engine::centralized::central_kms::RealCentralizedKms; +use crate::engine::centralized::central_kms::CentralizedKms; use crate::engine::context_manager::create_default_centralized_context_in_storage; use crate::engine::rng_source::test_rng_source; -use crate::engine::threshold::service::{RealThresholdKms, new_real_threshold_kms}; +use crate::engine::threshold::service::new_real_threshold_kms; +use crate::engine::threshold::threshold_kms::ThresholdKms; use crate::engine::{Shutdown, run_server}; use crate::grpc::MetaStoreStatusServiceImpl; use crate::util::rate_limiter::RateLimiterConfig; @@ -393,7 +394,7 @@ pub async fn setup_threshold_with_custom_peers< // Note: explicit some of the types to avoid clippy complaining let server: anyhow::Result<( - RealThresholdKms, + ThresholdKms, (HealthState, _), MetaStoreStatusServiceImpl, )> = new_real_threshold_kms( @@ -755,7 +756,7 @@ pub async fn setup_centralized_no_client< let config_path = format!("{}/config/default_centralized", env!("CARGO_MANIFEST_DIR")); let mut core_config: CoreConfig = init_conf(&config_path).expect("config must parse"); core_config.rate_limiter_conf = rate_limiter_conf; - let (kms, (health, health_service)) = RealCentralizedKms::new( + let (kms, (health, health_service)) = CentralizedKms::new( core_config, pub_storage, priv_storage, @@ -790,9 +791,8 @@ pub async fn setup_centralized_no_client< .await .expect("Could not start server"); }); - let service_name = , - > as NamedService>::NAME; + let service_name = + > as NamedService>::NAME; await_server_ready(service_name, listen_port).await; ServerHandle::new_centralized(arc_kms_clone, listen_port, tx, handle_health) } @@ -901,17 +901,16 @@ pub async fn setup_recovery_mode< .unwrap(); let config_path = format!("{}/config/default_centralized", env!("CARGO_MANIFEST_DIR")); let core_config: CoreConfig = init_conf(&config_path).expect("config must parse"); - let (kms, (health, health_service)) = - RealCentralizedKms::::new_from_base_kms( - core_config, - pub_storage, - priv_storage, - Some(backup_vault), - None, - base_kms, - ) - .await - .expect("a server without its signing key must boot in recovery mode"); + let (kms, (health, health_service)) = CentralizedKms::::new_from_base_kms( + core_config, + pub_storage, + priv_storage, + Some(backup_vault), + None, + base_kms, + ) + .await + .expect("a server without its signing key must boot in recovery mode"); let kms = Arc::new(kms); let server = Arc::clone(&kms); let handle_health = health.clone(); @@ -999,9 +998,8 @@ pub async fn setup_recovery_mode< } }; // The service name does not depend on the type parameters of the server. - let service_name = , - > as NamedService>::NAME; + let service_name = + > as NamedService>::NAME; await_server_ready(service_name, service_port).await; let uri = Uri::from_str(&format!( "{DEFAULT_PROTOCOL}://{DEFAULT_URL}:{service_port}" diff --git a/core/service/src/client/tests/centralized/misc_tests.rs b/core/service/src/client/tests/centralized/misc_tests.rs index 50f2a6b244..cca98e7bc3 100644 --- a/core/service/src/client/tests/centralized/misc_tests.rs +++ b/core/service/src/client/tests/centralized/misc_tests.rs @@ -6,7 +6,7 @@ use crate::client::tests::common::{PollConfig, retrying_poll}; use crate::client::tests::common::{get_pub_dec_resp, send_dec_reqs}; use crate::consts::TEST_CENTRAL_KEY_ID; -use crate::engine::centralized::central_kms::RealCentralizedKms; +use crate::engine::centralized::central_kms::CentralizedKms; use crate::testing::prelude::*; use crate::testing::utils::{get_health_client, get_status}; use kms_grpc::kms_service::v1::core_service_endpoint_server::CoreServiceEndpointServer; @@ -44,9 +44,8 @@ async fn test_central_health_endpoint_availability() -> Result<()> { let mut health_client = get_health_client(env.server.service_port) .await .expect("Failed to get health client"); - let service_name = , - > as NamedService>::NAME; + let service_name = + > as NamedService>::NAME; let request = tonic::Request::new(HealthCheckRequest { service: service_name.to_string(), }); @@ -103,9 +102,8 @@ async fn test_central_close_after_drop() -> Result<()> { let mut health_client = get_health_client(kms_server.service_port) .await .expect("Failed to get health client"); - let service_name = , - > as NamedService>::NAME; + let service_name = + > as NamedService>::NAME; let request = tonic::Request::new(HealthCheckRequest { service: service_name.to_string(), }); diff --git a/core/service/src/client/tests/threshold/misc_tests.rs b/core/service/src/client/tests/threshold/misc_tests.rs index e34f9f4515..720dd6b988 100644 --- a/core/service/src/client/tests/threshold/misc_tests.rs +++ b/core/service/src/client/tests/threshold/misc_tests.rs @@ -10,7 +10,7 @@ use crate::client::tests::common::send_dec_reqs; use crate::client::tests::common::{PollConfig, retrying_poll}; use crate::consts::TEST_THRESHOLD_KEY_ID_4P; use crate::consts::{DEFAULT_EPOCH_ID, DEFAULT_MPC_CONTEXT}; -use crate::engine::threshold::service::RealThresholdKms; +use crate::engine::threshold::threshold_kms::ThresholdKms; use crate::engine::utils::make_extra_data; use crate::testing::material::{material_subdir, threshold_material_subdir}; use crate::testing::prelude::*; @@ -59,9 +59,8 @@ async fn test_threshold_health_endpoint_availability() -> Result<()> { let servers = env.servers; // Wait for all core servers to be ready before sending requests - let core_service_name = , - > as NamedService>::NAME; + let core_service_name = + > as NamedService>::NAME; for cur_handle in servers.values() { await_server_ready(core_service_name, cur_handle.service_port).await; } @@ -183,9 +182,8 @@ async fn test_threshold_close_after_drop() -> Result<()> { let mut core_health_client = get_health_client(servers.get(&1).unwrap().service_port) .await .expect("Failed to get core health client"); - let core_service_name = , - > as NamedService>::NAME; + let core_service_name = + > as NamedService>::NAME; // Get health client for MPC threshold service on server 1 let mut threshold_health_client = get_health_client(servers.get(&1).unwrap().mpc_port.unwrap()) @@ -257,9 +255,8 @@ async fn test_threshold_shutdown() -> Result<()> { let mut servers = env.servers; // Ensure that the servers are ready - let core_service_name = , - > as NamedService>::NAME; + let core_service_name = + > as NamedService>::NAME; for cur_handle in servers.values() { await_server_ready(core_service_name, cur_handle.service_port).await; } diff --git a/core/service/src/engine/backup_operator.rs b/core/service/src/engine/backup_operator.rs index 560f5b4894..b4e1903cc8 100644 --- a/core/service/src/engine/backup_operator.rs +++ b/core/service/src/engine/backup_operator.rs @@ -28,8 +28,7 @@ use crate::{ signatures::{PrivateSigKey, PublicSigKey}, }, engine::{ - base::BaseKmsStruct, threshold::service::ThresholdFheKeys, traits::BackupOperator, - validation::RequestIdParsingErr, + base::BaseKmsStruct, threshold::service::ThresholdFheKeys, validation::RequestIdParsingErr, }, vault::{ Vault, VaultDataType, @@ -68,7 +67,7 @@ use tokio::sync::{Mutex, MutexGuard}; use tonic::{Request, Response}; use zeroize::Zeroizing; -pub struct RealBackupOperator< +pub(crate) struct RealBackupOperator< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, > { @@ -214,7 +213,7 @@ where .map_err(fail) } - pub fn new( + pub(crate) fn new( base_kms: BaseKmsStruct, crypto_storage: CryptoMaterialStorage, security_module: Option>, @@ -328,8 +327,7 @@ where } } -#[tonic::async_trait] -impl BackupOperator for RealBackupOperator +impl RealBackupOperator where PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, @@ -339,7 +337,7 @@ where /// A digest is attested rather than the key because the composite key exceeds the attestation /// document's [`crate::cryptography::attestation::NSM_ATTESTATION_FIELD_MAX_BYTES`] `public_key` /// field. - async fn get_operator_public_key( + pub(crate) async fn get_operator_public_key( &self, _request: Request, ) -> Result, MetricedError> { @@ -395,7 +393,7 @@ where } /// Restores the most recent custodian based backup. - async fn custodian_recovery_init( + pub(crate) async fn custodian_recovery_init( &self, request: Request, ) -> Result, MetricedError> { @@ -493,7 +491,7 @@ where /// /// Observe that the decryption key is NOT persisted on disc and in fact removed immediately after a call to `restore_from_backup` /// in order to minimize the possibility of leakage. - async fn custodian_backup_recovery( + pub(crate) async fn custodian_backup_recovery( &self, request: Request, ) -> Result, MetricedError> { @@ -693,7 +691,7 @@ where /// Observe that if secret sharing is used for backup (i.e. with a master key being shared with a set of custodians) /// then [`custodian_recovery`] _must_ be called first in order to ensure that the master key is restored, /// which is needed to allow decryption of the backup data. - async fn restore_from_backup( + pub(crate) async fn restore_from_backup( &self, _request: Request, ) -> Result, MetricedError> { @@ -728,7 +726,7 @@ where } } - async fn get_key_material_availability( + pub(crate) async fn get_key_material_availability( &self, _request: Request, ) -> Result, MetricedError> { diff --git a/core/service/src/engine/base.rs b/core/service/src/engine/base.rs index 4824a0bfae..64df8b4a55 100644 --- a/core/service/src/engine/base.rs +++ b/core/service/src/engine/base.rs @@ -1,12 +1,11 @@ pub use super::signed_payload::UserDecSignedPayload; use super::signed_payload::user_dec_payload; -use super::traits::BaseKms; use crate::consts::ID_LENGTH; use crate::consts::SAFE_SER_SIZE_LIMIT; use crate::cryptography::decompression; use crate::cryptography::internal_crypto_types::WrappedDKGParams; +use crate::cryptography::signatures::PublicSigKey; use crate::cryptography::signatures::internal_sign; -use crate::cryptography::signatures::{PublicSigKey, Signature}; use crate::cryptography::signing::SigningSchemeType; use crate::cryptography::signing::identity::NodeSigningIdentity; use crate::cryptography::signing::typed_signature::StoredTypedSignature; @@ -19,7 +18,7 @@ use alloy_primitives::U256; use alloy_primitives::{Address, B256, Bytes, FixedBytes, Uint}; use alloy_sol_types::Eip712Domain; use alloy_sol_types::SolStruct; -use hashing::{DomainSep, hash_element, hash_versioned, serialize_hash_element}; +use hashing::{DomainSep, hash_versioned, serialize_hash_element}; use kms_grpc::RequestId; use kms_grpc::kms::v1::{ CiphertextFormat, FheParameter, PublicDecryptionResponsePayload, TypedPlaintext, @@ -1193,26 +1192,6 @@ impl BaseKmsStruct { } } -impl BaseKms for BaseKmsStruct { - /// sign `msg` using the KMS' private signing key - fn sign(&self, dsep: &DomainSep, msg: &T) -> anyhow::Result - where - T: Serialize + AsRef<[u8]>, - { - match self.signing_identity.as_ref() { - None => anyhow::bail!("KMS has no signing key"), - Some(identity) => internal_sign(dsep, msg, identity.ecdsa()), - } - } - - fn digest(domain_separator: &DomainSep, msg: &T) -> anyhow::Result> - where - T: ?Sized + AsRef<[u8]>, - { - Ok(hash_element(domain_separator, msg)) - } -} - /// ABI encodes a list of typed plaintexts into a single byte vector for Ethereum compatibility. /// This follows the encoding pattern used in the JavaScript version for decrypted results and also supports `ebytes`. /// This function is NOT compatible with fhevm v0.9.0 and is only intended for future use with fhevm supporting `ebytes`. diff --git a/core/service/src/engine/centralized/central_kms.rs b/core/service/src/engine/centralized/central_kms.rs index ed9a2d1ec7..3fc1c6d70f 100644 --- a/core/service/src/engine/centralized/central_kms.rs +++ b/core/service/src/engine/centralized/central_kms.rs @@ -8,7 +8,7 @@ use crate::cryptography::attestation::SecurityModuleProxy; use crate::cryptography::decompression; use crate::cryptography::encryption::UnifiedPublicEncKey; use crate::cryptography::signatures::StoredTypedSignature; -use crate::cryptography::signatures::{PrivateSigKey, PublicSigKey, Signature}; +use crate::cryptography::signatures::{PrivateSigKey, PublicSigKey}; use crate::cryptography::signcryption::SigncryptFHEPlaintext; use crate::cryptography::signcryption::UnifiedSigncryptionKey; use crate::cryptography::signing::SigningSchemeType; @@ -25,8 +25,6 @@ use crate::engine::rng_source::RngSource; use crate::engine::storage_material_verification::{ PrivateLayout, verify_private_storage_layout, verify_storage_material, }; -use crate::engine::traits::{BackupOperator, ContextManager}; -use crate::engine::traits::{BaseKms, Kms}; use crate::engine::validation::DSEP_USER_DECRYPTION; use crate::grpc::metastore_status_service::CustodianMetaStore; use crate::util::key_setup::FhePublicKey; @@ -48,7 +46,6 @@ use crate::vault::storage::{ }; use crate::vault::{Vault, adopt_custodian_context, storage::Storage}; use aes_prng::AesRng; -use hashing::DomainSep; use kms_grpc::RequestId; use kms_grpc::identifiers::EpochId; use kms_grpc::kms::v1::TypedSigncryptedCiphertext; @@ -59,7 +56,6 @@ use kms_grpc::rpc_types::KMSType; use kms_grpc::rpc_types::PrivDataType; use observability::metrics::METRICS; use rand::{CryptoRng, Rng, RngCore}; -use serde::Serialize; use std::collections::HashMap; use std::sync::Arc; use std::{fmt, panic}; @@ -419,7 +415,7 @@ pub(crate) fn gen_centralized_crs( } // We only need to derive (de)serialize for test, which is why they're under a cfg_attr. -#[cfg_attr(test, derive(Serialize, serde::Deserialize))] +#[cfg_attr(test, derive(serde::Serialize, serde::Deserialize))] pub struct CentralizedKmsKeys { pub key_info: HashMap, pub sig_sk: PrivateSigKey, @@ -428,7 +424,7 @@ pub struct CentralizedKmsKeys { // We only need to derive (de)serialize for test, which is why they're under a cfg_attr. #[cfg(test)] -#[cfg_attr(test, derive(Serialize, serde::Deserialize))] +#[cfg_attr(test, derive(serde::Serialize, serde::Deserialize))] pub(crate) struct CentralizedTestingKeys { pub(crate) params: DKGParams, pub(crate) centralized_kms_keys: CentralizedKmsKeys, @@ -451,8 +447,6 @@ pub struct CentralizedPreprocBucket { pub struct CentralizedKms< PubS: Storage + Send + Sync + 'static, PrivS: StorageExt + Send + Sync + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, > { pub(crate) base_kms: BaseKmsStruct, pub(crate) crypto_storage: CentralizedCryptoMaterialStorage, @@ -475,8 +469,8 @@ pub struct CentralizedKms< // Map of ongoing CRS generation tasks, indexed by the CRS request ID pub(crate) ongoing_crs_gen: Arc>>, pub(crate) custodian_meta_map: Arc>, - pub(crate) context_manager: CM, - pub(crate) backup_operator: BO, + pub(crate) context_manager: CentralizedContextManager, + pub(crate) backup_operator: RealBackupOperator, // Rate limiting pub(crate) rate_limiter: RateLimiter, // Liveness and readiness that the gRPC health service reports @@ -484,12 +478,6 @@ pub struct CentralizedKms< // Task tacker to ensure that we keep track of all ongoing operations and can cancel them if needed (e.g. during shutdown). pub(crate) tracker: Arc, } -pub type RealCentralizedKms = CentralizedKms< - PubS, - PrivS, - CentralizedContextManager, - RealBackupOperator, ->; /// Perform asynchronous decryption and serialize the result pub fn central_public_decrypt< @@ -515,7 +503,7 @@ pub fn central_public_decrypt< let fhe_type = ct.fhe_type()?; let fhe_type_string = ct.fhe_type_string(); inner_timer.tag(TAG_TFHE_TYPE, fhe_type_string); - RealCentralizedKms::::public_decrypt( + CentralizedKms::::public_decrypt( keys, &ct.ciphertext, fhe_type, @@ -563,7 +551,7 @@ pub async fn async_user_decrypt< inner_timer.tag(TAG_TFHE_TYPE, fhe_type_string); let ct_format = typed_ciphertext.ciphertext_format(); let external_handle = typed_ciphertext.external_handle.clone(); - let signcrypted_ciphertext = RealCentralizedKms::::user_decrypt( + let signcrypted_ciphertext = CentralizedKms::::user_decrypt( keys, identity.ecdsa(), rng, @@ -615,41 +603,14 @@ pub async fn async_user_decrypt< } // impl fmt::Debug for CentralizedKms, we don't want to include the decryption key in the debug output -impl< - PubS: Storage + Sync + Send + 'static, - PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, -> fmt::Debug for CentralizedKms +impl fmt::Debug + for CentralizedKms { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("CentralizedKms").finish() // Don't include fhe_dec_key or signing key } } -impl< - PubS: Storage + Sync + Send + 'static, - PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, -> BaseKms for CentralizedKms -{ - fn sign>( - &self, - dsep: &DomainSep, - msg: &T, - ) -> anyhow::Result { - self.base_kms.sign(dsep, msg) - } - - fn digest>( - domain_separator: &DomainSep, - msg: &T, - ) -> anyhow::Result> { - BaseKmsStruct::digest(domain_separator, &msg) - } -} - macro_rules! deserialize_to_low_level_and_decrypt_helper { ($rust_type:ty,$fout:expr,$ct_format:expr,$serialized_high_level:expr,$keys:expr) => {{ match $ct_format { @@ -857,14 +818,10 @@ fn unsafe_decrypt( Ok(res) } -impl< - PubS: Storage + Sync + Send + 'static, - PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, -> Kms for CentralizedKms +impl + CentralizedKms { - fn public_decrypt( + pub(crate) fn public_decrypt( keys: &KmsFheKeyHandles, high_level_ct: &[u8], fhe_type: FheTypes, @@ -876,7 +833,8 @@ impl< } } - fn user_decrypt( + #[expect(clippy::too_many_arguments)] + pub(crate) fn user_decrypt( keys: &KmsFheKeyHandles, sig_key: &PrivateSigKey, rng: &mut (impl CryptoRng + RngCore), @@ -911,12 +869,8 @@ impl< } } -impl< - PubS: Storage + Sync + Send + 'static, - PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, -> CentralizedKms +impl + CentralizedKms { pub async fn new( config: CoreConfig, @@ -926,12 +880,12 @@ impl< security_module: Option>, signing_identity: NodeSigningIdentity, ) -> anyhow::Result<( - RealCentralizedKms, + CentralizedKms, (HealthState, HealthServer), )> { let rng_source = Arc::new(RngSource::new(security_module.clone())?); let base_kms = BaseKmsStruct::new(KMSType::Centralized, signing_identity, rng_source); - RealCentralizedKms::::new_from_base_kms( + CentralizedKms::::new_from_base_kms( config, public_storage, private_storage, @@ -955,7 +909,7 @@ impl< security_module: Option>, base_kms: BaseKmsStruct, ) -> anyhow::Result<( - RealCentralizedKms, + CentralizedKms, (HealthState, HealthServer), )> { let key_info_with_epoch: HashMap<(RequestId, EpochId), KmsFheKeyHandles> = @@ -1094,12 +1048,8 @@ impl< } } -impl< - PubS: Storage + Sync + Send + 'static, - PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, -> CentralizedKms +impl + CentralizedKms { /// Get a reference to the key generation MetaStore pub fn get_key_gen_meta_store(&self) -> &Arc>> { @@ -1127,12 +1077,8 @@ impl< } #[tonic::async_trait] -impl< - PubS: Storage + Sync + Send + 'static, - PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, -> Shutdown for CentralizedKms +impl Shutdown + for CentralizedKms { fn shutdown(&self) -> anyhow::Result> { let health = self.health.clone(); @@ -1150,12 +1096,8 @@ impl< } } -impl< - PubS: Storage + Sync + Send + 'static, - PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, -> Drop for CentralizedKms +impl Drop + for CentralizedKms { fn drop(&mut self) { // Let the shutdown run in the background @@ -1227,8 +1169,7 @@ pub(crate) mod tests { use crate::cryptography::signing::identity::NodeSigningIdentity; use crate::dummy_domain; use crate::engine::base::{KmsFheKeyHandles, derive_request_id}; - use crate::engine::centralized::central_kms::RealCentralizedKms; - use crate::engine::traits::Kms; + use crate::engine::centralized::central_kms::CentralizedKms; use crate::engine::validation::DSEP_USER_DECRYPTION; use crate::util::key_setup::test_tools::{EncryptionConfig, compute_cipher}; use crate::util::rate_limiter::RateLimiter; @@ -1272,7 +1213,7 @@ pub(crate) mod tests { } impl - RealCentralizedKms + CentralizedKms { pub(crate) fn set_bucket_size(&mut self, bucket_size: usize) { let config = crate::util::rate_limiter::RateLimiterConfig { @@ -1628,7 +1569,7 @@ pub(crate) mod tests { }; let config = init_conf("config/default_centralized.toml").unwrap(); let kms = { - let (inner, _health_service) = RealCentralizedKms::new( + let (inner, _health_service) = CentralizedKms::new( config, new_pub_ram_storage_from_existing_keys( &keys.pub_fhe_keys, @@ -1663,7 +1604,7 @@ pub(crate) mod tests { .read_centralized_fhe_keys(key_id, epoch_id) .await .unwrap(); - let raw_plaintext = RealCentralizedKms::::public_decrypt( + let raw_plaintext = CentralizedKms::::public_decrypt( &key_handle, &ct, fhe_type, @@ -1764,7 +1705,7 @@ pub(crate) mod tests { PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, >( - inner: &RealCentralizedKms, + inner: &CentralizedKms, key_id: &RequestId, epoch_id: &EpochId, params: DKGParams, @@ -1845,7 +1786,7 @@ pub(crate) mod tests { let kms = { let core_config: CoreConfig = init_conf("config/default_centralized.toml").unwrap(); - let (inner, _health_service) = RealCentralizedKms::::new( + let (inner, _health_service) = CentralizedKms::::new( core_config, new_pub_ram_storage_from_existing_keys( &keys.pub_fhe_keys, @@ -1904,7 +1845,7 @@ pub(crate) mod tests { }; let mut rng = kms.base_kms.new_rng(); - let raw_cipher = RealCentralizedKms::::user_decrypt( + let raw_cipher = CentralizedKms::::user_decrypt( &kms.crypto_storage .read_centralized_fhe_keys(key_id, epoch_id) .await @@ -1992,7 +1933,7 @@ pub(crate) mod tests { verf_key, crate::engine::rng_source::test_rng_source(), ); - let _booted = RealCentralizedKms::new_from_base_kms( + let _booted = CentralizedKms::new_from_base_kms( config, RamStorage::new(), RamStorage::new(), diff --git a/core/service/src/engine/centralized/endpoint.rs b/core/service/src/engine/centralized/endpoint.rs index 041c460a7e..871e82b9df 100644 --- a/core/service/src/engine/centralized/endpoint.rs +++ b/core/service/src/engine/centralized/endpoint.rs @@ -3,7 +3,7 @@ use crate::engine::centralized::service::{ abort_crs_gen_impl, abort_key_gen_impl, get_preprocessing_res_impl, init_impl, preprocessing_impl, }; -use crate::engine::traits::{BackupOperator, ContextManager}; +use crate::engine::traits::ContextManager; use crate::engine::utils::query_key_material_availability; use crate::vault::storage::{Storage, StorageExt}; use kms_grpc::kms::v1::{ @@ -34,12 +34,8 @@ use observability::metrics_names::OP_INSECURE_KEYGEN_REQUEST; use observability::{metrics::METRICS, metrics_names::*}; #[tonic::async_trait] -impl< - PubS: Storage + Sync + Send + 'static, - PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, -> CoreServiceEndpoint for CentralizedKms +impl + CoreServiceEndpoint for CentralizedKms { #[tracing::instrument(skip(self, request))] async fn key_gen_preproc( diff --git a/core/service/src/engine/centralized/service/crs_gen.rs b/core/service/src/engine/centralized/service/crs_gen.rs index 5c48bbdc5e..8dd79d5f75 100644 --- a/core/service/src/engine/centralized/service/crs_gen.rs +++ b/core/service/src/engine/centralized/service/crs_gen.rs @@ -20,7 +20,7 @@ use crate::cryptography::signing::SigningSchemeType; use crate::cryptography::signing::identity::NodeSigningIdentity; use crate::engine::base::{CrsGenMetadata, stored_scheme_signatures_to_proto}; use crate::engine::centralized::central_kms::{CentralizedKms, async_generate_crs}; -use crate::engine::traits::{BackupOperator, ContextManager}; +use crate::engine::traits::ContextManager; use crate::engine::utils::{MetricedError, signing_identity_for}; use crate::engine::validation::{ RequestIdParsingErr, parse_grpc_request_id, validate_crs_gen_request, @@ -36,10 +36,8 @@ use crate::vault::storage::{Storage, StorageExt}; pub async fn crs_gen_impl< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, >( - service: &CentralizedKms, + service: &CentralizedKms, request: Request, insecure: bool, ) -> Result, MetricedError> { @@ -133,10 +131,8 @@ pub async fn crs_gen_impl< pub async fn get_crs_gen_result_impl< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, >( - service: &CentralizedKms, + service: &CentralizedKms, request: Request, insecure: bool, ) -> Result, MetricedError> { @@ -201,10 +197,8 @@ pub async fn get_crs_gen_result_impl< pub async fn abort_crs_gen_impl< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, >( - service: &CentralizedKms, + service: &CentralizedKms, request: Request, ) -> Result, MetricedError> { let request_id = parse_grpc_request_id(&request.into_inner(), RequestIdParsingErr::CrsGenAbort) diff --git a/core/service/src/engine/centralized/service/decryption.rs b/core/service/src/engine/centralized/service/decryption.rs index b4832f841e..74df6962b7 100644 --- a/core/service/src/engine/centralized/service/decryption.rs +++ b/core/service/src/engine/centralized/service/decryption.rs @@ -4,7 +4,7 @@ use crate::engine::base::{PubDecCallValues, UserDecryptCallValues, sign_public_d use crate::engine::centralized::central_kms::{ CentralizedKms, async_user_decrypt, central_public_decrypt, }; -use crate::engine::traits::{BackupOperator, ContextManager}; +use crate::engine::traits::ContextManager; use crate::engine::utils::{MetricedError, format_unvalidated_id, signing_identity_for}; use crate::engine::validation::{ RequestIdParsingErr, parse_grpc_request_id, parse_optional_grpc_request_id, @@ -42,10 +42,8 @@ use tracing::Instrument; pub async fn user_decrypt_impl< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, >( - service: &CentralizedKms, + service: &CentralizedKms, request: Request, ) -> Result, MetricedError> { METRICS.increment_request_counter(OP_USER_DECRYPT_REQUEST); @@ -164,10 +162,8 @@ pub async fn user_decrypt_impl< pub async fn user_decrypt_sync_impl< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, >( - service: &CentralizedKms, + service: &CentralizedKms, request: Request, ) -> Result, MetricedError> { // `user_decrypt_impl` consumes the request, so keep the raw id for fetching the result below. @@ -194,10 +190,8 @@ pub async fn user_decrypt_sync_impl< pub async fn get_user_decryption_result_impl< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, >( - service: &CentralizedKms, + service: &CentralizedKms, request: Request, ) -> Result, MetricedError> { METRICS.increment_request_counter(OP_USER_DECRYPT_RESULT); @@ -250,10 +244,8 @@ pub async fn get_user_decryption_result_impl< pub async fn public_decrypt_impl< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, >( - service: &CentralizedKms, + service: &CentralizedKms, request: Request, ) -> Result, MetricedError> { METRICS.increment_request_counter(OP_PUBLIC_DECRYPT_REQUEST); @@ -406,10 +398,8 @@ pub async fn public_decrypt_impl< pub async fn public_decrypt_sync_impl< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, >( - service: &CentralizedKms, + service: &CentralizedKms, request: Request, ) -> Result, MetricedError> { // `public_decrypt_impl` consumes the request, so keep the raw id for fetching the result below @@ -436,10 +426,8 @@ pub async fn public_decrypt_sync_impl< pub async fn get_public_decryption_result_impl< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, >( - service: &CentralizedKms, + service: &CentralizedKms, request: Request, ) -> Result, MetricedError> { METRICS.increment_request_counter(OP_PUBLIC_DECRYPT_RESULT); @@ -510,7 +498,7 @@ pub(crate) mod tests { use crate::{ cryptography::signatures::PublicSigKey, engine::centralized::{ - central_kms::RealCentralizedKms, + central_kms::CentralizedKms, service::key_gen::tests::{setup_test_kms_with_preproc, test_standard_keygen}, }, util::key_setup::test_tools::{EncryptionConfig, TestingPlaintext, compute_cipher}, @@ -523,7 +511,7 @@ pub(crate) mod tests { rng: &mut AesRng, key_id: &RequestId, ) -> ( - RealCentralizedKms, + CentralizedKms, tfhe::CompactPublicKey, PublicSigKey, ) { diff --git a/core/service/src/engine/centralized/service/initiator.rs b/core/service/src/engine/centralized/service/initiator.rs index a5ed6d9277..4db3adbf0f 100644 --- a/core/service/src/engine/centralized/service/initiator.rs +++ b/core/service/src/engine/centralized/service/initiator.rs @@ -1,8 +1,6 @@ use crate::{ engine::{ - centralized::central_kms::CentralizedKms, - traits::{BackupOperator, ContextManager}, - utils::MetricedError, + centralized::central_kms::CentralizedKms, traits::ContextManager, utils::MetricedError, validation::validate_new_mpc_epoch_request, }, util::meta_store::{add_req_to_meta_store, update_req_in_meta_store}, @@ -19,7 +17,7 @@ use tonic::{Request, Response}; /// Thus initialization is only allowed once and the request ID supplied in [`InitRequest`] must be valid. /// /// # Arguments -/// - `service`: Reference to the `RealCentralizedKms` instance. +/// - `service`: Reference to the `CentralizedKms` instance. /// - `request`: The gRPC request containing an `InitRequest`. /// /// # Returns @@ -36,10 +34,8 @@ use tonic::{Request, Response}; pub async fn init_impl< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, >( - service: &CentralizedKms, + service: &CentralizedKms, request: Request, ) -> Result, MetricedError> { let inner = request.into_inner(); diff --git a/core/service/src/engine/centralized/service/key_gen.rs b/core/service/src/engine/centralized/service/key_gen.rs index e6570bf5fe..a252f22ed9 100644 --- a/core/service/src/engine/centralized/service/key_gen.rs +++ b/core/service/src/engine/centralized/service/key_gen.rs @@ -9,7 +9,7 @@ use crate::engine::centralized::central_kms::{ async_generate_fhe_keys, }; use crate::engine::keyset_configuration::InternalKeySetConfig; -use crate::engine::traits::{BackupOperator, ContextManager}; +use crate::engine::traits::ContextManager; use crate::engine::utils::{MetricedError, signing_identity_for}; use crate::engine::validation::{ RequestIdParsingErr, parse_grpc_request_id, validate_key_gen_request, @@ -43,10 +43,8 @@ use tracing::Instrument; pub async fn key_gen_impl< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, >( - service: &CentralizedKms, + service: &CentralizedKms, request: Request, insecure: bool, ) -> Result, MetricedError> { @@ -231,10 +229,8 @@ pub async fn key_gen_impl< pub async fn get_key_gen_result_impl< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, >( - service: &CentralizedKms, + service: &CentralizedKms, request: Request, insecure: bool, ) -> Result, MetricedError> { @@ -310,10 +306,8 @@ pub async fn get_key_gen_result_impl< pub async fn abort_key_gen_impl< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, >( - service: &CentralizedKms, + service: &CentralizedKms, request: Request, ) -> Result, MetricedError> { let preproc_id = parse_grpc_request_id(&request.into_inner(), RequestIdParsingErr::KeyGenAbort) @@ -502,7 +496,7 @@ pub(crate) mod tests { engine::{ base::derive_request_id, centralized::{ - central_kms::RealCentralizedKms, + central_kms::CentralizedKms, service::{preprocessing_impl, tests::setup_central_test_kms}, }, }, @@ -519,7 +513,7 @@ pub(crate) mod tests { pub(crate) async fn setup_test_kms_with_preproc( rng: &mut AesRng, preproc_id: &RequestId, - ) -> (RealCentralizedKms, PublicSigKey) { + ) -> (CentralizedKms, PublicSigKey) { let (kms, verf_key) = setup_central_test_kms(rng).await; // insert a preproc ID @@ -544,7 +538,7 @@ pub(crate) mod tests { } pub(crate) async fn test_standard_keygen( - kms: &RealCentralizedKms, + kms: &CentralizedKms, req_id: &RequestId, preproc_id: Option<&RequestId>, insecure: bool, diff --git a/core/service/src/engine/centralized/service/mod.rs b/core/service/src/engine/centralized/service/mod.rs index b0fbe19ea3..b2a655897e 100644 --- a/core/service/src/engine/centralized/service/mod.rs +++ b/core/service/src/engine/centralized/service/mod.rs @@ -59,7 +59,7 @@ mod tests { use crate::util::key_setup::store_server_signing_keys; use crate::{ cryptography::signatures::{NodeSigningIdentity, PublicSigKey, gen_sig_keys}, - engine::centralized::central_kms::RealCentralizedKms, + engine::centralized::central_kms::CentralizedKms, vault::storage::ram::RamStorage, }; use aes_prng::AesRng; @@ -69,7 +69,7 @@ mod tests { /// This also adds a dummy context pub(crate) async fn setup_central_test_kms( rng: &mut AesRng, - ) -> (RealCentralizedKms, PublicSigKey) { + ) -> (CentralizedKms, PublicSigKey) { let (verf_key, sig_key) = gen_sig_keys(rng); let mut public_storage = RamStorage::new(); let mut private_storage = RamStorage::new(); @@ -80,7 +80,7 @@ mod tests { .await .unwrap(); let core_config: CoreConfig = init_conf("config/default_centralized.toml").unwrap(); - let (kms, _health_service) = RealCentralizedKms::new( + let (kms, _health_service) = CentralizedKms::new( core_config, public_storage, private_storage, diff --git a/core/service/src/engine/centralized/service/preprocessing.rs b/core/service/src/engine/centralized/service/preprocessing.rs index e0800a86d4..2f5c21d8ab 100644 --- a/core/service/src/engine/centralized/service/preprocessing.rs +++ b/core/service/src/engine/centralized/service/preprocessing.rs @@ -2,7 +2,6 @@ use crate::{ engine::{ base::compute_preprocessing_signatures, centralized::central_kms::{CentralizedKms, CentralizedPreprocBucket}, - traits::{BackupOperator, ContextManager}, utils::{MetricedError, signing_identity_for}, validation::{RequestIdParsingErr, parse_grpc_request_id, validate_preproc_request}, }, @@ -44,10 +43,8 @@ use tonic::{Request, Response}; pub async fn preprocessing_impl< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, >( - service: &CentralizedKms, + service: &CentralizedKms, request: Request, ) -> Result, MetricedError> { let _permit = service.rate_limiter.start_preproc().await?; @@ -134,10 +131,8 @@ pub async fn preprocessing_impl< pub async fn get_preprocessing_res_impl< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, >( - service: &CentralizedKms, + service: &CentralizedKms, request: Request, ) -> Result, MetricedError> { tracing::warn!( diff --git a/core/service/src/engine/threshold/endpoint.rs b/core/service/src/engine/threshold/endpoint.rs index d24d58b611..176e7963b7 100644 --- a/core/service/src/engine/threshold/endpoint.rs +++ b/core/service/src/engine/threshold/endpoint.rs @@ -1,12 +1,8 @@ use crate::engine::threshold::bandwidth_bench::run_bandwidth_benchmark; use crate::engine::threshold::threshold_kms::ThresholdKms; -use crate::engine::threshold::traits::{ - CrsGenerator, KeyGenPreprocessor, KeyGenerator, PublicDecryptor, UserDecryptor, -}; -#[cfg(feature = "insecure")] -use crate::engine::threshold::traits::{InsecureCrsGenerator, InsecureKeyGenerator}; -use crate::engine::traits::{BackupOperator, ContextManager, EpochManager}; +use crate::engine::traits::ContextManager; use crate::engine::validation::{RequestIdParsingErr, parse_grpc_request_id}; +use crate::vault::storage::{Storage, StorageExt}; use kms_grpc::ContextId; use kms_grpc::kms::v1::*; use kms_grpc::kms_service::v1::core_service_endpoint_server::CoreServiceEndpoint; @@ -15,492 +11,567 @@ use threshold_networking::health_check::HealthCheckStatus; use tonic::{Request, Response, Status}; -macro_rules! impl_endpoint { - { impl CoreServiceEndpoint $implementations:tt } => { - #[cfg(not(feature="insecure"))] - #[tonic::async_trait] - impl< - EP: EpochManager + Sync + Send + 'static, - UD: UserDecryptor + Sync + Send + 'static, - PD: PublicDecryptor + Sync + Send + 'static, - KG: KeyGenerator + Sync + Send + 'static, - PP: KeyGenPreprocessor + Sync + Send + 'static, - CG: CrsGenerator + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, - > CoreServiceEndpoint for ThresholdKms $implementations - - #[cfg(feature="insecure")] - #[tonic::async_trait] - impl< - EP: EpochManager + Sync + Send + 'static, - UD: UserDecryptor + Sync + Send + 'static, - PD: PublicDecryptor + Sync + Send + 'static, - KG: KeyGenerator + Sync + Send + 'static, - IKG: InsecureKeyGenerator + Sync + Send + 'static, - PP: KeyGenPreprocessor + Sync + Send + 'static, - CG: CrsGenerator + Sync + Send + 'static, - ICG: InsecureCrsGenerator + Sync + Send + 'static, - CM: ContextManager + Sync + Send + 'static, - BO: BackupOperator + Sync + Send + 'static, - > CoreServiceEndpoint for ThresholdKms $implementations +// See the proto file for the documentation of each method. +#[tonic::async_trait] +impl CoreServiceEndpoint for ThresholdKms +where + PubS: Storage + Send + Sync + 'static, + PrivS: StorageExt + Send + Sync + 'static, +{ + #[tracing::instrument(skip(self, request))] + async fn key_gen_preproc( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(OP_KEYGEN_PREPROC_REQUEST); + self.keygen_preprocessor + .key_gen_preproc(request) + .await + .map_err(|e| e.into()) } -} - -impl_endpoint! { - // See the proto file for the documentation of each method. - impl CoreServiceEndpoint { - - #[tracing::instrument(skip(self, request))] - async fn key_gen_preproc( - &self, - request: Request, - ) -> Result, Status> { - METRICS.increment_request_counter(OP_KEYGEN_PREPROC_REQUEST); - self.keygen_preprocessor.key_gen_preproc(request).await.map_err(|e| e.into()) - } - - #[cfg(feature = "insecure")] - #[tracing::instrument(skip(self, request))] - async fn partial_key_gen_preproc( - &self, - request: Request, - ) -> Result, Status> { - METRICS.increment_request_counter(OP_KEYGEN_PREPROC_REQUEST); - self.keygen_preprocessor.partial_key_gen_preproc(request).await.map_err(|e| e.into()) - } - #[cfg(feature = "insecure")] - #[tracing::instrument(skip(self, request))] - async fn insecure_key_gen_preproc( - &self, - request: Request, - ) -> Result, Status> { - METRICS.increment_request_counter(OP_INSECURE_KEYGEN_PREPROC_REQUEST); - self.keygen_preprocessor.insecure_key_gen_preproc(request).await.map_err(|e| e.into()) - } - - #[cfg(feature = "insecure")] - #[tracing::instrument(skip(self, request))] - async fn get_insecure_key_gen_preproc_result( - &self, - request: Request, - ) -> Result, Status> { - METRICS.increment_request_counter(OP_INSECURE_KEYGEN_PREPROC_RESULT); - self.keygen_preprocessor.get_insecure_result(request).await.map_err(|e| e.into()) - } + #[cfg(feature = "insecure")] + #[tracing::instrument(skip(self, request))] + async fn partial_key_gen_preproc( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(OP_KEYGEN_PREPROC_REQUEST); + self.keygen_preprocessor + .partial_key_gen_preproc(request) + .await + .map_err(|e| e.into()) + } - #[tracing::instrument(skip(self, request))] - async fn get_key_gen_preproc_result( - &self, - request: Request, - ) -> Result, Status> { - METRICS.increment_request_counter(OP_KEYGEN_PREPROC_RESULT); - self.keygen_preprocessor.get_result(request).await.map_err(|e| e.into()) - } + #[cfg(feature = "insecure")] + #[tracing::instrument(skip(self, request))] + async fn insecure_key_gen_preproc( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(OP_INSECURE_KEYGEN_PREPROC_REQUEST); + self.keygen_preprocessor + .insecure_key_gen_preproc(request) + .await + .map_err(|e| e.into()) + } - #[tracing::instrument(skip(self, request))] - async fn key_gen(&self, request: Request) -> Result, Status> { - METRICS.increment_request_counter(OP_KEYGEN_REQUEST); - self.key_generator.key_gen(request).await.map_err(|e| e.into()) - } + #[cfg(feature = "insecure")] + #[tracing::instrument(skip(self, request))] + async fn get_insecure_key_gen_preproc_result( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(OP_INSECURE_KEYGEN_PREPROC_RESULT); + self.keygen_preprocessor + .get_insecure_result(request) + .await + .map_err(|e| e.into()) + } - #[tracing::instrument(skip(self, request))] - async fn get_key_gen_result( - &self, - request: Request, - ) -> Result, Status> { - METRICS.increment_request_counter(OP_KEYGEN_RESULT); - self.key_generator.get_result(request).await.map_err(|e| e.into()) - } + #[tracing::instrument(skip(self, request))] + async fn get_key_gen_preproc_result( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(OP_KEYGEN_PREPROC_RESULT); + self.keygen_preprocessor + .get_result(request) + .await + .map_err(|e| e.into()) + } - #[tracing::instrument(skip(self, request))] - async fn abort_key_gen(&self, request: Request) -> Result, Status> { - METRICS.increment_request_counter(OP_KEYGEN_ABORT); - let preproc_id = parse_grpc_request_id( - &request.into_inner(), - RequestIdParsingErr::KeyGenAbort, - )?; - // First cancel the potential key generation, then use this in the preprocessing cancellation to return the appropriate error - let key_gen_abort_res = self.key_generator.abort_key_gen(preproc_id).await; - self.keygen_preprocessor.abort_key_gen_preproc(preproc_id, key_gen_abort_res).await.map_err(|e| e.into()) - } + #[tracing::instrument(skip(self, request))] + async fn key_gen(&self, request: Request) -> Result, Status> { + METRICS.increment_request_counter(OP_KEYGEN_REQUEST); + self.key_generator + .key_gen(request) + .await + .map_err(|e| e.into()) + } - // NOTE: unlike other endpoints, the decryption counters are incremented inside the - // shared implementation, not here: one place instead of two (sync/async), and the - // sync path calls `get_result` directly, bypassing this dispatch, so incrementing - // here would skip that counter bump entirely. - - #[tracing::instrument(skip(self, request))] - async fn user_decrypt( - &self, - request: Request, - ) -> Result, Status> { - self.user_decryptor.user_decrypt(request).await.map_err(|e| e.into()) - } + #[tracing::instrument(skip(self, request))] + async fn get_key_gen_result( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(OP_KEYGEN_RESULT); + self.key_generator + .get_result(request) + .await + .map_err(|e| e.into()) + } - #[tracing::instrument(skip(self, request))] - async fn get_user_decryption_result( - &self, - request: Request, - ) -> Result, Status> { - self.user_decryptor.get_result(request).await.map_err(|e| e.into()) - } + #[tracing::instrument(skip(self, request))] + async fn abort_key_gen(&self, request: Request) -> Result, Status> { + METRICS.increment_request_counter(OP_KEYGEN_ABORT); + let preproc_id = + parse_grpc_request_id(&request.into_inner(), RequestIdParsingErr::KeyGenAbort)?; + // First cancel the potential key generation, then use this in the preprocessing cancellation to return the appropriate error + let key_gen_abort_res = self.key_generator.abort_key_gen(preproc_id).await; + self.keygen_preprocessor + .abort_key_gen_preproc(preproc_id, key_gen_abort_res) + .await + .map_err(|e| e.into()) + } - #[tracing::instrument(skip(self, request))] - async fn user_decrypt_sync( - &self, - request: Request, - ) -> Result, Status> { - self.user_decryptor.user_decrypt_sync(request).await.map_err(|e| e.into()) - } + // NOTE: unlike other endpoints, the decryption counters are incremented inside the + // shared implementation, not here: one place instead of two (sync/async), and the + // sync path calls `get_result` directly, bypassing this dispatch, so incrementing + // here would skip that counter bump entirely. + + #[tracing::instrument(skip(self, request))] + async fn user_decrypt( + &self, + request: Request, + ) -> Result, Status> { + self.user_decryptor + .user_decrypt(request) + .await + .map_err(|e| e.into()) + } - #[tracing::instrument(skip(self, request))] - async fn public_decrypt( - &self, - request: Request, - ) -> Result, Status> { - self.decryptor.public_decrypt(request).await.map_err(|e| e.into()) - } + #[tracing::instrument(skip(self, request))] + async fn get_user_decryption_result( + &self, + request: Request, + ) -> Result, Status> { + self.user_decryptor + .get_result(request) + .await + .map_err(|e| e.into()) + } - #[tracing::instrument(skip(self, request))] - async fn get_public_decryption_result( - &self, - request: Request, - ) -> Result, Status> { - self.decryptor.get_result(request).await.map_err(|e| e.into()) - } - - #[tracing::instrument(skip(self, request))] - async fn public_decrypt_sync( - &self, - request: Request, - ) -> Result, Status> { - self.decryptor.public_decrypt_sync(request).await.map_err(|e| e.into()) - } + #[tracing::instrument(skip(self, request))] + async fn user_decrypt_sync( + &self, + request: Request, + ) -> Result, Status> { + self.user_decryptor + .user_decrypt_sync(request) + .await + .map_err(|e| e.into()) + } - #[tracing::instrument(skip(self, request))] - async fn crs_gen(&self, request: Request) -> Result, Status> { - METRICS.increment_request_counter(OP_CRS_GEN_REQUEST); - self.crs_generator.crs_gen(request).await.map_err(|e| e.into()) - } + #[tracing::instrument(skip(self, request))] + async fn public_decrypt( + &self, + request: Request, + ) -> Result, Status> { + self.decryptor + .public_decrypt(request) + .await + .map_err(|e| e.into()) + } - #[tracing::instrument(skip(self, request))] - async fn get_crs_gen_result( - &self, - request: Request, - ) -> Result, Status> { - METRICS.increment_request_counter(OP_CRS_GEN_RESULT); - self.crs_generator.get_result(request).await.map_err(|e| e.into()) - } + #[tracing::instrument(skip(self, request))] + async fn get_public_decryption_result( + &self, + request: Request, + ) -> Result, Status> { + self.decryptor + .get_result(request) + .await + .map_err(|e| e.into()) + } - #[tracing::instrument(skip(self, request))] - async fn abort_crs_gen(&self, request: Request) -> Result, Status> { - METRICS.increment_request_counter(OP_CRS_GEN_ABORT); - self.crs_generator.abort_crs_gen(request).await.map_err(|e| e.into()) - } + #[tracing::instrument(skip(self, request))] + async fn public_decrypt_sync( + &self, + request: Request, + ) -> Result, Status> { + self.decryptor + .public_decrypt_sync(request) + .await + .map_err(|e| e.into()) + } - #[cfg(feature = "insecure")] - #[tracing::instrument(skip(self, request))] - async fn insecure_key_gen(&self, request: Request) -> Result, Status> { - METRICS.increment_request_counter(observability::metrics_names::OP_INSECURE_KEYGEN_REQUEST); - self.insecure_key_generator.insecure_key_gen(request).await.map_err(|e| e.into()) - } + #[tracing::instrument(skip(self, request))] + async fn crs_gen(&self, request: Request) -> Result, Status> { + METRICS.increment_request_counter(OP_CRS_GEN_REQUEST); + self.crs_generator + .crs_gen(request) + .await + .map_err(|e| e.into()) + } - #[cfg(feature = "insecure")] - #[tracing::instrument(skip(self, request))] - async fn get_insecure_key_gen_result( - &self, - request: Request, - ) -> Result, Status> { - METRICS.increment_request_counter(observability::metrics_names::OP_INSECURE_KEYGEN_RESULT); - self.insecure_key_generator.get_result(request).await.map_err(|e| e.into()) - } + #[tracing::instrument(skip(self, request))] + async fn get_crs_gen_result( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(OP_CRS_GEN_RESULT); + self.crs_generator + .get_result(request) + .await + .map_err(|e| e.into()) + } - #[cfg(feature = "insecure")] - #[tracing::instrument(skip(self, request))] - async fn insecure_crs_gen(&self, request: Request) -> Result, Status> { - METRICS.increment_request_counter(observability::metrics_names::OP_INSECURE_CRS_GEN_REQUEST); - self.insecure_crs_generator.insecure_crs_gen(request).await.map_err(|e| e.into()) - } + #[tracing::instrument(skip(self, request))] + async fn abort_crs_gen(&self, request: Request) -> Result, Status> { + METRICS.increment_request_counter(OP_CRS_GEN_ABORT); + self.crs_generator + .abort_crs_gen(request) + .await + .map_err(|e| e.into()) + } - #[cfg(feature = "insecure")] - #[tracing::instrument(skip(self, request))] - async fn get_insecure_crs_gen_result( - &self, - request: Request, - ) -> Result, Status> { - METRICS.increment_request_counter(observability::metrics_names::OP_INSECURE_CRS_GEN_RESULT); - self.insecure_crs_generator.get_result(request).await.map_err(|e| e.into()) - } + #[cfg(feature = "insecure")] + #[tracing::instrument(skip(self, request))] + async fn insecure_key_gen( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(observability::metrics_names::OP_INSECURE_KEYGEN_REQUEST); + self.insecure_key_generator + .insecure_key_gen(request) + .await + .map_err(|e| e.into()) + } - // TODO(#2868) refactor to use MetricedError - #[tracing::instrument(skip(self, request))] - async fn new_mpc_context( - &self, - request: Request, - ) -> Result, Status> { - METRICS.increment_request_counter(OP_NEW_MPC_CONTEXT); - self.context_manager.new_mpc_context(request).await.map_err(|e| e.into()) - } + #[cfg(feature = "insecure")] + #[tracing::instrument(skip(self, request))] + async fn get_insecure_key_gen_result( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(observability::metrics_names::OP_INSECURE_KEYGEN_RESULT); + self.insecure_key_generator + .get_result(request) + .await + .map_err(|e| e.into()) + } - #[tracing::instrument(skip(self, request))] - async fn destroy_mpc_context( - &self, - request: Request, - ) -> Result, Status> { - METRICS.increment_request_counter(OP_DESTROY_MPC_CONTEXT); - let inner = request.into_inner(); - - // Validate the whole request before mutating anything: a request with valid epoch_ids - // but a missing/malformed context_id must not erase epochs and only then fail context - // parsing. - // - // Note: we extract the inner message and later re-wrap it in `Request::new(inner)`. - let proto_context_id = inner - .context_id - .as_ref() - .ok_or_else(|| Status::invalid_argument("context_id is required"))?; - let context_id = parse_grpc_request_id::(proto_context_id, RequestIdParsingErr::Context)?; - - // Hold the exclusive context lease across both phases. This makes the epoch snapshot - // stable with respect to `NewMpcEpoch`, including creations that are still running PRSS - // and therefore have not registered their epoch in the session maker yet. - let _destruction_lease = self - .session_maker - .try_start_context_destruction(&context_id) - .await - .map_err(|e| { - Status::failed_precondition(format!( - "Cannot destroy MPC context {context_id}: {e}. Retry once epoch creation has settled." - )) - })?; - - // Destroy the associated epochs first: their secret key shares and PRSS randomness are security-sensitive - // material. Erase them before touching anything else so that if there is a problem, the worst transient - // state is "shares gone, context metadata lingers" rather than "context gone, shares still on disk". - // - // If any epoch fails to delete we return here and leave the context intact. The caller is expected to retry - // the whole `DestroyMpcContext` until both epochs and context are destroyed successfully. - let epochs_destroyed = self.epoch_manager.destroy_epochs_for_context(&context_id).await.map_err(Status::from)?; - - // Every epoch is now gone, so it is safe to remove the context. - self.context_manager - .destroy_mpc_context(Request::new(inner)) - .await - .map_err(Status::from)?; - Ok(Response::new(DestroyMpcContextResponse { - epoch_ids: epochs_destroyed.into_iter().map(|id| id.into()).collect(), - })) - } + #[cfg(feature = "insecure")] + #[tracing::instrument(skip(self, request))] + async fn insecure_crs_gen( + &self, + request: Request, + ) -> Result, Status> { + METRICS + .increment_request_counter(observability::metrics_names::OP_INSECURE_CRS_GEN_REQUEST); + self.insecure_crs_generator + .insecure_crs_gen(request) + .await + .map_err(|e| e.into()) + } - #[tracing::instrument(skip_all)] - async fn new_mpc_epoch( - &self, - request: Request, - ) -> Result, Status> { - let inner = request.into_inner(); - METRICS.increment_request_counter(OP_NEW_EPOCH); - Ok(self.epoch_manager - .new_mpc_epoch(Request::new(inner.clone())) - .await?) - } + #[cfg(feature = "insecure")] + #[tracing::instrument(skip(self, request))] + async fn get_insecure_crs_gen_result( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(observability::metrics_names::OP_INSECURE_CRS_GEN_RESULT); + self.insecure_crs_generator + .get_result(request) + .await + .map_err(|e| e.into()) + } - #[tracing::instrument(skip_all)] - async fn get_epoch_result( - &self, - request: Request, - ) -> Result, Status> { - METRICS.increment_request_counter(OP_GET_EPOCH_RESULT); - Ok(self.epoch_manager.get_epoch_result(request).await?) - } + // TODO(#2868) refactor to use MetricedError + #[tracing::instrument(skip(self, request))] + async fn new_mpc_context( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(OP_NEW_MPC_CONTEXT); + self.context_manager + .new_mpc_context(request) + .await + .map_err(|e| e.into()) + } - async fn destroy_mpc_epoch( - &self, - request: Request, - ) -> Result, Status> { - METRICS.increment_request_counter(OP_DESTROY_EPOCH); - Ok(self.epoch_manager.destroy_mpc_epoch(request).await?) - } + #[tracing::instrument(skip(self, request))] + async fn destroy_mpc_context( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(OP_DESTROY_MPC_CONTEXT); + let inner = request.into_inner(); + + // Validate the whole request before mutating anything: a request with valid epoch_ids + // but a missing/malformed context_id must not erase epochs and only then fail context + // parsing. + // + // Note: we extract the inner message and later re-wrap it in `Request::new(inner)`. + let proto_context_id = inner + .context_id + .as_ref() + .ok_or_else(|| Status::invalid_argument("context_id is required"))?; + let context_id = + parse_grpc_request_id::(proto_context_id, RequestIdParsingErr::Context)?; + + // Hold the exclusive context lease across both phases. This makes the epoch snapshot + // stable with respect to `NewMpcEpoch`, including creations that are still running PRSS + // and therefore have not registered their epoch in the session maker yet. + let _destruction_lease = self + .session_maker + .try_start_context_destruction(&context_id) + .await + .map_err(|e| { + Status::failed_precondition(format!( + "Cannot destroy MPC context {context_id}: {e}. Retry once epoch creation has settled." + )) + })?; + + // Destroy the associated epochs first: their secret key shares and PRSS randomness are security-sensitive + // material. Erase them before touching anything else so that if there is a problem, the worst transient + // state is "shares gone, context metadata lingers" rather than "context gone, shares still on disk". + // + // If any epoch fails to delete we return here and leave the context intact. The caller is expected to retry + // the whole `DestroyMpcContext` until both epochs and context are destroyed successfully. + let epochs_destroyed = self + .epoch_manager + .destroy_epochs_for_context(&context_id) + .await + .map_err(Status::from)?; + + // Every epoch is now gone, so it is safe to remove the context. + self.context_manager + .destroy_mpc_context(Request::new(inner)) + .await + .map_err(Status::from)?; + Ok(Response::new(DestroyMpcContextResponse { + epoch_ids: epochs_destroyed.into_iter().map(|id| id.into()).collect(), + })) + } - #[tracing::instrument(skip(self, request))] - async fn new_custodian_context( - &self, - request: Request, - ) -> Result, Status> { - METRICS.increment_request_counter(OP_NEW_CUSTODIAN_CONTEXT); - self.context_manager.new_custodian_context(request).await.map_err(|e| e.into()) - } + #[tracing::instrument(skip_all)] + async fn new_mpc_epoch( + &self, + request: Request, + ) -> Result, Status> { + let inner = request.into_inner(); + METRICS.increment_request_counter(OP_NEW_EPOCH); + Ok(self + .epoch_manager + .new_mpc_epoch(Request::new(inner.clone())) + .await?) + } - #[tracing::instrument(skip(self, request))] - async fn destroy_custodian_context( - &self, - request: Request, - ) -> Result, Status> { - METRICS.increment_request_counter(OP_DESTROY_CUSTODIAN_CONTEXT); - self.context_manager.destroy_custodian_context(request).await.map_err(|e| e.into()) - } + #[tracing::instrument(skip_all)] + async fn get_epoch_result( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(OP_GET_EPOCH_RESULT); + Ok(self.epoch_manager.get_epoch_result(request).await?) + } - #[tracing::instrument(skip(self, request))] - async fn get_operator_public_key( - &self, - request: Request, - ) -> Result, Status> { - METRICS.increment_request_counter(OP_FETCH_PK); - self.backup_operator.get_operator_public_key(request).await.map_err(|e| e.into()) - } + async fn destroy_mpc_epoch( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(OP_DESTROY_EPOCH); + Ok(self.epoch_manager.destroy_mpc_epoch(request).await?) + } - #[tracing::instrument(skip(self, request))] - async fn custodian_backup_recovery( - &self, - request: Request, - ) -> Result, Status> { - METRICS.increment_request_counter(OP_CUSTODIAN_BACKUP_RECOVERY); - self.backup_operator.custodian_backup_recovery(request).await.map_err(|e| e.into()) - } + #[tracing::instrument(skip(self, request))] + async fn new_custodian_context( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(OP_NEW_CUSTODIAN_CONTEXT); + self.context_manager + .new_custodian_context(request) + .await + .map_err(|e| e.into()) + } - #[tracing::instrument(skip(self, request))] - async fn restore_from_backup( - &self, - request: Request, - ) -> Result, Status> { - METRICS.increment_request_counter(OP_RESTORE_FROM_BACKUP); - self.backup_operator.restore_from_backup(request).await.map_err(|e| e.into()) - } + #[tracing::instrument(skip(self, request))] + async fn destroy_custodian_context( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(OP_DESTROY_CUSTODIAN_CONTEXT); + self.context_manager + .destroy_custodian_context(request) + .await + .map_err(|e| e.into()) + } - #[tracing::instrument(skip(self, _request))] - async fn get_key_material_availability( - &self, - _request: Request, - ) -> Result, Status> { - // Get preprocessing IDs from the preprocessor - let preprocessing_ids = self.keygen_preprocessor.get_all_preprocessing_ids().await?; + #[tracing::instrument(skip(self, request))] + async fn get_operator_public_key( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(OP_FETCH_PK); + self.backup_operator + .get_operator_public_key(request) + .await + .map_err(|e| e.into()) + } - // Get storage references from backup_operator - let backup_response = self.backup_operator.get_key_material_availability(Request::new(Empty {})).await?; - let mut response = backup_response.into_inner(); + #[tracing::instrument(skip(self, request))] + async fn custodian_backup_recovery( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(OP_CUSTODIAN_BACKUP_RECOVERY); + self.backup_operator + .custodian_backup_recovery(request) + .await + .map_err(|e| e.into()) + } - // Update the response with preprocessing IDs - response.preprocessing_ids = preprocessing_ids; + #[tracing::instrument(skip(self, request))] + async fn restore_from_backup( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(OP_RESTORE_FROM_BACKUP); + self.backup_operator + .restore_from_backup(request) + .await + .map_err(|e| e.into()) + } - Ok(Response::new(response)) - } + #[tracing::instrument(skip(self, _request))] + async fn get_key_material_availability( + &self, + _request: Request, + ) -> Result, Status> { + // Get preprocessing IDs from the preprocessor + let preprocessing_ids = self.keygen_preprocessor.get_all_preprocessing_ids().await?; + + // Get storage references from backup_operator + let backup_response = self + .backup_operator + .get_key_material_availability(Request::new(Empty {})) + .await?; + let mut response = backup_response.into_inner(); + + // Update the response with preprocessing IDs + response.preprocessing_ids = preprocessing_ids; + + Ok(Response::new(response)) + } - #[tracing::instrument(skip(self, request))] - async fn custodian_recovery_init( - &self, - request: Request, - ) -> Result, Status> { - METRICS.increment_request_counter(OP_CUSTODIAN_RECOVERY_INIT); - self.backup_operator.custodian_recovery_init(request).await.map_err(|e| e.into()) - } + #[tracing::instrument(skip(self, request))] + async fn custodian_recovery_init( + &self, + request: Request, + ) -> Result, Status> { + METRICS.increment_request_counter(OP_CUSTODIAN_RECOVERY_INIT); + self.backup_operator + .custodian_recovery_init(request) + .await + .map_err(|e| e.into()) + } - #[tracing::instrument(skip(self, _request))] - async fn get_health_status( - &self, - _request: Request, - ) -> Result, Status> { - - // Get own key material directly from backup_operator (avoid redundant gRPC call to self) - let backup_response = self.backup_operator.get_key_material_availability(Request::new(Empty {})).await?; - let mut own_material = backup_response.into_inner(); - - // Add preprocessing IDs from the preprocessor - own_material.preprocessing_ids = self.keygen_preprocessor.get_all_preprocessing_ids().await?; - - let health_check_sessions = self.session_maker.get_healthcheck_session_all_contexts().await - .map_err(|e| {tonic::Status::internal(format!("Failed to get health check sessions: {}", e))})?; - - let mut peers_from_all_contexts = Vec::new(); - for (context_id,health_check_session) in health_check_sessions { - let my_role = health_check_session.get_my_role().one_based() as u32; - let total_nodes = health_check_session.get_num_parties() as u32; - let min_nodes_for_healthy = (2 * total_nodes) / 3 + 1; // 2/3 majority + 1 - let min_threshold = (total_nodes / 3) + 1; // Minimum threshold to be able to reconstruct anything - let health_check_results = health_check_session.run_healthcheck().await; - let mut peers_status = Vec::new(); - let mut nodes_reachable = 1; //I am reachable - if let Ok(results) = health_check_results { - for ((role, identity), result) in results.into_iter() { - let peer_status = match result { - HealthCheckStatus::Ok(latency) => { - nodes_reachable += 1; - PeerHealth { - peer_id: role.one_based() as u32, - endpoint: identity.hostname().to_string(), - reachable: true, - latency_ms: latency.as_millis() as u32, - error: String::new(), - } - } - HealthCheckStatus::Error((latency, error)) => PeerHealth { + #[tracing::instrument(skip(self, _request))] + async fn get_health_status( + &self, + _request: Request, + ) -> Result, Status> { + // Get own key material directly from backup_operator (avoid redundant gRPC call to self) + let backup_response = self + .backup_operator + .get_key_material_availability(Request::new(Empty {})) + .await?; + let mut own_material = backup_response.into_inner(); + + // Add preprocessing IDs from the preprocessor + own_material.preprocessing_ids = + self.keygen_preprocessor.get_all_preprocessing_ids().await?; + + let health_check_sessions = self + .session_maker + .get_healthcheck_session_all_contexts() + .await + .map_err(|e| { + tonic::Status::internal(format!("Failed to get health check sessions: {}", e)) + })?; + + let mut peers_from_all_contexts = Vec::new(); + for (context_id, health_check_session) in health_check_sessions { + let my_role = health_check_session.get_my_role().one_based() as u32; + let total_nodes = health_check_session.get_num_parties() as u32; + let min_nodes_for_healthy = (2 * total_nodes) / 3 + 1; // 2/3 majority + 1 + let min_threshold = (total_nodes / 3) + 1; // Minimum threshold to be able to reconstruct anything + let health_check_results = health_check_session.run_healthcheck().await; + let mut peers_status = Vec::new(); + let mut nodes_reachable = 1; //I am reachable + if let Ok(results) = health_check_results { + for ((role, identity), result) in results.into_iter() { + let peer_status = match result { + HealthCheckStatus::Ok(latency) => { + nodes_reachable += 1; + PeerHealth { peer_id: role.one_based() as u32, endpoint: identity.hostname().to_string(), - reachable: false, + reachable: true, latency_ms: latency.as_millis() as u32, - error: format!("Error : {}", error.message()), - }, - HealthCheckStatus::TimeOut(elapsed) => {PeerHealth { - peer_id: role.one_based() as u32, - endpoint: identity.hostname().to_string(), - reachable: false, - latency_ms: 0, - error: format!("Timeout after {:?} s", elapsed.as_secs()), - }}, - }; - peers_status.push(peer_status); - } - } else { - tracing::warn!("Health check failed for context {:?}", context_id); + error: String::new(), + } + } + HealthCheckStatus::Error((latency, error)) => PeerHealth { + peer_id: role.one_based() as u32, + endpoint: identity.hostname().to_string(), + reachable: false, + latency_ms: latency.as_millis() as u32, + error: format!("Error : {}", error.message()), + }, + HealthCheckStatus::TimeOut(elapsed) => PeerHealth { + peer_id: role.one_based() as u32, + endpoint: identity.hostname().to_string(), + reachable: false, + latency_ms: 0, + error: format!("Timeout after {:?} s", elapsed.as_secs()), + }, + }; + peers_status.push(peer_status); } - - // Determine overall health status - let status = if nodes_reachable >= total_nodes { - HealthStatus::Optimal.into() // HEALTH_STATUS_OPTIMAL - all nodes online and reachable - } else if nodes_reachable >= min_nodes_for_healthy { - HealthStatus::Healthy.into() // HEALTH_STATUS_HEALTHY - sufficient 2/3 majority but not all nodes - } else if nodes_reachable > min_threshold { - HealthStatus::Degraded.into() // HEALTH_STATUS_DEGRADED - above minimum threshold but below 2/3 - } else { - HealthStatus::Unhealthy.into() // HEALTH_STATUS_UNHEALTHY - insufficient nodes for operations - }; - - let peers_from_context = PeersFromContext { - context_id: Some(context_id.into()), - my_party_id: my_role, - threshold_required: min_threshold, - nodes_reachable, - status, - peers: peers_status, - }; - peers_from_all_contexts.push(peers_from_context) + } else { + tracing::warn!("Health check failed for context {:?}", context_id); } - - let response = HealthStatusResponse { - peers_from_all_contexts, - my_fhe_key_ids: own_material.fhe_key_ids, - my_crs_ids: own_material.crs_ids, - my_preprocessing_key_ids: own_material.preprocessing_ids, - my_storage_info: own_material.storage_info, - node_type: NodeType::Threshold.into(), + // Determine overall health status + let status = if nodes_reachable >= total_nodes { + HealthStatus::Optimal.into() // HEALTH_STATUS_OPTIMAL - all nodes online and reachable + } else if nodes_reachable >= min_nodes_for_healthy { + HealthStatus::Healthy.into() // HEALTH_STATUS_HEALTHY - sufficient 2/3 majority but not all nodes + } else if nodes_reachable > min_threshold { + HealthStatus::Degraded.into() // HEALTH_STATUS_DEGRADED - above minimum threshold but below 2/3 + } else { + HealthStatus::Unhealthy.into() // HEALTH_STATUS_UNHEALTHY - insufficient nodes for operations }; - Ok(Response::new(response)) + let peers_from_context = PeersFromContext { + context_id: Some(context_id.into()), + my_party_id: my_role, + threshold_required: min_threshold, + nodes_reachable, + status, + peers: peers_status, + }; + peers_from_all_contexts.push(peers_from_context) } + let response = HealthStatusResponse { + peers_from_all_contexts, + my_fhe_key_ids: own_material.fhe_key_ids, + my_crs_ids: own_material.crs_ids, + my_preprocessing_key_ids: own_material.preprocessing_ids, + my_storage_info: own_material.storage_info, + node_type: NodeType::Threshold.into(), + }; - #[tracing::instrument(skip(self, request))] - async fn bandwidth_benchmark( - &self, - request: Request, - ) -> Result, Status> { - run_bandwidth_benchmark( - request, - self.session_maker.clone(), - std::sync::Arc::clone(&self.bandwidth_bench_limiter), - &self.bandwidth_bench_config, - ).await - } + Ok(Response::new(response)) + } + + #[tracing::instrument(skip(self, request))] + async fn bandwidth_benchmark( + &self, + request: Request, + ) -> Result, Status> { + run_bandwidth_benchmark( + request, + self.session_maker.clone(), + std::sync::Arc::clone(&self.bandwidth_bench_limiter), + &self.bandwidth_bench_config, + ) + .await } } diff --git a/core/service/src/engine/threshold/mod.rs b/core/service/src/engine/threshold/mod.rs index 2b7df489a0..e2c02eaa86 100644 --- a/core/service/src/engine/threshold/mod.rs +++ b/core/service/src/engine/threshold/mod.rs @@ -2,4 +2,3 @@ pub mod bandwidth_bench; mod endpoint; pub mod service; pub mod threshold_kms; -pub mod traits; diff --git a/core/service/src/engine/threshold/service/crs_generator.rs b/core/service/src/engine/threshold/service/crs_generator.rs index 503df5c771..9150520f89 100644 --- a/core/service/src/engine/threshold/service/crs_generator.rs +++ b/core/service/src/engine/threshold/service/crs_generator.rs @@ -37,10 +37,7 @@ use crate::{ BaseKmsStruct, CrsGenMetadata, DSEP_PUBDATA_CRS, compute_info_crs, stored_scheme_signatures_to_proto, }, - threshold::{ - service::session::{ImmutableSessionMaker, validate_context_and_epoch}, - traits::CrsGenerator, - }, + threshold::service::session::{ImmutableSessionMaker, validate_context_and_epoch}, validation::{RequestIdParsingErr, parse_grpc_request_id, validate_crs_gen_request}, }, util::{ @@ -56,7 +53,7 @@ use crate::{ // === Insecure Feature-Specific Imports === cfg_if::cfg_if! { if #[cfg(feature = "insecure")] { - use crate::engine::{centralized::central_kms::async_generate_crs, threshold::traits::InsecureCrsGenerator}; + use crate::engine::centralized::central_kms::async_generate_crs; use threshold_execution::{tfhe_internals::test_feature::transfer_crs}; } } @@ -68,7 +65,7 @@ cfg_if::cfg_if! { } } -pub struct RealCrsGenerator< +pub(crate) struct RealCrsGenerator< PubS: Storage + Send + Sync + 'static, PrivS: StorageExt + Send + Sync + 'static, C: Ceremony + Send + Sync + 'static, @@ -478,28 +475,27 @@ impl< } } -#[tonic::async_trait] impl< PubS: Storage + Send + Sync + 'static, PrivS: StorageExt + Send + Sync + 'static, C: Ceremony + Send + Sync + 'static, -> CrsGenerator for RealCrsGenerator +> RealCrsGenerator { - async fn crs_gen( + pub(crate) async fn crs_gen( &self, request: Request, ) -> Result, MetricedError> { self.inner_crs_gen_from_request(request, false).await } - async fn get_result( + pub(crate) async fn get_result( &self, request: Request, ) -> Result, MetricedError> { self.inner_get_result(request, false).await } - async fn abort_crs_gen( + pub(crate) async fn abort_crs_gen( &self, request: Request, ) -> Result, MetricedError> { @@ -508,7 +504,7 @@ impl< } #[cfg(feature = "insecure")] -pub struct RealInsecureCrsGenerator< +pub(crate) struct RealInsecureCrsGenerator< PubS: Storage + Send + Sync + 'static, PrivS: StorageExt + Send + Sync + 'static, C: Ceremony + Send + Sync + 'static, @@ -540,14 +536,13 @@ impl< } #[cfg(feature = "insecure")] -#[tonic::async_trait] impl< PubS: Storage + Send + Sync + 'static, PrivS: StorageExt + Send + Sync + 'static, C: Ceremony + Send + Sync + 'static, -> InsecureCrsGenerator for RealInsecureCrsGenerator +> RealInsecureCrsGenerator { - async fn insecure_crs_gen( + pub(crate) async fn insecure_crs_gen( &self, request: Request, ) -> Result, MetricedError> { @@ -557,7 +552,7 @@ impl< .await } - async fn get_result( + pub(crate) async fn get_result( &self, request: Request, ) -> Result, MetricedError> { @@ -565,13 +560,6 @@ impl< .inner_get_result(request, true) .await } - - async fn abort_crs_gen( - &self, - request: Request, - ) -> Result, MetricedError> { - self.real_crs_generator.inner_abort_crs_gen(request).await - } } #[cfg(test)] diff --git a/core/service/src/engine/threshold/service/epoch_manager.rs b/core/service/src/engine/threshold/service/epoch_manager.rs index f58dbcbd68..76ae485bdb 100644 --- a/core/service/src/engine/threshold/service/epoch_manager.rs +++ b/core/service/src/engine/threshold/service/epoch_manager.rs @@ -1,6 +1,6 @@ //! This file provides methods to manage epochs, which includes dealing with PRSS setups //! as well as resharing of the secret keys. -//! The main struct is [`RealThresholdEpochManager`] which implements the [`EpochManager`] trait. +//! The main struct is [`RealThresholdEpochManager`]. //! //! __PRSS__ //! @@ -88,7 +88,6 @@ use crate::{ }, session::{ImmutableSessionMaker, PRSSSetupCombined, SessionMaker}, }, - traits::EpochManager, utils::{MetricedError, signing_identity_for}, validation::{ RequestIdParsingErr, ResharingParams, VerifiedNewMpcEpochRequest, @@ -423,7 +422,7 @@ pub enum EpochOutput { } /// The Epoch Manager takes over the role of the Initiator and Resharer -pub struct RealThresholdEpochManager< +pub(crate) struct RealThresholdEpochManager< PubS: Storage + Send + Sync + 'static, PrivS: StorageExt + Send + Sync + 'static, Init: PRSSInit + PRSSInit, @@ -1844,7 +1843,6 @@ impl< } } -#[tonic::async_trait] impl< PubS: Storage + Send + Sync + 'static, PrivS: StorageExt + Send + Sync + 'static, @@ -1853,9 +1851,9 @@ impl< + Default + 'static, Reshare: ReshareSecretKeys + Default + 'static, -> EpochManager for RealThresholdEpochManager +> RealThresholdEpochManager { - async fn new_mpc_epoch( + pub(crate) async fn new_mpc_epoch( &self, request: Request, ) -> Result, MetricedError> { @@ -2011,7 +2009,7 @@ impl< Ok(Response::new(Empty {})) } - async fn destroy_mpc_epoch( + pub(crate) async fn destroy_mpc_epoch( &self, request: Request, ) -> Result, MetricedError> { @@ -2024,7 +2022,9 @@ impl< self.destroy_epoch_with_lease(&epoch_id).await } - async fn destroy_epochs_for_context( + /// Destroy every epoch associated with `context_id`. + /// This is needed for context destruction since the Context manager does not compose with the epoch manager. + pub(crate) async fn destroy_epochs_for_context( &self, context_id: &ContextId, ) -> Result, MetricedError> { @@ -2032,7 +2032,13 @@ impl< self.destroy_mpc_epochs(&epochs_to_destroy).await } - async fn destroy_mpc_epochs( + /// Destroy every epoch in `epoch_ids` as part of a context destruction. + /// + /// Idempotent and best-effort: an epoch that is not present (already destroyed, or never + /// created on this party) is skipped; every listed epoch is attempted even if an earlier one + /// fails; the first deletion error is returned only after all attempts, so a failed run can be + /// retried until no shares remain. + pub(crate) async fn destroy_mpc_epochs( &self, epoch_ids: &[EpochId], ) -> Result, MetricedError> { @@ -2064,7 +2070,7 @@ impl< } } - async fn get_epoch_result( + pub(crate) async fn get_epoch_result( &self, request: Request, ) -> Result, MetricedError> { diff --git a/core/service/src/engine/threshold/service/key_generator.rs b/core/service/src/engine/threshold/service/key_generator.rs index d583ac8c22..8914fb7755 100644 --- a/core/service/src/engine/threshold/service/key_generator.rs +++ b/core/service/src/engine/threshold/service/key_generator.rs @@ -58,12 +58,9 @@ use crate::{ }, keyset_configuration::InternalKeySetConfig, material_integrity::verify_public_key_digest_from_bytes, - threshold::{ - service::{ - PublicKeyMaterial, ThresholdFheKeys, - session::{ImmutableSessionMaker, validate_context_and_epoch}, - }, - traits::KeyGenerator, + threshold::service::{ + PublicKeyMaterial, ThresholdFheKeys, + session::{ImmutableSessionMaker, validate_context_and_epoch}, }, utils::{MetricedError, signing_identity_for}, validation::{ @@ -115,8 +112,6 @@ enum ThresholdKeyGenResult { // === Insecure Feature-Specific Imports === #[cfg(feature = "insecure")] -use crate::engine::threshold::traits::InsecureKeyGenerator; -#[cfg(feature = "insecure")] use threshold_execution::runtime::sessions::session_parameters::GenericParameterHandles; #[cfg(feature = "insecure")] use threshold_execution::tfhe_internals::{ @@ -129,7 +124,7 @@ use threshold_execution::tfhe_internals::{ }, }; -pub struct RealKeyGenerator< +pub(crate) struct RealKeyGenerator< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, KG: OnlineDistributedKeyGen + 'static, @@ -155,7 +150,7 @@ pub struct RealKeyGenerator< } #[cfg(feature = "insecure")] -pub struct RealInsecureKeyGenerator< +pub(crate) struct RealInsecureKeyGenerator< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, KG: OnlineDistributedKeyGen + 'static, @@ -188,14 +183,13 @@ impl< } } #[cfg(feature = "insecure")] -#[tonic::async_trait] impl< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, KG: OnlineDistributedKeyGen + 'static, -> InsecureKeyGenerator for RealInsecureKeyGenerator +> RealInsecureKeyGenerator { - async fn insecure_key_gen( + pub(crate) async fn insecure_key_gen( &self, request: Request, ) -> Result, MetricedError> { @@ -203,7 +197,7 @@ impl< self.real_key_generator.inner_key_gen(request, true).await } - async fn get_result( + pub(crate) async fn get_result( &self, request: Request, ) -> Result, MetricedError> { @@ -211,20 +205,10 @@ impl< .inner_get_result(request, true) .await } - - async fn abort_key_gen(&self, preproc_id: RequestId) -> Status { - self.real_key_generator - .inner_abort_key_gen(preproc_id) - .await - } } -// This is an enum to determine whether to start the dkg -// in a secure mode. If the secure mode is selected, -// a preprocessing handle must be given. -// This is essentially the same as an Option, but it's -// more clear to label the variants as `Secure` -// and `Insecure`. +// This is an enum to determine whether to start the dkg in a secure mode. If the secure mode is selected, a +// preprocessing handle must be given. #[expect(clippy::type_complexity)] #[derive(Clone)] pub enum PreprocHandleWithMode { @@ -234,6 +218,7 @@ pub enum PreprocHandleWithMode { Arc>>>, ), ), + #[cfg_attr(not(feature = "insecure"), expect(dead_code))] Insecure(RequestId), } @@ -1855,28 +1840,29 @@ impl< } } -#[tonic::async_trait] impl< PubS: Storage + Sync + Send + 'static, PrivS: StorageExt + Sync + Send + 'static, KG: OnlineDistributedKeyGen + 'static, -> KeyGenerator for RealKeyGenerator +> RealKeyGenerator { - async fn key_gen( + pub(crate) async fn key_gen( &self, request: Request, ) -> Result, MetricedError> { self.inner_key_gen(request, false).await } - async fn get_result( + pub(crate) async fn get_result( &self, request: tonic::Request, ) -> Result, MetricedError> { self.inner_get_result(request, false).await } - async fn abort_key_gen(&self, preproc_id: RequestId) -> Status { + // Note that the Status is returned, since this call is never directly mapped to the gRPC end-point, + // but instead used in conjunction with `abort_key_gen_preproc` in [`RealPreprocessor`] + pub(crate) async fn abort_key_gen(&self, preproc_id: RequestId) -> Status { self.inner_abort_key_gen(preproc_id).await } } diff --git a/core/service/src/engine/threshold/service/kms_impl.rs b/core/service/src/engine/threshold/service/kms_impl.rs index 3b78efa52f..8a7b898880 100644 --- a/core/service/src/engine/threshold/service/kms_impl.rs +++ b/core/service/src/engine/threshold/service/kms_impl.rs @@ -32,16 +32,9 @@ use tfhe::{ xof_key_set::CompressedXofKeySet, }; use tfhe_versionable::{Upgrade, Version, VersionsDispatch}; -use threshold_execution::endpoints::reshare_sk::SecureReshareSecretKeys; use threshold_execution::{ - endpoints::keygen::SecureOnlineDistributedKeyGen128, - online::preprocessing::{ - DKGPreprocessing, create_memory_factory, - orchestration::producer_traits::SecureSmallProducerFactory, - }, - small_execution::prss::RobustSecurePrssInit, + online::preprocessing::{DKGPreprocessing, create_memory_factory}, tfhe_internals::{parameters::DKGParams, private_keysets::PrivateKeySet}, - zk::ceremony::SecureCeremony, }; use threshold_networking::{ grpc::{GrpcNetworkingManager, GrpcServer, TlsExtensionGetter}, @@ -79,11 +72,7 @@ use crate::{ verify_storage_material, }, threshold::{ - service::{ - public_decryptor::SecureNoiseFloodDecryptor, - session::{ImmutableSessionMaker, SessionMaker}, - user_decryptor::SecureNoiseFloodPartialDecryptor, - }, + service::session::{ImmutableSessionMaker, SessionMaker}, threshold_kms::ThresholdKms, }, traits::PrivateKeyMaterialMetadata, @@ -470,44 +459,6 @@ pub(crate) fn new_insecure_preproc_bucket( }) } -#[cfg(not(feature = "insecure"))] -pub type RealThresholdKms = ThresholdKms< - RealThresholdEpochManager, - RealUserDecryptor, - RealPublicDecryptor, - RealKeyGenerator< - PubS, - PrivS, - SecureOnlineDistributedKeyGen128<{ ResiduePolyF4Z128::EXTENSION_DEGREE }>, - >, - RealPreprocessor>, - RealCrsGenerator, - ThresholdContextManager, - RealBackupOperator, ->; - -#[cfg(feature = "insecure")] -pub type RealThresholdKms = ThresholdKms< - RealThresholdEpochManager, - RealUserDecryptor, - RealPublicDecryptor, - RealKeyGenerator< - PubS, - PrivS, - SecureOnlineDistributedKeyGen128<{ ResiduePolyF4Z128::EXTENSION_DEGREE }>, - >, - RealInsecureKeyGenerator< - PubS, - PrivS, - SecureOnlineDistributedKeyGen128<{ ResiduePolyF4Z128::EXTENSION_DEGREE }>, - >, - RealPreprocessor>, - RealCrsGenerator, - RealInsecureCrsGenerator, // doesn't matter which ceremony we use here - ThresholdContextManager, - RealBackupOperator, ->; - #[expect(clippy::too_many_arguments)] pub async fn new_real_threshold_kms( config: CoreConfig, @@ -520,7 +471,7 @@ pub async fn new_real_threshold_kms( tls_config: Option<(ServerConfig, ClientConfig, Arc)>, shutdown_signal: F, ) -> anyhow::Result<( - RealThresholdKms, + ThresholdKms, (HealthState, HealthServer), MetaStoreStatusServiceImpl, )> @@ -689,7 +640,7 @@ where // We are only serving after initialization health .reporter() - .set_not_serving::>>() + .set_not_serving::>>() .await; let mpc_server_health = health.clone(); diff --git a/core/service/src/engine/threshold/service/mod.rs b/core/service/src/engine/threshold/service/mod.rs index 2ee457798f..5e207dbb6b 100644 --- a/core/service/src/engine/threshold/service/mod.rs +++ b/core/service/src/engine/threshold/service/mod.rs @@ -15,16 +15,16 @@ mod kms_impl; pub use kms_impl::*; // Module components -mod crs_generator; +pub(crate) mod crs_generator; pub(crate) mod epoch_manager; // Re-exported so the epoch data type is reachable outside the crate (e.g. the // backward-compatibility generators and tests) without exposing the rest of the // epoch manager internals. pub use epoch_manager::EpochData; //mod initiator; -mod key_generator; -mod preprocessor; -mod public_decryptor; +pub(crate) mod key_generator; +pub(crate) mod preprocessor; +pub(crate) mod public_decryptor; pub(crate) mod reshare_utils; pub mod session; -mod user_decryptor; +pub(crate) mod user_decryptor; diff --git a/core/service/src/engine/threshold/service/preprocessor.rs b/core/service/src/engine/threshold/service/preprocessor.rs index 2f9228cc6c..46b3769633 100644 --- a/core/service/src/engine/threshold/service/preprocessor.rs +++ b/core/service/src/engine/threshold/service/preprocessor.rs @@ -44,10 +44,7 @@ use crate::{ base::{ BaseKmsStruct, compute_preprocessing_signatures, stored_scheme_signatures_to_proto, }, - threshold::{ - service::session::{ImmutableSessionMaker, validate_context_and_epoch}, - traits::KeyGenPreprocessor, - }, + threshold::service::session::{ImmutableSessionMaker, validate_context_and_epoch}, utils::{MetricedError, signing_identity_for}, validation::{RequestIdParsingErr, parse_grpc_request_id, validate_preproc_request}, }, @@ -63,8 +60,9 @@ use crate::{ // === Current Module Imports === use super::{BucketMetaStore, PreprocMaterial}; -pub struct RealPreprocessor>> -{ +pub(crate) struct RealPreprocessor< + P: ProducerFactory>, +> { // TODO eventually add mode to allow for nlarge as well. pub(crate) base_kms: BaseKmsStruct, pub preproc_buckets: Arc>>, @@ -610,11 +608,10 @@ impl>> Rea } } -#[tonic::async_trait] impl> + Send + Sync> - KeyGenPreprocessor for RealPreprocessor

+ RealPreprocessor

{ - async fn key_gen_preproc( + pub(crate) async fn key_gen_preproc( &self, request: Request, ) -> Result, MetricedError> { @@ -627,7 +624,7 @@ impl> + Se } #[cfg(feature = "insecure")] - async fn partial_key_gen_preproc( + pub(crate) async fn partial_key_gen_preproc( &self, request: Request, ) -> Result, MetricedError> { @@ -644,8 +641,10 @@ impl> + Se .await } + /// Insecure (dummy) preprocessing that records metadata but no preprocessing + /// material in the meta store, to be consumed by the insecure key generation. #[cfg(feature = "insecure")] - async fn insecure_key_gen_preproc( + pub(crate) async fn insecure_key_gen_preproc( &self, request: Request, ) -> Result, MetricedError> { @@ -653,7 +652,7 @@ impl> + Se .await } - async fn get_result( + pub(crate) async fn get_result( &self, request: Request, ) -> Result, MetricedError> { @@ -661,8 +660,10 @@ impl> + Se .await } + /// Same as [`Self::get_result`] but for preprocessing started via + /// [`Self::insecure_key_gen_preproc`]. #[cfg(feature = "insecure")] - async fn get_insecure_result( + pub(crate) async fn get_insecure_result( &self, request: Request, ) -> Result, MetricedError> { @@ -670,7 +671,7 @@ impl> + Se .await } - async fn abort_key_gen_preproc( + pub(crate) async fn abort_key_gen_preproc( &self, preproc_id: RequestId, key_gen_cancel_res: Status, @@ -700,7 +701,7 @@ impl> + Se } } - async fn get_all_preprocessing_ids(&self) -> Result, MetricedError> { + pub(crate) async fn get_all_preprocessing_ids(&self) -> Result, MetricedError> { let guarded_meta_store = self.preproc_buckets.read().await; Ok(guarded_meta_store .get_successful_completed_request_ids() diff --git a/core/service/src/engine/threshold/service/public_decryptor.rs b/core/service/src/engine/threshold/service/public_decryptor.rs index 74cee93446..f426a5026c 100644 --- a/core/service/src/engine/threshold/service/public_decryptor.rs +++ b/core/service/src/engine/threshold/service/public_decryptor.rs @@ -52,10 +52,7 @@ use crate::{ BaseKmsStruct, PubDecCallValues, deserialize_to_low_level, sign_public_decryption_result, }, - threshold::{ - service::session::{ImmutableSessionMaker, validate_context_and_epoch}, - traits::PublicDecryptor, - }, + threshold::service::session::{ImmutableSessionMaker, validate_context_and_epoch}, utils::{MetricedError, format_handle, format_unvalidated_id, signing_identity_for}, validation::{ RequestIdParsingErr, parse_grpc_request_id, parse_optional_grpc_request_id, @@ -119,7 +116,7 @@ impl NoiseFloodDecryptor for SecureNoiseFloodDecryptor { } } -pub struct RealPublicDecryptor< +pub(crate) struct RealPublicDecryptor< PubS: Storage + Send + Sync + 'static, PrivS: StorageExt + Send + Sync + 'static, Dec: NoiseFloodDecryptor< @@ -239,7 +236,6 @@ impl< } } -#[tonic::async_trait] impl< PubS: Storage + Send + Sync + 'static, PrivS: StorageExt + Send + Sync + 'static, @@ -249,7 +245,7 @@ impl< SmallSession, >, > + 'static, -> PublicDecryptor for RealPublicDecryptor +> RealPublicDecryptor { // `context_id`/`epoch_id` are only known after request validation, so they start empty and are // recorded below. Every event and error in this request — including the ones emitted by the @@ -263,7 +259,7 @@ impl< context_id = tracing::field::Empty, epoch_id = tracing::field::Empty ))] - async fn public_decrypt( + pub(crate) async fn public_decrypt( &self, request: Request, ) -> Result, MetricedError> { @@ -651,7 +647,7 @@ impl< Ok(Response::new(Empty {})) } - async fn public_decrypt_sync( + pub(crate) async fn public_decrypt_sync( &self, request: Request, ) -> Result, MetricedError> { @@ -675,7 +671,7 @@ impl< self.get_result(Request::new(req_id.into())).await } - async fn get_result( + pub(crate) async fn get_result( &self, request: Request, ) -> Result, MetricedError> { diff --git a/core/service/src/engine/threshold/service/user_decryptor.rs b/core/service/src/engine/threshold/service/user_decryptor.rs index 2754438077..8459a0d9ef 100644 --- a/core/service/src/engine/threshold/service/user_decryptor.rs +++ b/core/service/src/engine/threshold/service/user_decryptor.rs @@ -66,10 +66,7 @@ use crate::{ BaseKmsStruct, UserDecryptCallValues, deserialize_to_low_level, sign_user_decryption_result, }, - threshold::{ - service::session::{ImmutableSessionMaker, validate_context_and_epoch}, - traits::UserDecryptor, - }, + threshold::service::session::{ImmutableSessionMaker, validate_context_and_epoch}, utils::{MetricedError, format_handle, format_unvalidated_id, signing_identity_for}, validation::{ DSEP_USER_DECRYPTION, RequestIdParsingErr, parse_grpc_request_id, @@ -133,7 +130,7 @@ impl NoiseFloodPartialDecryptor for SecureNoiseFloodPartialDecryptor { } } -pub struct RealUserDecryptor< +pub(crate) struct RealUserDecryptor< PubS: Storage + Send + Sync + 'static, PrivS: StorageExt + Send + Sync + 'static, Dec: NoiseFloodPartialDecryptor< @@ -445,7 +442,6 @@ impl< } } -#[tonic::async_trait] impl< PubS: Storage + Send + Sync + 'static, PrivS: StorageExt + Send + Sync + 'static, @@ -455,7 +451,7 @@ impl< SmallSession, >, > + 'static, -> UserDecryptor for RealUserDecryptor +> RealUserDecryptor { // Mirrors the public decryption span: `context_id`/`epoch_id` are only known after request // validation, so they start empty and are recorded below. The spawned decryption task inherits @@ -466,7 +462,7 @@ impl< context_id = tracing::field::Empty, epoch_id = tracing::field::Empty ))] - async fn user_decrypt( + pub(crate) async fn user_decrypt( &self, request: Request, ) -> Result, MetricedError> { @@ -621,7 +617,7 @@ impl< Ok(Response::new(Empty {})) } - async fn user_decrypt_sync( + pub(crate) async fn user_decrypt_sync( &self, request: Request, ) -> Result, MetricedError> { @@ -645,7 +641,7 @@ impl< self.get_result(Request::new(req_id.into())).await } - async fn get_result( + pub(crate) async fn get_result( &self, request: Request, ) -> Result, MetricedError> { diff --git a/core/service/src/engine/threshold/threshold_kms.rs b/core/service/src/engine/threshold/threshold_kms.rs index 72cccf4b31..b6ef0bf7b5 100644 --- a/core/service/src/engine/threshold/threshold_kms.rs +++ b/core/service/src/engine/threshold/threshold_kms.rs @@ -1,39 +1,59 @@ use crate::conf::BandwidthBenchmarkConfig; use crate::engine::Shutdown; +use crate::engine::backup_operator::RealBackupOperator; +use crate::engine::context_manager::ThresholdContextManager; use crate::engine::threshold::bandwidth_bench::new_bandwidth_bench_limiter; +use crate::engine::threshold::service::crs_generator::RealCrsGenerator; +#[cfg(feature = "insecure")] +use crate::engine::threshold::service::crs_generator::RealInsecureCrsGenerator; +use crate::engine::threshold::service::epoch_manager::RealThresholdEpochManager; +#[cfg(feature = "insecure")] +use crate::engine::threshold::service::key_generator::RealInsecureKeyGenerator; +use crate::engine::threshold::service::key_generator::RealKeyGenerator; +use crate::engine::threshold::service::preprocessor::RealPreprocessor; +use crate::engine::threshold::service::public_decryptor::{ + RealPublicDecryptor, SecureNoiseFloodDecryptor, +}; use crate::engine::threshold::service::session::ImmutableSessionMaker; +use crate::engine::threshold::service::user_decryptor::{ + RealUserDecryptor, SecureNoiseFloodPartialDecryptor, +}; use crate::retry_loop; +use crate::vault::storage::{Storage, StorageExt}; +use algebra::galois_rings::degree_4::ResiduePolyF4Z128; +use algebra::structure_traits::Ring; use kms_grpc::kms_service::v1::core_service_endpoint_server::CoreServiceEndpointServer; use observability::health::HealthState; use std::sync::Arc; +use threshold_execution::endpoints::keygen::SecureOnlineDistributedKeyGen128; +use threshold_execution::endpoints::reshare_sk::SecureReshareSecretKeys; +use threshold_execution::online::preprocessing::orchestration::producer_traits::SecureSmallProducerFactory; +use threshold_execution::small_execution::prss::RobustSecurePrssInit; +use threshold_execution::zk::ceremony::SecureCeremony; use tokio::sync::Semaphore; use tokio::task::JoinHandle; use tokio_util::task::TaskTracker; +type KeyGen = SecureOnlineDistributedKeyGen128<{ ResiduePolyF4Z128::EXTENSION_DEGREE }>; + +/// The threshold KMS gRPC service: one party of the MPC network. pub struct ThresholdKms< - EP: Sync, - UD: Sync, - PD: Sync, - KG: Sync, - #[cfg(feature = "insecure")] IKG: Sync, - PP: Sync, - CG: Sync, - #[cfg(feature = "insecure")] ICG: Sync, - CM: Sync, - BO: Sync, + PubS: Storage + Send + Sync + 'static, + PrivS: StorageExt + Send + Sync + 'static, > { - pub(crate) epoch_manager: EP, - pub(crate) user_decryptor: UD, - pub(crate) decryptor: PD, - pub(crate) key_generator: KG, + pub(crate) epoch_manager: + RealThresholdEpochManager, + pub(crate) user_decryptor: RealUserDecryptor, + pub(crate) decryptor: RealPublicDecryptor, + pub(crate) key_generator: RealKeyGenerator, #[cfg(feature = "insecure")] - pub(crate) insecure_key_generator: IKG, - pub(crate) keygen_preprocessor: PP, - pub(crate) crs_generator: CG, + pub(crate) insecure_key_generator: RealInsecureKeyGenerator, + pub(crate) keygen_preprocessor: RealPreprocessor>, + pub(crate) crs_generator: RealCrsGenerator, #[cfg(feature = "insecure")] - pub(crate) insecure_crs_generator: ICG, - pub(crate) context_manager: CM, - pub(crate) backup_operator: BO, + pub(crate) insecure_crs_generator: RealInsecureCrsGenerator, // doesn't matter which ceremony we use here + pub(crate) context_manager: ThresholdContextManager, + pub(crate) backup_operator: RealBackupOperator, pub(crate) session_maker: ImmutableSessionMaker, /// Bounds concurrent `bandwidth_benchmark` runs (the endpoint takes no rate-limiter permit). pub(crate) bandwidth_bench_limiter: Arc, @@ -44,32 +64,36 @@ pub struct ThresholdKms< mpc_abort_handle: JoinHandle>, } -#[cfg(feature = "insecure")] -impl< - EP: Sync, - UD: Sync, - PD: Sync, - KG: Sync, - IKG: Sync, - PP: Sync, - CG: Sync, - ICG: Sync, - CM: Sync, - BO: Sync, -> ThresholdKms +impl ThresholdKms +where + PubS: Storage + Send + Sync + 'static, + PrivS: StorageExt + Send + Sync + 'static, { #[expect(clippy::too_many_arguments)] pub(crate) fn new( - epoch_manager: EP, - user_decryptor: UD, - decryptor: PD, - key_generator: KG, - insecure_key_generator: IKG, - keygen_preprocessor: PP, - crs_generator: CG, - insecure_crs_generator: ICG, - context_manager: CM, - backup_operator: BO, + epoch_manager: RealThresholdEpochManager< + PubS, + PrivS, + RobustSecurePrssInit, + SecureReshareSecretKeys, + >, + user_decryptor: RealUserDecryptor, + decryptor: RealPublicDecryptor, + key_generator: RealKeyGenerator, + #[cfg(feature = "insecure")] insecure_key_generator: RealInsecureKeyGenerator< + PubS, + PrivS, + KeyGen, + >, + keygen_preprocessor: RealPreprocessor>, + crs_generator: RealCrsGenerator, + #[cfg(feature = "insecure")] insecure_crs_generator: RealInsecureCrsGenerator< + PubS, + PrivS, + SecureCeremony, + >, + context_manager: ThresholdContextManager, + backup_operator: RealBackupOperator, tracker: Arc, session_maker: ImmutableSessionMaker, bandwidth_bench_config: BandwidthBenchmarkConfig, @@ -81,9 +105,11 @@ impl< user_decryptor, decryptor, key_generator, + #[cfg(feature = "insecure")] insecure_key_generator, keygen_preprocessor, crs_generator, + #[cfg(feature = "insecure")] insecure_crs_generator, context_manager, backup_operator, @@ -99,20 +125,11 @@ impl< } } -#[cfg(feature = "insecure")] #[tonic::async_trait] -impl< - EP: Sync, - UD: Sync, - PD: Sync, - KG: Sync, - IKG: Sync, - PP: Sync, - CG: Sync, - ICG: Sync, - CM: Sync, - BO: Sync, -> Shutdown for ThresholdKms +impl Shutdown for ThresholdKms +where + PubS: Storage + Send + Sync + 'static, + PrivS: StorageExt + Send + Sync + 'static, { fn shutdown(&self) -> anyhow::Result> { let health = self.health.clone(); @@ -153,116 +170,13 @@ impl< } } -#[cfg(feature = "insecure")] -impl< - EP: Sync, - UD: Sync, - PD: Sync, - KG: Sync, - IKG: Sync, - PP: Sync, - CG: Sync, - ICG: Sync, - CM: Sync, - BO: Sync, -> Drop for ThresholdKms +impl Drop for ThresholdKms +where + PubS: Storage + Send + Sync + 'static, + PrivS: StorageExt + Send + Sync + 'static, { fn drop(&mut self) { // Start the shutdown and let it finish in the background let _ = self.shutdown(); } } - -#[cfg(not(feature = "insecure"))] -impl - ThresholdKms -{ - #[expect(clippy::too_many_arguments)] - pub(crate) fn new( - epoch_manager: EP, - user_decryptor: UD, - decryptor: PD, - key_generator: KG, - keygen_preprocessor: PP, - crs_generator: CG, - context_manager: CM, - backup_operator: BO, - tracker: Arc, - session_maker: ImmutableSessionMaker, - bandwidth_bench_config: BandwidthBenchmarkConfig, - health: HealthState, - mpc_abort_handle: JoinHandle>, - ) -> Self { - Self { - epoch_manager, - user_decryptor, - decryptor, - key_generator, - keygen_preprocessor, - crs_generator, - context_manager, - backup_operator, - tracker, - session_maker, - bandwidth_bench_limiter: new_bandwidth_bench_limiter( - bandwidth_bench_config.max_concurrent_runs, - ), - bandwidth_bench_config, - health, - mpc_abort_handle, - } - } -} - -#[tonic::async_trait] -#[cfg(not(feature = "insecure"))] -impl Shutdown - for ThresholdKms -{ - fn shutdown(&self) -> anyhow::Result> { - let health = self.health.clone(); - let tracker = Arc::clone(&self.tracker); - let mpc_abort_handle = self.mpc_abort_handle.abort_handle(); - let handle = { - let new_handle_clone = mpc_abort_handle.clone(); - tokio::task::spawn(async move { - health.mark_shutting_down().await; - health - .reporter() - .set_not_serving::>() - .await; - tracing::trace!("Set not serving"); - tracker.close(); - tracker.wait().await; - mpc_abort_handle.abort(); - let res: anyhow::Result<()> = retry_loop!( - || { - let new_handle_clone = new_handle_clone.clone(); - async move { - if !new_handle_clone.is_finished() { - return Err(anyhow::anyhow!("MPC server not done")); - } - Ok(()) - } - }, - 100, - 200 - ); - if let Err(e) = res { - tracing::error!("Error waiting for MPC server to finish: {:?}", e); - } - tracing::info!("Threshold Core service endpoint server shutdown complete."); - }) - }; - Ok(handle) - } -} - -#[cfg(not(feature = "insecure"))] -impl Drop - for ThresholdKms -{ - fn drop(&mut self) { - let _ = self.shutdown(); - } -} diff --git a/core/service/src/engine/threshold/traits.rs b/core/service/src/engine/threshold/traits.rs deleted file mode 100644 index cef015fbb7..0000000000 --- a/core/service/src/engine/threshold/traits.rs +++ /dev/null @@ -1,142 +0,0 @@ -use crate::engine::utils::MetricedError; -use kms_grpc::kms::v1::*; -use tonic::{Request, Response, Status}; - -#[tonic::async_trait] -pub trait UserDecryptor { - async fn user_decrypt( - &self, - request: Request, - ) -> Result, MetricedError>; - async fn get_result( - &self, - request: Request, - ) -> Result, MetricedError>; - async fn user_decrypt_sync( - &self, - request: Request, - ) -> Result, MetricedError>; -} - -#[tonic::async_trait] -pub trait PublicDecryptor { - async fn public_decrypt( - &self, - request: Request, - ) -> Result, MetricedError>; - async fn get_result( - &self, - request: Request, - ) -> Result, MetricedError>; - async fn public_decrypt_sync( - &self, - request: Request, - ) -> Result, MetricedError>; -} - -#[tonic::async_trait] -pub trait KeyGenerator { - async fn key_gen( - &self, - request: Request, - ) -> Result, MetricedError>; - async fn get_result( - &self, - request: Request, - ) -> Result, MetricedError>; - // Note that the Status is returned, since this call is never directly mapped to the gRPC end-point, - // but instead used in conjunction with `abort_key_gen_preproc` in [`KeyGenPreprocessor`] - async fn abort_key_gen(&self, preproc_id: kms_grpc::RequestId) -> Status; -} - -#[cfg(feature = "insecure")] -#[tonic::async_trait] -pub trait InsecureKeyGenerator { - async fn insecure_key_gen( - &self, - request: Request, - ) -> Result, MetricedError>; - async fn get_result( - &self, - request: Request, - ) -> Result, MetricedError>; - // Note that the Status is returned, since this call is never directly mapped to the gRPC end-point, - // but instead used in conjunction with `abort_key_gen_preproc` in [`KeyGenPreprocessor`] - async fn abort_key_gen(&self, preproc_id: kms_grpc::RequestId) -> Status; -} - -#[tonic::async_trait] -pub trait KeyGenPreprocessor { - async fn key_gen_preproc( - &self, - request: Request, - ) -> Result, MetricedError>; - - #[cfg(feature = "insecure")] - async fn partial_key_gen_preproc( - &self, - request: Request, - ) -> Result, MetricedError>; - - /// Insecure (dummy) preprocessing that records metadata but no preprocessing - /// material in the meta store, to be consumed by the insecure key generation. - #[cfg(feature = "insecure")] - async fn insecure_key_gen_preproc( - &self, - request: Request, - ) -> Result, MetricedError>; - - async fn get_result( - &self, - request: Request, - ) -> Result, MetricedError>; - - /// Same as [`Self::get_result`] but for preprocessing started via - /// [`Self::insecure_key_gen_preproc`]. - #[cfg(feature = "insecure")] - async fn get_insecure_result( - &self, - request: Request, - ) -> Result, MetricedError>; - - async fn get_all_preprocessing_ids(&self) -> Result, MetricedError>; - - async fn abort_key_gen_preproc( - &self, - preproc_id: kms_grpc::RequestId, - key_gen_cancel_res: Status, - ) -> Result, MetricedError>; -} - -#[tonic::async_trait] -pub trait CrsGenerator { - async fn crs_gen( - &self, - request: Request, - ) -> Result, MetricedError>; - async fn get_result( - &self, - request: Request, - ) -> Result, MetricedError>; - async fn abort_crs_gen( - &self, - request: Request, - ) -> Result, MetricedError>; -} - -#[cfg(feature = "insecure")] -#[tonic::async_trait] -pub trait InsecureCrsGenerator { - async fn insecure_crs_gen( - &self, - request: Request, - ) -> Result, MetricedError>; - async fn get_result( - &self, - request: Request, - ) -> Result, MetricedError>; - async fn abort_crs_gen( - &self, - request: Request, - ) -> Result, MetricedError>; -} diff --git a/core/service/src/engine/traits.rs b/core/service/src/engine/traits.rs index 21952d6ad1..40bcf4a7d4 100644 --- a/core/service/src/engine/traits.rs +++ b/core/service/src/engine/traits.rs @@ -1,53 +1,11 @@ -use hashing::DomainSep; +use kms_grpc::ContextId; use kms_grpc::kms::v1::*; -use kms_grpc::{ContextId, EpochId}; -use rand::CryptoRng; -use rand::RngCore; -use serde::Serialize; -use tfhe::FheTypes; use tonic::Request; use tonic::Response; -use crate::cryptography::encryption::UnifiedPublicEncKey; -use crate::cryptography::signatures::{PrivateSigKey, Signature}; use crate::engine::base::KeyGenMetadata; use crate::engine::utils::MetricedError; -use super::base::KmsFheKeyHandles; - -pub trait BaseKms { - fn sign>( - &self, - dsep: &DomainSep, - msg: &T, - ) -> anyhow::Result; - fn digest>( - domain_separator: &DomainSep, - msg: &T, - ) -> anyhow::Result>; -} -/// The [Kms] trait represents either a dummy KMS, an HSM, or an MPC network. -pub trait Kms: BaseKms { - fn public_decrypt( - keys: &KmsFheKeyHandles, - ct: &[u8], - fhe_type: FheTypes, - ct_format: CiphertextFormat, - ) -> anyhow::Result; - #[expect(clippy::too_many_arguments)] - fn user_decrypt( - keys: &KmsFheKeyHandles, - sig_key: &PrivateSigKey, - rng: &mut (impl CryptoRng + RngCore), - ct: &[u8], - ct_type: FheTypes, - ct_format: CiphertextFormat, - digest_link: &[u8], - enc_key: &UnifiedPublicEncKey, - client_address: &[u8], - ) -> anyhow::Result>; -} - #[tonic::async_trait] pub trait ContextManager { async fn new_mpc_context( @@ -78,70 +36,6 @@ pub trait ContextManager { async fn mpc_context_exists_in_cache(&self, context_id: &ContextId) -> bool; } -#[tonic::async_trait] -pub trait EpochManager { - async fn new_mpc_epoch( - &self, - request: Request, - ) -> Result, MetricedError>; - - async fn destroy_mpc_epoch( - &self, - request: Request, - ) -> Result, MetricedError>; - - /// Destroy every epoch in `epoch_ids` as part of a context destruction. - /// - /// Idempotent and best-effort: an epoch that is not present (already destroyed, or never - /// created on this party) is skipped; every listed epoch is attempted even if an earlier one - /// fails; the first deletion error is returned only after all attempts, so a failed run can be - /// retried until no shares remain. - async fn destroy_mpc_epochs( - &self, - epoch_ids: &[EpochId], - ) -> Result, MetricedError>; - - /// Destroy every epoch associated with `context_id`. - /// This is needed for context destruction since the Context manager does not compose with the epoch manager. - async fn destroy_epochs_for_context( - &self, - context_id: &ContextId, - ) -> Result, MetricedError>; - - async fn get_epoch_result( - &self, - request: Request, - ) -> Result, MetricedError>; -} - -#[tonic::async_trait] -pub trait BackupOperator { - async fn get_operator_public_key( - &self, - request: Request, - ) -> Result, MetricedError>; - - async fn custodian_recovery_init( - &self, - request: Request, - ) -> Result, MetricedError>; - - async fn custodian_backup_recovery( - &self, - request: Request, - ) -> Result, MetricedError>; - - async fn restore_from_backup( - &self, - request: Request, - ) -> Result, MetricedError>; - - async fn get_key_material_availability( - &self, - request: Request, - ) -> Result, MetricedError>; -} - pub trait PrivateKeyMaterialMetadata { fn get_metadata(&self) -> &KeyGenMetadata; }