chore: rework keygen config

Author: kc1212

core-client/src/keygen.rs

@@ -47,8 +47,8 @@ pub(crate) async fn do_keygen(
         cc_conf.num_majority
     };
 
-    //NOTE: If we do not use dummy_domain here, then
-    //this needs changing too in the KeyGenResult command.
+    // NOTE: If we do not use dummy_domain here, then
+    // this needs changing too in the KeyGenResult command.
     let keyset_config = if shared_config.compressed || shared_config.keyset_type.is_some() {
         Some(kms_grpc::kms::v1::KeySetConfig {
             keyset_type: shared_config
@@ -58,7 +58,7 @@ pub(crate) async fn do_keygen(
                 .unwrap_or(kms_grpc::kms::v1::KeySetType::Standard as i32),
             standard_keyset_config: Some(kms_grpc::kms::v1::StandardKeySetConfig {
                 compute_key_type: 0,  // CPU
-                secret_key_config: 0, // Generate
+                secret_key_config: 0, // Generate all secret keys
                 compressed_key_config: if shared_config.compressed {
                     kms_grpc::kms::v1::CompressedKeyConfig::CompressedAll.into()
                 } else {

core/grpc/proto/kms.v1.proto

@@ -46,17 +46,24 @@ enum ComputeKeyType {
 
 message KeySetAddedInfo {
   // Must be set if KeyGenSecretKeyConfig::UseExistingCompressionSecretKey is used
-  RequestId compression_keyset_id = 1;
+  RequestId existing_compression_keyset_id = 1;
+
+  // Can be set if KeyGenSecretKeyConfig::UseExistingCompressionSecretKey is used,
+  // if it's not set then the epoch ID from the KeyGenRequest is used.
+  RequestId compression_epoch_id = 2;
 
   // Must be set if KeySetType::DecompressionOnly is used
-  RequestId from_keyset_id_decompression_only = 2;
+  RequestId from_keyset_id_decompression_only = 3;
 
   // Must be set if KeySetType::DecompressionOnly is used
-  RequestId to_keyset_id_decompression_only = 3;
+  RequestId to_keyset_id_decompression_only = 4;
 
   // Must be set if KeyGenSecretKeyConfig::UseExisting is used
-  RequestId existing_keyset_id = 4;
-  RequestId existing_epoch_id = 5;
+  RequestId existing_keyset_id = 5;
+
+  // Can be set if KeyGenSecretKeyConfig::UseExisting is used,
+  // if it's not set then the epoch ID from the KeyGenRequest is used.
+  RequestId existing_epoch_id = 6;
 }
 
 // The keyset configuration message.
@@ -111,9 +118,16 @@ message PartialKeyGenPreprocRequest {
 }
 
 enum KeyGenSecretKeyConfig {
-  // The default must be represented by 0.
+  // The default must be represented by 0, which is to generate all the secret keys.
   GenerateAll = 0;
+
+  // Use existing secret shares to generate the keyset, the key ID and epoch ID
+  // that points to the shares must be given by KeySetAddedInfo.
   UseExisting = 1;
+
+  // Use the secret share from an existing compression secret key,
+  // but generate the other shares. The share of the existing compresssion secret key
+  // must be given by KeySetAddedInfo.
   UseExistingCompressionSecretKey = 2;
 }
 
@@ -178,7 +192,8 @@ message KeyGenRequest {
   // See https://github.com/zama-ai/kms-internal/issues/2530
   RequestId context_id = 7;
 
-  // The epoch number placeholder (zama-ai/kms-internal#2743).
+  // The epoch ID for the key generation, i.e., the shares produced
+  // under this keygen will be stored and registered under this epoch.
   RequestId epoch_id = 8;
 }
 

core/service/src/client/tests/common.rs

@@ -59,7 +59,8 @@ pub(crate) fn compressed_from_existing_keygen_config(
             }),
         }),
         Some(KeySetAddedInfo {
-            compression_keyset_id: None,
+            existing_compression_keyset_id: None,
+            compression_epoch_id: None,
             from_keyset_id_decompression_only: None,
             to_keyset_id_decompression_only: None,
             existing_keyset_id: Some((*existing_keyset_id).into()),
@@ -80,7 +81,8 @@ pub(crate) fn decompression_keygen_config(
             standard_keyset_config: None,
         }),
         Some(KeySetAddedInfo {
-            compression_keyset_id: None,
+            existing_compression_keyset_id: None,
+            compression_epoch_id: None,
             from_keyset_id_decompression_only: Some((*from_keyset_id).into()),
             to_keyset_id_decompression_only: Some((*to_keyset_id).into()),
             existing_keyset_id: None,

core/service/src/client/tests/threshold/key_gen_tests_isolated.rs

@@ -391,15 +391,10 @@ async fn secure_threshold_keygen_crash_preprocessing_isolated() -> Result<()> {
 /// gRPC service layer.
 ///
 /// **Workflow:**
-/// 1. Standard keygen (preprocessing + online) to produce secret shares
-/// 2. Preprocessing for compressed keygen from existing
+/// 1. Standard keygen (preprocessing + online) to produce the first keyset
+/// 2. Preprocessing for compressed keygen from existing shares
 /// 3. Compressed keygen from existing shares
-/// 4. Verify both keygens completed on all parties
-///
-/// **Requires:**
-/// - `slow_tests` feature flag
-///
-/// **Run with:** `cargo test --lib --features slow_tests,testing secure_threshold_compressed_keygen_from_existing_isolated`
+/// 4. Verify both keygens completed on all parties using ddec
 #[tokio::test]
 #[cfg(feature = "slow_tests")]
 async fn secure_threshold_compressed_keygen_from_existing_isolated() -> Result<()> {
@@ -559,13 +554,8 @@ async fn secure_threshold_compressed_keygen_from_existing_isolated() -> Result<(
 /// 3. Generate decompression key from keyset 1 to keyset 2 (secure mode with preprocessing)
 /// 4. Verify all keys generated successfully
 /// 5. Perform public decryption to validate keys are functional
-///
-/// **Requires:**
-/// - `slow_tests` and `insecure` feature flags (PRSS generation at runtime)
-///
-/// **Run with:** `cargo test --lib --features slow_tests,testing,insecure test_insecure_threshold_decompression_keygen_isolated`
 #[tokio::test]
-#[cfg(all(feature = "slow_tests", feature = "insecure"))]
+#[cfg(feature = "slow_tests")]
 async fn test_insecure_threshold_decompression_keygen_isolated() -> Result<()> {
     use crate::client::tests::threshold::public_decryption_tests::run_decryption_threshold;
     use crate::consts::PUBLIC_STORAGE_PREFIX_THRESHOLD_ALL;
@@ -650,7 +640,8 @@ async fn test_insecure_threshold_decompression_keygen_isolated() -> Result<()> {
                 standard_keyset_config: None,
             }),
             keyset_added_info: Some(KeySetAddedInfo {
-                compression_keyset_id: None,
+                existing_compression_keyset_id: None,
+                compression_epoch_id: None,
                 from_keyset_id_decompression_only: Some(key_id_1.into()),
                 to_keyset_id_decompression_only: Some(key_id_2.into()),
                 existing_keyset_id: None,

core/service/src/engine/centralized/central_kms.rs

@@ -334,7 +334,7 @@ pub fn generate_fhe_keys(
                 }
             }
             KeyGenSecretKeyConfig::UseExisting => {
-                todo!()
+                anyhow::bail!("This is not implemented yet")
             }
         };
 

core/service/src/engine/centralized/service/key_gen.rs

@@ -294,7 +294,7 @@ pub(crate) async fn key_gen_background<
     }
     match internal_keyset_config.keyset_config() {
         KeySetConfig::Standard(standard_key_set_config) => {
-            let compression_id = match internal_keyset_config.get_compression_id() {
+            let compression_id = match internal_keyset_config.get_existing_compression_key_id() {
                 Ok(compression_id) => compression_id,
                 Err(e) => {
                     let _ = update_err_req_in_meta_store(