feat: core/service keygen endpoint

Author: kc1212

core/grpc/proto/kms.v1.proto

@@ -34,6 +34,7 @@ message StandardKeySetConfig {
 
   // Configuration for generating compressed vs uncompressed keys.
   // Compressed keys use XOF-seeded compression and are smaller.
+  // (We may rename this to KeyGenPublicKeyConfig.)
   CompressedKeyConfig compressed_key_config = 3;
 }
 
@@ -44,14 +45,18 @@ enum ComputeKeyType {
 }
 
 message KeySetAddedInfo {
-  // Must be set if KeyGenSecretKeyConfig::UseExisting is used
+  // Must be set if KeyGenSecretKeyConfig::UseExistingCompressionSecretKey is used
   RequestId compression_keyset_id = 1;
 
   // Must be set if KeySetType::DecompressionOnly is used
   RequestId from_keyset_id_decompression_only = 2;
 
   // Must be set if KeySetType::DecompressionOnly is used
   RequestId to_keyset_id_decompression_only = 3;
+
+  // Must be set if KeyGenSecretKeyConfig::UseExisting is used
+  RequestId existing_keyset_id = 4;
+  RequestId existing_epoch_id = 5;
 }
 
 // The keyset configuration message.

core/grpc/src/rpc_types.rs

@@ -1134,7 +1134,7 @@ mod tests {
         assert_eq!(v1::ComputeKeyType::default(), v1::ComputeKeyType::Cpu);
         assert_eq!(
             v1::KeyGenSecretKeyConfig::default(),
-            v1::KeyGenSecretKeyConfig::Generate
+            v1::KeyGenSecretKeyConfig::GenerateAll
         );
     }
 

core/service/src/client/tests/centralized/key_gen_tests.rs

@@ -59,9 +59,9 @@ async fn test_decompression_key_gen_centralized() {
             standard_keyset_config: None,
         }),
         Some(KeySetAddedInfo {
-            compression_keyset_id: None,
             from_keyset_id_decompression_only: Some(request_id_1.into()),
             to_keyset_id_decompression_only: Some(request_id_2.into()),
+            ..Default::default()
         }),
     )
     .await;
@@ -103,9 +103,9 @@ async fn default_decompression_key_gen_centralized() {
             standard_keyset_config: None,
         }),
         Some(KeySetAddedInfo {
-            compression_keyset_id: None,
             from_keyset_id_decompression_only: Some(request_id_1.into()),
             to_keyset_id_decompression_only: Some(request_id_2.into()),
+            ..Default::default()
         }),
     )
     .await;

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

@@ -43,6 +43,31 @@ pub(crate) fn compressed_keygen_config() -> (Option<KeySetConfig>, Option<KeySet
     )
 }
 
+/// Returns compressed keygen config that reuses existing secret key shares
+#[cfg(feature = "slow_tests")]
+pub(crate) fn compressed_from_existing_keygen_config(
+    existing_keyset_id: &RequestId,
+    existing_epoch_id: &kms_grpc::identifiers::EpochId,
+) -> (Option<KeySetConfig>, Option<KeySetAddedInfo>) {
+    (
+        Some(KeySetConfig {
+            keyset_type: KeySetType::Standard.into(),
+            standard_keyset_config: Some(kms_grpc::kms::v1::StandardKeySetConfig {
+                compute_key_type: 0,
+                secret_key_config: kms_grpc::kms::v1::KeyGenSecretKeyConfig::UseExisting.into(),
+                compressed_key_config: CompressedKeyConfig::CompressedAll.into(),
+            }),
+        }),
+        Some(KeySetAddedInfo {
+            compression_keyset_id: None,
+            from_keyset_id_decompression_only: None,
+            to_keyset_id_decompression_only: None,
+            existing_keyset_id: Some((*existing_keyset_id).into()),
+            existing_epoch_id: Some((*existing_epoch_id).into()),
+        }),
+    )
+}
+
 /// Returns decompression-only keygen config
 #[cfg(feature = "slow_tests")]
 pub(crate) fn decompression_keygen_config(
@@ -58,6 +83,8 @@ pub(crate) fn decompression_keygen_config(
             compression_keyset_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,
+            existing_epoch_id: None,
         }),
     )
 }

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

@@ -293,11 +293,16 @@ pub async fn threshold_insecure_key_gen_isolated(
 /// * `Ok(())` if preprocessing and key generation succeeded on all parties
 /// * `Err` if any party failed
 #[cfg(feature = "slow_tests")]
+#[allow(clippy::too_many_arguments)]
 pub async fn threshold_key_gen_secure_isolated(
     clients: &HashMap<u32, CoreServiceEndpointClient<Channel>>,
     preproc_id: &kms_grpc::RequestId,
     keygen_id: &kms_grpc::RequestId,
     params: kms_grpc::kms::v1::FheParameter,
+    keyset_config: Option<kms_grpc::kms::v1::KeySetConfig>,
+    keyset_added_info: Option<kms_grpc::kms::v1::KeySetAddedInfo>,
+    context_id: Option<kms_grpc::kms::v1::RequestId>,
+    epoch_id: Option<kms_grpc::kms::v1::RequestId>,
 ) -> anyhow::Result<()> {
     use crate::dummy_domain;
     use crate::testing::helpers::domain_to_msg;
@@ -314,9 +319,9 @@ pub async fn threshold_key_gen_secure_isolated(
             request_id: Some((*preproc_id).into()),
             params: params as i32,
             domain: Some(domain_msg.clone()),
-            keyset_config: None,
-            context_id: None,
-            epoch_id: None,
+            keyset_config,
+            context_id: context_id.clone(),
+            epoch_id: epoch_id.clone(),
         };
         preproc_tasks.spawn(async move {
             cur_client
@@ -353,10 +358,10 @@ pub async fn threshold_key_gen_secure_isolated(
             params: Some(params as i32),
             preproc_id: Some((*preproc_id).into()),
             domain: Some(domain_msg.clone()),
-            keyset_config: None,
-            keyset_added_info: None,
-            context_id: None,
-            epoch_id: None,
+            keyset_config,
+            keyset_added_info: keyset_added_info.clone(),
+            context_id: context_id.clone(),
+            epoch_id: epoch_id.clone(),
         };
         keygen_tasks
             .spawn(async move { cur_client.key_gen(tonic::Request::new(keygen_req)).await });

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

@@ -389,6 +389,7 @@ async fn decrypt_after_recovery(amount_custodians: usize, threshold: u32) {
         &mut kms_servers,
         &mut kms_clients,
         &mut internal_client,
+        None,
         &req_key_id,
         None,
         vec![TestingPlaintext::U8(u8::MAX)],

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

@@ -971,6 +971,7 @@ pub(crate) async fn preproc_and_keygen(
                 &mut kms_servers,
                 &mut kms_clients,
                 &mut internal_client,
+                None,
                 key_id,
                 None,
                 vec![TestingPlaintext::U8(u8::MAX)],
@@ -1036,6 +1037,7 @@ pub(crate) async fn preproc_and_keygen(
                 &mut kms_servers,
                 &mut kms_clients,
                 &mut internal_client,
+                None,
                 key_id,
                 None,
                 vec![TestingPlaintext::U8(u8::MAX)],

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

@@ -137,8 +137,17 @@ async fn secure_threshold_keygen_isolated() -> Result<()> {
     let keygen_id = derive_request_id("secure_threshold_keygen")?;
 
     // Run secure key generation with preprocessing
-    threshold_key_gen_secure_isolated(&env.clients, &preproc_id, &keygen_id, FheParameter::Test)
-        .await?;
+    threshold_key_gen_secure_isolated(
+        &env.clients,
+        &preproc_id,
+        &keygen_id,
+        FheParameter::Test,
+        None,
+        None,
+        None,
+        None,
+    )
+    .await?;
 
     // Verify key was generated on all parties
     for client in env.all_clients() {
@@ -374,6 +383,170 @@ async fn secure_threshold_keygen_crash_preprocessing_isolated() -> Result<()> {
     Ok(())
 }
 
+/// Test secure threshold compressed key generation from existing secret shares.
+///
+/// Generates a standard keyset first, then performs compressed key generation
+/// reusing the existing secret key shares from the first keygen. This validates
+/// the end-to-end flow of compressed keygen from existing secrets through the
+/// gRPC service layer.
+///
+/// **Workflow:**
+/// 1. Standard keygen (preprocessing + online) to produce secret shares
+/// 2. Preprocessing for compressed keygen from existing
+/// 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`
+#[tokio::test]
+#[cfg(feature = "slow_tests")]
+async fn secure_threshold_compressed_keygen_from_existing_isolated() -> Result<()> {
+    use crate::client::tests::common::compressed_from_existing_keygen_config;
+    use crate::consts::DEFAULT_EPOCH_ID;
+
+    let env = ThresholdTestEnv::builder()
+        .with_test_name("compressed_from_existing_keygen")
+        .with_party_count(4)
+        .with_threshold(1)
+        .with_prss()
+        .build()
+        .await?;
+
+    let clients = &env.clients;
+
+    // Step 1: Standard keygen (preprocessing + online)
+    let preproc_id_1 = derive_request_id("compressed_existing_preproc_1")?;
+    let keygen_id_1 = derive_request_id("compressed_existing_keygen_1")?;
+
+    threshold_key_gen_secure_isolated(
+        clients,
+        &preproc_id_1,
+        &keygen_id_1,
+        FheParameter::Test,
+        None,
+        None,
+        None,
+        None,
+    )
+    .await?;
+
+    // Verify standard keygen completed on all parties
+    for client in env.all_clients() {
+        let mut cur_client = client.clone();
+        let result = cur_client
+            .get_key_gen_result(tonic::Request::new(keygen_id_1.into()))
+            .await?;
+        assert_eq!(result.into_inner().request_id, Some(keygen_id_1.into()));
+    }
+
+    // Step 2: Compressed keygen from existing secret shares (preprocessing + online)
+    let preproc_id_2 = derive_request_id("compressed_existing_preproc_2")?;
+    let keygen_id_2 = derive_request_id("compressed_existing_keygen_2")?;
+
+    let (keyset_config, keyset_added_info) =
+        compressed_from_existing_keygen_config(&keygen_id_1, &DEFAULT_EPOCH_ID);
+
+    threshold_key_gen_secure_isolated(
+        clients,
+        &preproc_id_2,
+        &keygen_id_2,
+        FheParameter::Test,
+        keyset_config,
+        keyset_added_info,
+        None,
+        None,
+    )
+    .await?;
+
+    // Verify compressed keygen completed on all parties
+    for client in env.all_clients() {
+        let mut cur_client = client.clone();
+        let result = cur_client
+            .get_key_gen_result(tonic::Request::new(keygen_id_2.into()))
+            .await?;
+        assert_eq!(result.into_inner().request_id, Some(keygen_id_2.into()));
+    }
+
+    // Do distributed decryption to verify the generated key is ok
+    // TODO this could be refactored
+    use crate::client::tests::threshold::public_decryption_tests::run_decryption_threshold;
+    use crate::util::key_setup::test_tools::{EncryptionConfig, TestingPlaintext};
+    let material_dir = env.material_dir;
+    let mut servers = env.servers;
+    let mut clients = env.clients;
+
+    let material_path = material_dir.path();
+    let pub_storage_prefixes = &crate::consts::PUBLIC_STORAGE_PREFIX_THRESHOLD_ALL[0..4];
+
+    // Create internal client for decryption
+    let mut pub_storage_map = std::collections::HashMap::new();
+    for (i, prefix) in pub_storage_prefixes.iter().enumerate() {
+        pub_storage_map.insert(
+            (i + 1) as u32,
+            FileStorage::new(Some(material_path), StorageType::PUB, prefix.as_deref())?,
+        );
+    }
+    let client_storage = FileStorage::new(Some(material_path), StorageType::CLIENT, None)?;
+    let mut internal_client = crate::client::client_wasm::Client::new_client(
+        client_storage,
+        pub_storage_map,
+        &crate::consts::TEST_PARAM,
+        None,
+    )
+    .await?;
+
+    // Run ddec with the new keyset
+    run_decryption_threshold(
+        4,
+        &mut servers,
+        &mut clients,
+        &mut internal_client,
+        None,
+        &keygen_id_2,
+        None,
+        vec![TestingPlaintext::U32(66)],
+        EncryptionConfig {
+            compression: true,
+            precompute_sns: true,
+        },
+        None,
+        1,
+        Some(material_path),
+        true,
+    )
+    .await;
+
+    // Run ddec by encrypting using the old public key but
+    // still the new shares from the new keyset
+    run_decryption_threshold(
+        4,
+        &mut servers,
+        &mut clients,
+        &mut internal_client,
+        Some(&keygen_id_1),
+        &keygen_id_2,
+        None,
+        vec![TestingPlaintext::U32(55)],
+        EncryptionConfig {
+            compression: true,
+            precompute_sns: true,
+        },
+        None,
+        1,
+        Some(material_path),
+        false, // we do not used compressed_keys since that was the old public key
+    )
+    .await;
+
+    for (_, server) in servers {
+        server.assert_shutdown().await;
+    }
+
+    Ok(())
+}
+
 /// Test insecure threshold decompression key generation with decryption validation.
 ///
 /// Generates two regular keysets using insecure mode, then generates a decompression key
@@ -480,6 +653,8 @@ async fn test_insecure_threshold_decompression_keygen_isolated() -> Result<()> {
                 compression_keyset_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,
+                existing_epoch_id: None,
             }),
             context_id: None,
             epoch_id: None,
@@ -561,12 +736,13 @@ async fn test_insecure_threshold_decompression_keygen_isolated() -> Result<()> {
         &mut servers,
         &mut clients,
         &mut internal_client,
+        None,
         &key_id_1,
         None,
         vec![TestingPlaintext::U32(42)],
         EncryptionConfig {
             compression: true,
-            precompute_sns: false,
+            precompute_sns: true,
         },
         None,
         1,