feat: copy migrated key from new key ID to the old one

Author: kc1212

core/grpc/proto/kms.v1.proto

@@ -60,6 +60,16 @@ message KeySetAddedInfo {
   // This option is only used when KeyGenSecretKeyConfig::UseExisting is used.
   // Otherwise it is ignored.
   bool use_existing_key_tag = 5;
+
+  // If true, after a compressed keygen from existing shares
+  // (UseExisting + CompressedKeyConfig::All), copy the resulting
+  // CompressedXofKeySet to the original key ID and replace the original
+  // ThresholdFheKeys with the migrated one. The CompactPublicKey is
+  // preserved at both the original and the new key ID.
+  //
+  // Only used when KeyGenSecretKeyConfig::UseExisting is used with
+  // CompressedKeyConfig::All. Otherwise it is ignored.
+  bool copy_compressed_key_to_original = 4;
 }
 
 // The keyset configuration message.

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

@@ -78,6 +78,7 @@ pub(crate) fn keygen_config_from_existing(
             to_keyset_id_decompression_only: None,
             existing_keyset_id: Some((*existing_keyset_id).into()),
             use_existing_key_tag,
+            copy_compressed_key_to_original: false,
         }),
     )
 }
@@ -98,6 +99,7 @@ pub(crate) fn decompression_keygen_config(
             to_keyset_id_decompression_only: Some((*to_keyset_id).into()),
             existing_keyset_id: None,
             use_existing_key_tag: false,
+            copy_compressed_key_to_original: false,
         }),
     )
 }

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

@@ -807,6 +807,7 @@ async fn test_insecure_threshold_decompression_keygen_isolated() -> Result<()> {
                 to_keyset_id_decompression_only: Some(key_id_2.into()),
                 existing_keyset_id: None,
                 use_existing_key_tag: false,
+                copy_compressed_key_to_original: false,
             }),
             context_id: None,
             epoch_id: None,

core/service/src/engine/keyset_configuration.rs

@@ -182,6 +182,12 @@ impl InternalKeySetConfig {
             .is_some_and(|info| info.use_existing_key_tag)
     }
 
+    pub fn copy_compressed_key_to_original(&self) -> bool {
+        self.keyset_added_info
+            .as_ref()
+            .is_some_and(|info| info.copy_compressed_key_to_original)
+    }
+
     pub fn validate(&self) -> anyhow::Result<()> {
         match &self.keyset_config {
             ddec_keyset_config::KeySetConfig::Standard(inner) => match inner.secret_key_config {

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

@@ -1486,7 +1486,7 @@ impl<
                         threshold_fhe_keys,
                         &compressed_keyset,
                         &compact_pk,
-                        meta_store,
+                        Arc::clone(&meta_store),
                     )
                     .await
                 {
@@ -1495,6 +1495,50 @@ impl<
                     );
                     return;
                 }
+
+                // If requested, copy the compressed key to the original key ID.
+                //
+                // Note: at this point the *new* keygen has already been committed
+                // and its meta_store entry is Done — that part of the request
+                // succeeded. The copy is a follow-up on a *different* key id
+                // (old_key_id), so a failure here does not invalidate the
+                // new_key_id material. We log loudly so operators can detect
+                // partial success and retry the migration, but we do not try
+                // to mark the new keygen itself as failed.
+                //
+                // TODO if copying fails, do we still call `update_backup_vault`?
+                if matches!(
+                    keyset_config.secret_key_config,
+                    ddec_keyset_config::KeyGenSecretKeyConfig::UseExisting
+                ) && internal_keyset_config.copy_compressed_key_to_original()
+                {
+                    let old_key_id = internal_keyset_config
+                        .get_existing_keyset_id()
+                        .expect("validated");
+                    // UseExisting reads the old private shares at the current
+                    // epoch_id (see key_gen_from_existing_private_keyset), so
+                    // the copy targets the same (old_key_id, epoch_id) pair.
+                    if let Err(e) = crypto_storage
+                        .copy_compressed_key_to_original(
+                            req_id,
+                            epoch_id,
+                            &old_key_id,
+                            epoch_id,
+                            &sk,
+                            &eip712_domain,
+                            Arc::clone(&meta_store),
+                        )
+                        .await
+                    {
+                        tracing::error!(
+                            "Compressed keygen for {req_id} committed successfully, but the \
+                             follow-up copy to original key id {old_key_id} failed: {e}. \
+                             The new keys at {req_id} are valid; \
+                             the migration to {old_key_id} must be retried."
+                        );
+                        return;
+                    }
+                }
             }
         }
         // Update the backup and handle potential failures by incrementing backup errors in the metrics