feat: propagate tag during keygen from existing shares (#444)

Closes zama-ai/kms-internal#2920

Author: kc1212

core-client/src/keygen.rs

@@ -85,6 +85,7 @@ pub(crate) async fn do_keygen(
             .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()
             });
     let dkg_req = internal_client.key_gen_request(

core-client/src/lib.rs

@@ -595,6 +595,10 @@ pub struct SharedKeyGenParameters {
     /// 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)]
+    pub use_existing_key_tag: bool,
     pub context_id: Option<ContextId>,
     pub epoch_id: Option<EpochId>,
 }

core-client/tests/integration/integration_test.rs

@@ -1632,6 +1632,7 @@ async fn test_threshold_mpc_context_switch_6(ctx: &DockerComposeThresholdTestNoI
                 compressed: false,
                 existing_keyset_id: None,
                 existing_epoch_id: None,
+                use_existing_key_tag: false,
             },
             200,
         )
@@ -1842,6 +1843,7 @@ async fn test_threshold_reshare(ctx: &DockerComposeThresholdTestNoInitSixParty)
             compressed: false,
             existing_keyset_id: None,
             existing_epoch_id: None,
+            use_existing_key_tag: false,
         },
         200,
     )

core-client/tests/integration/integration_test_isolated.rs

@@ -2028,6 +2028,7 @@ async fn real_preproc_and_keygen_compressed_isolated(
                 compressed: true,
                 existing_keyset_id: None,
                 existing_epoch_id: None,
+                use_existing_key_tag: false,
                 context_id: None,
                 epoch_id: None,
             },

core/grpc/proto/kms.v1.proto

@@ -57,6 +57,13 @@ message KeySetAddedInfo {
   // 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.
+  //
+  // This option is only used when KeyGenSecretKeyConfig::UseExisting is used.
+  // Otherwise it is ignored.
+  bool use_existing_key_tag = 5;
 }
 
 // The keyset configuration message.

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

@@ -48,6 +48,7 @@ pub(crate) fn compressed_keygen_config() -> (Option<KeySetConfig>, Option<KeySet
 pub(crate) fn compressed_from_existing_keygen_config(
     existing_keyset_id: &RequestId,
     existing_epoch_id: &kms_grpc::identifiers::EpochId,
+    use_existing_key_tag: bool,
 ) -> (Option<KeySetConfig>, Option<KeySetAddedInfo>) {
     (
         Some(KeySetConfig {
@@ -63,6 +64,7 @@ pub(crate) fn compressed_from_existing_keygen_config(
             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,
         }),
     )
 }
@@ -83,6 +85,7 @@ pub(crate) fn decompression_keygen_config(
             to_keyset_id_decompression_only: Some((*to_keyset_id).into()),
             existing_keyset_id: None,
             existing_epoch_id: None,
+            use_existing_key_tag: false,
         }),
     )
 }

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

@@ -477,7 +477,7 @@ async fn secure_threshold_compressed_keygen_from_existing_isolated() -> Result<(
     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);
+        compressed_from_existing_keygen_config(&keygen_id_1, &DEFAULT_EPOCH_ID, true);
 
     threshold_key_gen_secure_isolated(
         clients,
@@ -519,6 +519,30 @@ async fn secure_threshold_compressed_keygen_from_existing_isolated() -> Result<(
             FileStorage::new(Some(material_path), StorageType::PUB, prefix.as_deref())?,
         );
     }
+
+    // Verify tag propagation: keys from keygen_id_2 should carry keygen_id_1's tag
+    {
+        use crate::vault::storage::crypto_material::CryptoMaterialReader;
+        use tfhe::prelude::Tagged;
+        let expected_tag: tfhe::Tag = keygen_id_1.into();
+        for (&party_id, storage) in &pub_storage_map {
+            let compressed_keyset: tfhe::xof_key_set::CompressedXofKeySet =
+                CryptoMaterialReader::read_from_storage(storage, &keygen_id_2).await?;
+
+            let (pk, server_key) = compressed_keyset.decompress().unwrap().into_raw_parts();
+            assert_eq!(
+                pk.tag(),
+                &expected_tag,
+                "Public key for party {party_id} should have tag propagated from existing keyset"
+            );
+            assert_eq!(
+                server_key.tag(),
+                &expected_tag,
+                "Server key for party {party_id} should have tag propagated from existing keyset"
+            );
+        }
+    }
+
     let client_storage = FileStorage::new(Some(material_path), StorageType::CLIENT, None)?;
     let mut internal_client = crate::client::client_wasm::Client::new_client(
         client_storage,
@@ -712,6 +736,7 @@ async fn test_insecure_threshold_decompression_keygen_isolated() -> Result<()> {
                 to_keyset_id_decompression_only: Some(key_id_2.into()),
                 existing_keyset_id: None,
                 existing_epoch_id: None,
+                use_existing_key_tag: false,
             }),
             context_id: None,
             epoch_id: None,

core/service/src/engine/keyset_configuration.rs

@@ -175,6 +175,12 @@ impl InternalKeySetConfig {
         Ok(keyset_id)
     }
 
+    pub fn use_existing_key_tag(&self) -> bool {
+        self.keyset_added_info
+            .as_ref()
+            .is_some_and(|info| info.use_existing_key_tag)
+    }
+
     pub fn get_existing_epoch_id(&self) -> anyhow::Result<Option<EpochId>> {
         let added_info = self
             .keyset_added_info

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

@@ -18,6 +18,7 @@ use observability::{
     },
 };
 use tfhe::integer::compression_keys::DecompressionKey;
+use tfhe::prelude::Tagged;
 use tfhe::xof_key_set::CompressedXofKeySet;
 use threshold_fhe::{
     algebra::{
@@ -77,14 +78,18 @@ use crate::{
         },
         rate_limiter::RateLimiter,
     },
-    vault::storage::{crypto_material::ThresholdCryptoMaterialStorage, Storage, StorageExt},
+    vault::storage::{
+        crypto_material::{CryptoMaterialReader, ThresholdCryptoMaterialStorage},
+        Storage, StorageExt,
+    },
 };
 
 // === Current Module Imports ===
 use super::BucketMetaStore;
 
 const DKG_Z64_SESSION_COUNTER: u64 = 1;
 const DKG_Z128_SESSION_COUNTER: u64 = 2;
+const ERR_FAILED_TO_READ_EXISTING_TAG: &str = "Failed to read existing tag";
 
 struct DkgSessions {
     session_z64: SmallSession<ResiduePolyF4Z64>,
@@ -342,6 +347,14 @@ impl<
         // we must validate the parameter before passing it into the background process
         internal_keyset_config.validate()?;
 
+        // Read existing key tag from public storage if needed
+        let existing_key_tag: Option<tfhe::Tag> = if internal_keyset_config.use_existing_key_tag() {
+            let existing_keyset_id = internal_keyset_config.get_existing_keyset_id()?;
+            Some(Self::read_existing_key_tag(&crypto_storage, &existing_keyset_id).await?)
+        } else {
+            None
+        };
+
         let keygen_background = async move {
             // Remove the preprocessing material, even if the request was cancelled we cannot reuse the preprocessing
             match &preproc_handle_w_mode {
@@ -366,6 +379,7 @@ impl<
                     // Nothing to remove
                 }
             }
+
             match internal_keyset_config.keyset_config() {
                 ddec_keyset_config::KeySetConfig::Standard(inner_config) => {
                     Self::key_gen_background(
@@ -382,6 +396,7 @@ impl<
                         eip712_domain_copy,
                         permit,
                         op_tag,
+                        existing_key_tag,
                     )
                     .await
                 }
@@ -1005,13 +1020,13 @@ impl<
 
     #[allow(clippy::too_many_arguments)]
     async fn compressed_key_gen_from_existing_private_keyset<P>(
-        req_id: &RequestId,
         dkg_sessions: &mut DkgSessions,
         crypto_storage: ThresholdCryptoMaterialStorage<PubS, PrivS>,
         params: DKGParams,
         existing_keyset_id: RequestId,
         existing_epoch_id: EpochId,
         preprocessing: &mut P,
+        tag: tfhe::Tag,
     ) -> anyhow::Result<(
         CompressedXofKeySet,
         PrivateKeySet<{ ResiduePolyF4Z128::EXTENSION_DEGREE }>,
@@ -1040,7 +1055,7 @@ impl<
             &mut dkg_sessions.session_z128,
             preprocessing,
             params,
-            req_id.into(),
+            tag,
             &existing_private_keys,
         )
         .await?;
@@ -1049,13 +1064,13 @@ impl<
 
     #[allow(clippy::too_many_arguments)]
     async fn key_gen_from_existing_private_keyset<P>(
-        req_id: &RequestId,
         dkg_sessions: &mut DkgSessions,
         crypto_storage: ThresholdCryptoMaterialStorage<PubS, PrivS>,
         params: DKGParams,
         existing_keyset_id: RequestId,
         existing_epoch_id: EpochId,
         preprocessing: &mut P,
+        tag: tfhe::Tag,
     ) -> anyhow::Result<(
         FhePubKeySet,
         PrivateKeySet<{ ResiduePolyF4Z128::EXTENSION_DEGREE }>,
@@ -1084,7 +1099,7 @@ impl<
             &mut dkg_sessions.session_z128,
             preprocessing,
             params,
-            req_id.into(),
+            tag,
             &existing_private_keys,
         )
         .await?;
@@ -1106,6 +1121,7 @@ impl<
         eip712_domain: alloy_sol_types::Eip712Domain,
         permit: OwnedSemaphorePermit,
         op_tag: &'static str,
+        existing_key_tag: Option<tfhe::Tag>,
     ) {
         let _permit = permit;
         let start = Instant::now();
@@ -1209,14 +1225,15 @@ impl<
                                 .get_existing_epoch_id()
                                 .expect("validated")
                                 .unwrap_or(*epoch_id);
+                            let tag: tfhe::Tag = existing_key_tag.unwrap_or_else(|| req_id.into());
                             Self::key_gen_from_existing_private_keyset(
-                                req_id,
                                 &mut dkg_sessions,
                                 crypto_storage.clone(),
                                 params,
                                 existing_keyset_id,
                                 existing_epoch_id,
                                 preproc_handle.as_mut(),
+                                tag,
                             )
                             .await
                             .map(|(pk, sk)| ThresholdKeyGenResult::Uncompressed(pk, sk))
@@ -1233,14 +1250,15 @@ impl<
                                 .get_existing_epoch_id()
                                 .expect("validated")
                                 .unwrap_or(*epoch_id);
+                            let tag: tfhe::Tag = existing_key_tag.unwrap_or_else(|| req_id.into());
                             Self::compressed_key_gen_from_existing_private_keyset(
-                                req_id,
                                 &mut dkg_sessions,
                                 crypto_storage.clone(),
                                 params,
                                 existing_keyset_id,
                                 existing_epoch_id,
                                 preproc_handle.as_mut(),
+                                tag,
                             )
                             .await
                             .map(|(compressed_keyset, sk)| {
@@ -1392,6 +1410,36 @@ impl<
             start.elapsed().as_millis()
         );
     }
+
+    /// Reads the tag from an existing keyset in public storage.
+    /// Tries CompressedXofKeySet first, then falls back to ServerKey.
+    async fn read_existing_key_tag(
+        crypto_storage: &ThresholdCryptoMaterialStorage<PubS, PrivS>,
+        existing_keyset_id: &RequestId,
+    ) -> anyhow::Result<tfhe::Tag> {
+        let pub_storage = crypto_storage.inner.public_storage.lock().await;
+
+        let res = if let Ok(compressed_keyset) =
+            <CompressedXofKeySet as CryptoMaterialReader>::read_from_storage(
+                &*pub_storage,
+                existing_keyset_id,
+            )
+            .await
+        {
+            Ok(compressed_keyset
+                .clone()
+                .into_raw_parts()
+                .1
+                .into_raw_parts()
+                .1)
+        } else {
+            CryptoMaterialReader::read_from_storage(&*pub_storage, existing_keyset_id)
+                .await
+                .map(|server_key: tfhe::ServerKey| server_key.tag().clone())
+        };
+
+        res.map_err(|e| anyhow::anyhow!("{}: {e}", ERR_FAILED_TO_READ_EXISTING_TAG))
+    }
 }
 
 #[tonic::async_trait]
@@ -1800,6 +1848,54 @@ mod tests {
         // we don't have a trait for meta store
     }
 
+    #[tokio::test]
+    async fn use_existing_key_tag_with_wrong_keyset_id() {
+        // When use_existing_key_tag is true but existing_keyset_id points to a
+        // non-existent key in storage, launch_dkg should fail with Internal
+        // because read_existing_key_tag cannot find any key material.
+        let (prep_ids, kg) = setup_key_generator::<
+            DroppingOnlineDistributedKeyGen128<{ ResiduePolyF4Z128::EXTENSION_DEGREE }>,
+        >()
+        .await;
+        let prep_id = prep_ids[0];
+        let key_id = RequestId::new_random(&mut OsRng);
+        let wrong_keyset_id = RequestId::new_random(&mut OsRng);
+
+        let domain = alloy_to_protobuf_domain(&dummy_domain()).unwrap();
+        let keyset_config = KeySetConfig {
+            keyset_type: kms_grpc::kms::v1::KeySetType::Standard as i32,
+            standard_keyset_config: Some(kms_grpc::kms::v1::StandardKeySetConfig {
+                compute_key_type: 0,
+                secret_key_config: kms_grpc::kms::v1::KeyGenSecretKeyConfig::UseExisting as i32,
+                compressed_key_config: 0,
+            }),
+        };
+        let keyset_added_info = KeySetAddedInfo {
+            existing_keyset_id: Some(wrong_keyset_id.into()),
+            use_existing_key_tag: true,
+            ..Default::default()
+        };
+
+        let request = tonic::Request::new(KeyGenRequest {
+            request_id: Some(key_id.into()),
+            params: Some(FheParameter::Test as i32),
+            preproc_id: Some(prep_id.into()),
+            domain: Some(domain),
+            keyset_config: Some(keyset_config),
+            keyset_added_info: Some(keyset_added_info),
+            context_id: Some((*DEFAULT_MPC_CONTEXT).into()),
+            epoch_id: None,
+        });
+
+        let res = kg.key_gen(request).await.unwrap_err();
+        assert_eq!(res.code(), tonic::Code::Internal);
+
+        assert!(res
+            .internal_err()
+            .to_string()
+            .contains(ERR_FAILED_TO_READ_EXISTING_TAG));
+    }
+
     #[tokio::test]
     async fn sunshine() {
         let (prep_ids, kg) = setup_key_generator::<