chore: remove existing_epoch_id

Author: kc1212

core-client/src/keygen.rs

@@ -111,7 +111,6 @@ pub(crate) async fn do_keygen(
             .existing_keyset_id
             .map(|id| kms_grpc::kms::v1::KeySetAddedInfo {
                 existing_keyset_id: Some(id.into()),
-                existing_epoch_id: shared_config.existing_epoch_id.map(Into::into),
                 use_existing_key_tag: shared_config.use_existing_key_tag,
                 ..Default::default()
             });

core-client/src/lib.rs

@@ -615,9 +615,6 @@ pub struct SharedKeyGenParameters {
     /// instead of running full distributed keygen.
     #[clap(long)]
     pub existing_keyset_id: Option<RequestId>,
-    /// Epoch ID for the existing keyset (optional, defaults to the request's epoch).
-    #[clap(long)]
-    pub existing_epoch_id: Option<EpochId>,
     /// Reuse the tag from the existing keyset instead of using the new key ID as tag.
     /// This is only used when generating a key from existing shares.
     #[clap(long, default_value_t = false)]
@@ -1263,16 +1260,16 @@ pub async fn fetch_ctxt_from_file(
 /// compressed storage or the legacy uncompressed layout.
 ///
 /// If `uncompressed_keys` is explicitly `true`, fetches `[PublicKey, ServerKey]` only.
-/// Otherwise, tries `[CompressedXofKeySet]` first; on failure, falls back to
-/// `[PublicKey, ServerKey]`.
+/// Otherwise, tries the current compressed layout `[CompressedXofKeySet, PublicKey]`
+/// first; on failure, falls back to the legacy `[PublicKey, ServerKey]`.
 /// Returns the fetched party confs and a boolean indicating whether uncompressed keys were found.
 async fn fetch_keys_auto_detect(
     key_id: &str,
     uncompressed_keys: bool,
     cc_conf: &CoreClientConfig,
     destination_prefix: &Path,
 ) -> anyhow::Result<(Vec<CoreConf>, bool)> {
-    let compressed_key_types = vec![PubDataType::CompressedXofKeySet];
+    let compressed_key_types = vec![PubDataType::CompressedXofKeySet, PubDataType::PublicKey];
     let key_types = vec![PubDataType::PublicKey, PubDataType::ServerKey];
 
     if uncompressed_keys {
@@ -1293,7 +1290,7 @@ async fn fetch_keys_auto_detect(
         Ok(confs) => Ok((confs, false)),
         Err(_) => {
             tracing::info!(
-                "CompressedXofKeySet not found, trying legacy [PublicKey, ServerKey]..."
+                "Compressed layout [CompressedXofKeySet, PublicKey] not found, trying legacy [PublicKey, ServerKey]..."
             );
             let confs =
                 fetch_public_elements(key_id, &key_types, cc_conf, destination_prefix, false)
@@ -2387,7 +2384,13 @@ fn print_timings(cmd: &str, durations: &mut [tokio::time::Duration], start: toki
 #[cfg(test)]
 mod tests {
     use super::*;
+    use kms_lib::engine::base::derive_request_id;
+    use kms_lib::util::key_setup::test_tools::load_pk_from_pub_storage;
+    use kms_lib::vault::storage::{StorageType, file::FileStorage, store_versioned_at_request_id};
     use std::env;
+    use tempfile::tempdir;
+    use tfhe::core_crypto::prelude::NormalizedHammingWeightBound;
+    use tfhe::xof_key_set::CompressedXofKeySet;
 
     #[test]
     fn test_parse_hex() {
@@ -2563,4 +2566,120 @@ mod tests {
         );
         assert!(PreviousEpochParameters::from_str(&input_string).is_err());
     }
+
+    #[tokio::test]
+    async fn fetch_keys_auto_detect_downloads_public_key_for_compressed_layout() {
+        let remote_root = tempdir().unwrap();
+        let destination_root = tempdir().unwrap();
+        let object_folder = "PUB-p1";
+        let key_id = derive_request_id("fetch_keys_auto_detect_downloads_public_key").unwrap();
+
+        let params = kms_lib::consts::TEST_PARAM;
+        let config = params.to_tfhe_config();
+        let max_norm_hwt = params
+            .get_params_basics_handle()
+            .get_sk_deviations()
+            .map(|x| x.pmax)
+            .unwrap_or(1.0);
+        let max_norm_hwt = NormalizedHammingWeightBound::new(max_norm_hwt).unwrap();
+        let (_client_key, compressed_keyset) = CompressedXofKeySet::generate(
+            config,
+            vec![1, 2, 3, 4],
+            params.get_params_basics_handle().get_sec() as u32,
+            max_norm_hwt,
+            key_id.into(),
+        )
+        .unwrap();
+        let (public_key, _server_key) = compressed_keyset
+            .clone()
+            .decompress()
+            .unwrap()
+            .into_raw_parts();
+
+        let mut remote_storage = FileStorage::new(
+            Some(remote_root.path()),
+            StorageType::PUB,
+            Some(object_folder),
+        )
+        .unwrap();
+        store_versioned_at_request_id(
+            &mut remote_storage,
+            &key_id,
+            &compressed_keyset,
+            &PubDataType::CompressedXofKeySet.to_string(),
+        )
+        .await
+        .unwrap();
+        store_versioned_at_request_id(
+            &mut remote_storage,
+            &key_id,
+            &public_key,
+            &PubDataType::PublicKey.to_string(),
+        )
+        .await
+        .unwrap();
+
+        let cc_conf = CoreClientConfig {
+            kms_type: KmsType::Threshold,
+            cores: vec![CoreConf {
+                party_id: 1,
+                address: "127.0.0.1:0".to_string(),
+                s3_endpoint: format!("file://{}", remote_root.path().display()),
+                object_folder: object_folder.to_string(),
+                #[cfg(feature = "testing")]
+                private_object_folder: None,
+                #[cfg(feature = "testing")]
+                config_path: None,
+            }],
+            decryption_mode: None,
+            num_majority: 1,
+            num_reconstruct: 1,
+            fhe_params: Some(FheParameter::Test),
+        };
+
+        let (party_confs, detected_uncompressed) = fetch_keys_auto_detect(
+            &key_id.to_string(),
+            false,
+            &cc_conf,
+            destination_root.path(),
+        )
+        .await
+        .unwrap();
+
+        assert_eq!(party_confs.len(), 1);
+        assert!(
+            !detected_uncompressed,
+            "compressed layout should not fall back to legacy uncompressed keys"
+        );
+
+        let downloaded_pk_path = destination_root
+            .path()
+            .join(object_folder)
+            .join(PubDataType::PublicKey.to_string())
+            .join(key_id.to_string());
+        assert!(
+            downloaded_pk_path.exists(),
+            "compressed auto-detect should download the authoritative standalone PublicKey"
+        );
+
+        let _downloaded_pk =
+            load_pk_from_pub_storage(Some(destination_root.path()), &key_id, Some(object_folder))
+                .await;
+        let (ciphertext, _format, _fhe_type) = compute_cipher_from_stored_key(
+            Some(destination_root.path()),
+            TestingPlaintext::U8(42),
+            &key_id,
+            Some(object_folder),
+            EncryptionConfig {
+                compression: true,
+                precompute_sns: false,
+            },
+            false,
+        )
+        .await;
+        assert!(
+            !ciphertext.is_empty(),
+            "encryption should succeed from freshly fetched compressed key material"
+        );
+    }
 }

core-client/src/mpc_epoch.rs

@@ -261,7 +261,7 @@ pub(crate) async fn do_new_epoch(
 
             let (party_confs_successful, is_compressed) = match fetch_public_elements(
                 &key_id.to_string(),
-                &[PubDataType::CompressedXofKeySet],
+                &[PubDataType::CompressedXofKeySet, PubDataType::PublicKey],
                 cc_conf,
                 destination_prefix,
                 true,

core-client/tests/integration/integration_test.rs

@@ -22,7 +22,7 @@ use kms_lib::engine::base::derive_request_id;
 use kms_lib::engine::base::safe_serialize_hash_element_versioned;
 use kms_lib::util::key_setup::test_tools::load_material_from_pub_storage;
 use serial_test::serial;
-use std::fs::create_dir_all;
+use std::fs::{create_dir_all, remove_file};
 use std::io::Write;
 use std::path::Path;
 use std::path::PathBuf;
@@ -1791,7 +1791,6 @@ async fn test_threshold_mpc_context_switch_6(ctx: &DockerComposeThresholdTestNoI
                 epoch_id: Some(epoch_id_2),
                 uncompressed: false,
                 existing_keyset_id: None,
-                existing_epoch_id: None,
                 use_existing_key_tag: false,
                 extra_data: None,
             },
@@ -1889,10 +1888,39 @@ async fn test_threshold_reshare(ctx: &DockerComposeThresholdTestNoInitSixParty)
         .build()
         .init_conf()
         .unwrap();
+    let clear_local_public_key_cache = |cc_conf: &CoreClientConfig, key_id: &str| {
+        for core in &cc_conf.cores {
+            let path = test_path
+                .join(&core.object_folder)
+                .join(PubDataType::PublicKey.to_string())
+                .join(key_id);
+            if path.exists() {
+                remove_file(&path).unwrap();
+            }
+            assert!(
+                !path.exists(),
+                "expected no cached PublicKey at {:?} before exercising fresh download",
+                path
+            );
+        }
+    };
+    let assert_some_public_key_cached = |cc_conf: &CoreClientConfig, key_id: &str| {
+        assert!(
+            cc_conf.cores.iter().any(|core| {
+                test_path
+                    .join(&core.object_folder)
+                    .join(PubDataType::PublicKey.to_string())
+                    .join(key_id)
+                    .exists()
+            }),
+            "expected at least one PublicKey artifact to be downloaded for key {}",
+            key_id
+        );
+    };
 
     let party_confs = fetch_public_elements(
         &key_id,
-        &[PubDataType::CompressedXofKeySet],
+        &[PubDataType::CompressedXofKeySet, PubDataType::PublicKey],
         &cc_conf,
         test_path,
         false,
@@ -1931,6 +1959,8 @@ async fn test_threshold_reshare(ctx: &DockerComposeThresholdTestNoInitSixParty)
     let crs_digest =
         hex::encode(safe_serialize_hash_element_versioned(&DSEP_PUBDATA_CRS, &crs).unwrap());
 
+    clear_local_public_key_cache(&cc_conf, &key_id.to_string());
+
     // create and store second mpc context
     // create and store mpc context
     println!("Creating second MPC context");
@@ -1982,6 +2012,15 @@ async fn test_threshold_reshare(ctx: &DockerComposeThresholdTestNoInitSixParty)
     };
     let result = execute_cmd(&epoch_config, test_path).await.unwrap();
     println!("Resharing completed successfully {:?}", result);
+    assert_some_public_key_cached(&cc_conf, &key_id.to_string());
+
+    let cc_conf_set_2: CoreClientConfig = observability::conf::Settings::builder()
+        .path(config_path_set_2.to_str().unwrap())
+        .env_prefix("CORE_CLIENT")
+        .build()
+        .init_conf()
+        .unwrap();
+    clear_local_public_key_cache(&cc_conf_set_2, &key_id.to_string());
 
     // Try and do a decrypt with the new set and new epoch
     let ddec_command = CCCommand::PublicDecrypt(CipherArguments::FromArgs(CipherParameters {
@@ -2012,6 +2051,7 @@ async fn test_threshold_reshare(ctx: &DockerComposeThresholdTestNoInitSixParty)
 
     let result = execute_cmd(&ddec_config, test_path).await.unwrap();
     println!("Decrypt in reshared context succeeded : {:?}", result);
+    assert_some_public_key_cached(&cc_conf_set_2, &key_id.to_string());
 
     // Delete the old epoch and verify it's gone
     println!("Destroying old epoch");
@@ -2026,7 +2066,6 @@ async fn test_threshold_reshare(ctx: &DockerComposeThresholdTestNoInitSixParty)
             epoch_id: Some(epoch_id_set_1),
             uncompressed: false,
             existing_keyset_id: None,
-            existing_epoch_id: None,
             use_existing_key_tag: false,
             extra_data: None,
         },

core-client/tests/integration/integration_test_isolated.rs

@@ -3284,7 +3284,7 @@ async fn test_threshold_reshare() -> Result<()> {
 
     let ids = fetch_public_elements(
         &key_id_str,
-        &[PubDataType::CompressedXofKeySet],
+        &[PubDataType::CompressedXofKeySet, PubDataType::PublicKey],
         &cc_conf,
         test_path,
         false,
@@ -3350,6 +3350,28 @@ async fn test_threshold_reshare() -> Result<()> {
 
     info!("Resharing result: {:?}", resharing_result);
     assert_eq!(resharing_result.len(), 2);
+    let ddec_config = cmd_config(
+        &config_path,
+        CCCommand::PublicDecrypt(CipherArguments::FromArgs(CipherParameters {
+            to_encrypt: "0x123456".to_string(),
+            data_type: FheType::Euint64,
+            no_compression: false,
+            no_precompute_sns: true,
+            key_id: KeyId::from_str(&key_id.to_string()).unwrap(),
+            context_id: Some(context_id),
+            epoch_id: Some(new_epoch_id),
+            batch_size: 1,
+            num_requests: 1,
+            ciphertext_output_path: None,
+            parallel_requests: 1,
+            inter_request_delay_ms: 0,
+            uncompressed_keys: false,
+            extra_data: None,
+        })),
+        200,
+    );
+    let decrypt_result = execute_cmd(&ddec_config, test_path).await.unwrap();
+    info!("Decrypt in reshared epoch succeeded: {:?}", decrypt_result);
 
     // The second element is the previous epoch_id used for reshare
     assert_eq!(resharing_result[1].0.unwrap(), epoch_id.into());

core/grpc/proto/kms.v1.proto

@@ -54,10 +54,6 @@ message KeySetAddedInfo {
   // Must be set if KeyGenSecretKeyConfig::UseExisting is used
   RequestId existing_keyset_id = 3;
 
-  // 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 = 4;
-
   // If true, use the tag from the existing keyset's PublicKeyMaterial
   // instead of deriving the tag from the new request ID.
   //

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

@@ -62,7 +62,6 @@ pub(crate) fn uncompressed_keygen_config() -> (Option<KeySetConfig>, Option<KeyS
 #[cfg(feature = "slow_tests")]
 pub(crate) fn keygen_config_from_existing(
     existing_keyset_id: &RequestId,
-    existing_epoch_id: &kms_grpc::identifiers::EpochId,
     use_existing_key_tag: bool,
 ) -> (Option<KeySetConfig>, Option<KeySetAddedInfo>) {
     (
@@ -78,7 +77,6 @@ pub(crate) fn keygen_config_from_existing(
             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()),
             use_existing_key_tag,
         }),
     )
@@ -99,7 +97,6 @@ pub(crate) fn decompression_keygen_config(
             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,
             use_existing_key_tag: false,
         }),
     )