Merge branch 'main' into tore/fix/network-concurrency

Author: jot2re

Cargo.lock

@@ -72,7 +72,6 @@ reqwest = { version = "=0.12.22", default-features = false, features = [
     "rustls-tls",
 ] }
 schemars = "=0.8.22"
-secp256k1 = "=0.31.1"
 serde = { version = "=1.0.219", features = ["derive", "rc"] }
 serde_json = "=1.0.140"
 serial_test = "=3.2.0"
@@ -100,6 +99,7 @@ tracing-opentelemetry = "=0.30.0"
 tracing-subscriber = { version = "=0.3.19", features = ["fmt", "std"] }
 trait-variant = "0.1.2"
 validator = { version = "=0.20.0", features = ["derive"] }
+zeroize = { version = "=1.8.1", features = ["zeroize_derive"] }
 x509-parser = { version = "=0.17.0", features = ["verify"] }
 
 alloy-dyn-abi = "=1.1.2"

Cargo.toml

@@ -4,8 +4,8 @@ image_id = "ami-0a56cf46caf2fd41c"
 instance_type = "m6i.4xlarge"
 
 [backend.aws.big-instance-service]
-region = "eu-west-3"
-image_id = "ami-0a56cf46caf2fd41c"
+region = "eu-west-1"
+image_id = "ami-04f62ae93462b459a"
 instance_type = "c7i.8xlarge"
 
 [backend.aws.docker-big-instance]

ci/slab.toml

@@ -6,6 +6,7 @@ import "kms.v1.proto";
 service CoreServiceEndpoint {
   // Perform the threshold KMS initialization.
   // This call returns an error on the centralized KMS.
+  // This is a synchronous call and an successful response suggests the initialization completed.
   rpc Init(kms.v1.InitRequest) returns (kms.v1.Empty);
 
   // Start generating preprocessing materials for key generation asynchronously.

core/grpc/proto/kms-service.v1.proto

@@ -109,6 +109,7 @@ tracing = { workspace = true, features = ["log"] }
 trait-variant.workspace = true
 url = { version = "=2.5.4", features = ["serde"] }
 validator.workspace = true
+zeroize.workspace = true
 ordermap = "=0.5.7"
 x509-parser = { workspace = true, optional = true }
 

core/service/Cargo.toml

@@ -6,6 +6,7 @@ use serde::{Deserialize, Serialize};
 use tfhe::named::Named;
 use tfhe::Versionize;
 use tfhe_versionable::VersionsDispatch;
+use zeroize::Zeroize;
 
 use crate::consts::SAFE_SER_SIZE_LIMIT;
 use crate::cryptography::hybrid_ml_kem::{self};
@@ -21,12 +22,40 @@ struct InnerBackupPrivateKey {
     decapsulation_key: <ml_kem::kem::Kem<MlKemParams> as ml_kem::KemCore>::DecapsulationKey,
 }
 
+impl Drop for InnerBackupPrivateKey {
+    fn drop(&mut self) {
+        // Directly zeroize the underlying key bytes without creating copies
+        // This is more secure as it avoids temporary allocations of sensitive data
+        let key_bytes_ptr = self.decapsulation_key.as_bytes().as_ptr() as *mut u8;
+        let key_len = self.decapsulation_key.as_bytes().len();
+
+        // SAFETY: We're zeroizing the memory that belongs to this struct
+        // The pointer is valid and the length is correct from as_bytes()
+        unsafe {
+            std::ptr::write_bytes(key_bytes_ptr, 0, key_len);
+        }
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize, VersionsDispatch)]
 pub enum BackupPrivateKeyVersioned {
     V0(BackupPrivateKey),
 }
 
-#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Versionize)]
+// NOTE:
+// The `Versionize` derive macro cannot be combined with a custom `Drop` implementation
+// or `ZeroizeOnDrop` derivation without conflicts.
+//
+// Security-wise, the current design is solid:
+// - Actual cryptographic key material resides in `InnerBackupPrivateKey`, not in the
+//   serialized `Vec<u8>` representation.
+// - Most `BackupPrivateKey` instances are short-lived and immediately used for decryption;
+//   `decrypt()` converts them into `InnerBackupPrivateKey`.
+// - `InnerBackupPrivateKey` implements a custom `Drop` with explicit memory zeroization.
+// - `BackupPrivateKey` implements `Zeroize` for manual wiping if needed.
+// - `try_from()` ensures that the intermediate `decaps_key_buf` (which *does* contain
+//   sensitive data) is securely zeroized after use.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Versionize, Zeroize)]
 #[versionize(BackupPrivateKeyVersioned)]
 pub struct BackupPrivateKey {
     decapsulation_key: Vec<u8>,
@@ -59,6 +88,10 @@ impl TryFrom<&BackupPrivateKey> for InnerBackupPrivateKey {
 
         let decapsulation_key =
             <MlKemType as ml_kem::KemCore>::DecapsulationKey::from_bytes(&decaps_key_buf);
+
+        // Zeroize the key buffer to prevent memory disclosure attacks
+        decaps_key_buf.zeroize();
+
         Ok(Self { decapsulation_key })
     }
 }

core/service/src/cryptography/backup_pke.rs

@@ -15,10 +15,11 @@ use observability::{
 use threshold_fhe::{
     algebra::base_ring::Z64,
     execution::{
-        runtime::session::{BaseSession, ParameterHandles, ToBaseSession},
+        runtime::session::{BaseSession, ParameterHandles},
         tfhe_internals::parameters::DKGParams,
         zk::ceremony::{compute_witness_dim, Ceremony},
     },
+    networking::NetworkMode,
 };
 use tokio::sync::{Mutex, OwnedSemaphorePermit, RwLock};
 use tokio_util::{sync::CancellationToken, task::TaskTracker};
@@ -182,11 +183,11 @@ impl<
         }
 
         let session_id = req_id.derive_session_id()?;
+        // CRS ceremony requires a sync network
         let session = self
             .session_preparer
-            .prepare_ddec_data_from_sessionid_z128(session_id)
-            .await?
-            .to_base_session();
+            .make_base_session(session_id, NetworkMode::Sync)
+            .await?;
 
         let meta_store = Arc::clone(&self.crs_meta_store);
         let meta_store_cancelled = Arc::clone(&self.crs_meta_store);
@@ -664,6 +665,19 @@ mod tests {
         };
 
         let request = Request::new(req);
+        let _ = crs_gen.crs_gen(request).await.unwrap();
+
+        // Send a the request again, it should return an error
+        // because the session has already been used.
+        let req = CrsGenRequest {
+            params: 0,
+            max_num_bits: None,
+            request_id: Some(req_id.into()),
+            domain: None,
+        };
+
+        let request = Request::new(req);
+
         assert_eq!(
             crs_gen.crs_gen(request).await.unwrap_err().code(),
             tonic::Code::Aborted

core/service/src/engine/threshold/service/crs_generator.rs

@@ -1,5 +1,5 @@
 // === Standard Library ===
-use std::sync::Arc;
+use std::{marker::PhantomData, sync::Arc};
 
 // === External Crates ===
 use kms_grpc::{
@@ -12,7 +12,7 @@ use threshold_fhe::{
     algebra::galois_rings::degree_4::{ResiduePolyF4Z128, ResiduePolyF4Z64},
     execution::{
         runtime::session::ParameterHandles,
-        small_execution::prss::{PRSSInit, PRSSSetup, RobustSecurePrssInit},
+        small_execution::prss::{PRSSInit, PRSSSetup},
     },
     networking::NetworkMode,
 };
@@ -33,16 +33,26 @@ use crate::{
 // === Current Module Imports ===
 use super::{session::SessionPreparer, RealThresholdKms};
 
-pub struct RealInitiator<PrivS: Storage + Send + Sync + 'static> {
+pub struct RealInitiator<
+    PrivS: Storage + Send + Sync + 'static,
+    Init: PRSSInit<ResiduePolyF4Z64> + PRSSInit<ResiduePolyF4Z128>,
+> {
     // TODO eventually add mode to allow for nlarge as well.
     pub prss_setup_z128: Arc<RwLock<Option<PRSSSetup<ResiduePolyF4Z128>>>>,
     pub prss_setup_z64: Arc<RwLock<Option<PRSSSetup<ResiduePolyF4Z64>>>>,
     pub private_storage: Arc<Mutex<PrivS>>,
     pub session_preparer: SessionPreparer,
-    pub health_reporter: Arc<RwLock<HealthReporter>>,
+    pub health_reporter: HealthReporter,
+    pub(crate) _init: PhantomData<Init>,
 }
 
-impl<PrivS: Storage + Send + Sync + 'static> RealInitiator<PrivS> {
+impl<
+        PrivS: Storage + Send + Sync + 'static,
+        Init: PRSSInit<ResiduePolyF4Z64, OutputType = PRSSSetup<ResiduePolyF4Z64>>
+            + PRSSInit<ResiduePolyF4Z128, OutputType = PRSSSetup<ResiduePolyF4Z128>>
+            + Default,
+    > RealInitiator<PrivS, Init>
+{
     pub async fn init_prss_from_disk(&self, req_id: &RequestId) -> anyhow::Result<()> {
         let prss_setup_z128_from_file = {
             let guarded_private_storage = self.private_storage.lock().await;
@@ -109,8 +119,6 @@ impl<PrivS: Storage + Send + Sync + 'static> RealInitiator<PrivS> {
         {
             // Notice that this is a hack to get the health reporter to report serving. The type `PrivS` has no influence on the service name.
             self.health_reporter
-                .write()
-                .await
                 .set_serving::<CoreServiceEndpointServer<RealThresholdKms<PrivS, PrivS>>>()
                 .await;
         }
@@ -141,13 +149,14 @@ impl<PrivS: Storage + Send + Sync + 'static> RealInitiator<PrivS> {
             "Role assignments: {:?}",
             base_session.parameters.role_assignments()
         );
-        let prss_setup_obj_z128: PRSSSetup<ResiduePolyF4Z128> = RobustSecurePrssInit::default()
-            .init(&mut base_session)
-            .await?;
 
-        let prss_setup_obj_z64: PRSSSetup<ResiduePolyF4Z64> = RobustSecurePrssInit::default()
-            .init(&mut base_session)
-            .await?;
+        // It seems we cannot do something like
+        // `Init::default().init(&mut base_session).await?;`
+        // as the type inference gets confused even when using the correct return type.
+        let prss_setup_obj_z128: PRSSSetup<ResiduePolyF4Z128> =
+            PRSSInit::<ResiduePolyF4Z128>::init(&Init::default(), &mut base_session).await?;
+        let prss_setup_obj_z64: PRSSSetup<ResiduePolyF4Z64> =
+            PRSSInit::<ResiduePolyF4Z64>::init(&Init::default(), &mut base_session).await?;
 
         let mut guarded_prss_setup = self.prss_setup_z128.write().await;
         *guarded_prss_setup = Some(prss_setup_obj_z128.clone());
@@ -186,8 +195,6 @@ impl<PrivS: Storage + Send + Sync + 'static> RealInitiator<PrivS> {
         {
             // Notice that this is a hack to get the health reporter to report serving. The type `PrivS` has no influence on the service name.
             self.health_reporter
-                .write()
-                .await
                 .set_serving::<CoreServiceEndpointServer<RealThresholdKms<PrivS, PrivS>>>()
                 .await;
         }
@@ -197,7 +204,13 @@ impl<PrivS: Storage + Send + Sync + 'static> RealInitiator<PrivS> {
 }
 
 #[tonic::async_trait]
-impl<PrivS: Storage + Send + Sync + 'static> Initiator for RealInitiator<PrivS> {
+impl<
+        PrivS: Storage + Send + Sync + 'static,
+        Init: PRSSInit<ResiduePolyF4Z64, OutputType = PRSSSetup<ResiduePolyF4Z64>>
+            + PRSSInit<ResiduePolyF4Z128, OutputType = PRSSSetup<ResiduePolyF4Z128>>
+            + Default,
+    > Initiator for RealInitiator<PrivS, Init>
+{
     async fn init(&self, request: Request<v1::InitRequest>) -> Result<Response<Empty>, Status> {
         let inner = request.into_inner();
         let request_id = tonic_some_or_err(
@@ -219,14 +232,38 @@ impl<PrivS: Storage + Send + Sync + 'static> Initiator for RealInitiator<PrivS>
 
 #[cfg(test)]
 mod tests {
+    use super::*;
+
     use crate::{
         client::test_tools::{self},
         consts::PRSS_INIT_REQ_ID,
-        engine::base::derive_request_id,
         util::key_setup::test_tools::purge,
-        vault::storage::{file::FileStorage, StorageType},
+        vault::storage::{file::FileStorage, ram, StorageType},
+    };
+    use aes_prng::AesRng;
+    use kms_grpc::kms::v1::InitRequest;
+    use rand::SeedableRng;
+    use threshold_fhe::{
+        execution::runtime::party::Role,
+        malicious_execution::small_execution::malicious_prss::{EmptyPrss, FailingPrss},
     };
-    use threshold_fhe::execution::runtime::party::Role;
+
+    impl<
+            Init: PRSSInit<ResiduePolyF4Z64, OutputType = PRSSSetup<ResiduePolyF4Z64>>
+                + PRSSInit<ResiduePolyF4Z128, OutputType = PRSSSetup<ResiduePolyF4Z128>>,
+        > RealInitiator<ram::RamStorage, Init>
+    {
+        fn init_test(session_preparer: SessionPreparer) -> Self {
+            Self {
+                prss_setup_z128: Arc::new(RwLock::new(None)),
+                prss_setup_z64: Arc::new(RwLock::new(None)),
+                private_storage: Arc::new(Mutex::new(ram::RamStorage::new())),
+                session_preparer,
+                health_reporter: HealthReporter::new(),
+                _init: PhantomData,
+            }
+        }
+    }
 
     #[tokio::test]
     #[serial_test::serial]
@@ -316,4 +353,75 @@ mod tests {
             "Initializing threshold KMS server with PRSS Setup Z64 from disk"
         ));
     }
+
+    #[tokio::test]
+    async fn sunshine() {
+        let session_preparer = SessionPreparer::new_test_session(false);
+        let initiator = RealInitiator::<ram::RamStorage, EmptyPrss>::init_test(session_preparer);
+
+        let mut rng = AesRng::seed_from_u64(42);
+        let req_id = RequestId::new_random(&mut rng);
+        initiator
+            .init(tonic::Request::new(InitRequest {
+                request_id: Some(req_id.into()),
+            }))
+            .await
+            .unwrap();
+    }
+    #[tokio::test]
+    async fn invalid_argument() {
+        let session_preparer = SessionPreparer::new_test_session(false);
+        let initiator = RealInitiator::<ram::RamStorage, EmptyPrss>::init_test(session_preparer);
+
+        let bad_req_id = kms_grpc::kms::v1::RequestId {
+            request_id: "bad req id".to_string(),
+        };
+        assert_eq!(
+            initiator
+                .init(tonic::Request::new(InitRequest {
+                    request_id: Some(bad_req_id)
+                }))
+                .await
+                .unwrap_err()
+                .code(),
+            tonic::Code::InvalidArgument
+        );
+    }
+
+    #[tokio::test]
+    async fn aborted() {
+        let session_preparer = SessionPreparer::new_test_session(false);
+        let initiator = RealInitiator::<ram::RamStorage, EmptyPrss>::init_test(session_preparer);
+
+        assert_eq!(
+            initiator
+                .init(tonic::Request::new(InitRequest {
+                    // this is set to none
+                    request_id: None
+                }))
+                .await
+                .unwrap_err()
+                .code(),
+            tonic::Code::Aborted
+        );
+    }
+
+    #[tokio::test]
+    async fn internal() {
+        let session_preparer = SessionPreparer::new_test_session(false);
+        let initiator = RealInitiator::<ram::RamStorage, FailingPrss>::init_test(session_preparer);
+
+        let mut rng = AesRng::seed_from_u64(42);
+        let req_id = RequestId::new_random(&mut rng);
+        assert_eq!(
+            initiator
+                .init(tonic::Request::new(InitRequest {
+                    request_id: Some(req_id.into())
+                }))
+                .await
+                .unwrap_err()
+                .code(),
+            tonic::Code::Internal
+        );
+    }
 }