Reshare in a spawned task and fix propagation pallet rotate keyshares…

… endpoint lookup key (#1185) * Reshare happens in a spawned task, fix local service rotate keyshares lookup key * Reshare happens in a spawned task, fix local service rotate keyshares lookup key * Clippy * Wait for protocol to finish during reshare test * Fix integration test * Changelog
entropyxyz · Nov 27, 2024 · b6884ec · b6884ec
1 parent d62db3b
commit b6884ec
Show file tree

Hide file tree

Showing 5 changed files with 76 additions and 41 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -41,6 +41,10 @@ runtime
 - Update programs to accept multiple oracle data ([#1153](https://github.com/entropyxyz/entropy-core/pull/1153/))
 - Use context, not block number in TDX quote input data ([#1179](https://github.com/entropyxyz/entropy-core/pull/1179))
 
+### Fixed
+
+- Reshare in a spawned task and fix propagation pallet rotate keyshares endpoint lookup key ([#1185](https://github.com/entropyxyz/entropy-core/pull/1185))
+
 ## [0.3.0](https://github.com/entropyxyz/entropy-core/compare/release/v0.2.0...release/v0.3.0) - 2024-10-22
 
 ### Breaking Changes

diff --git a/crates/threshold-signature-server/src/validator/api.rs b/crates/threshold-signature-server/src/validator/api.rs
@@ -44,6 +44,7 @@ use subxt::{
     OnlineClient,
 };
 use synedrion::{KeyResharingInputs, NewHolder, OldHolder};
+use x25519_dalek::StaticSecret;
 
 /// HTTP POST endpoint called by the off-chain worker (propagation pallet) during network reshare.
 ///
@@ -61,22 +62,50 @@ pub async fn new_reshare(
     let rpc = get_rpc(&app_state.configuration.endpoint).await?;
     validate_new_reshare(&api, &rpc, &data, &app_state.kv_store).await?;
 
+    let (signer, x25519_secret_key) = get_signer_and_x25519_secret(&app_state.kv_store)
+        .await
+        .map_err(|e| ValidatorErr::UserError(e.to_string()))?;
+
     let next_signers_query = entropy::storage().staking_extension().next_signers();
     let next_signers = query_chain(&api, &rpc, next_signers_query, None)
         .await?
         .ok_or_else(|| ValidatorErr::ChainFetch("Error getting next signers"))?
         .next_signers;
-
     let validators_info = get_validators_info(&api, &rpc, next_signers)
         .await
         .map_err(|e| ValidatorErr::UserError(e.to_string()))?;
 
-    let (signer, x25519_secret_key) = get_signer_and_x25519_secret(&app_state.kv_store)
-        .await
-        .map_err(|e| ValidatorErr::UserError(e.to_string()))?;
+    let is_proper_signer = validators_info
+        .iter()
+        .any(|validator_info| validator_info.tss_account == *signer.account_id());
+
+    if !is_proper_signer {
+        return Ok(StatusCode::MISDIRECTED_REQUEST);
+    }
 
+    // Do reshare in a separate task so we can already respond
+    tokio::spawn(async move {
+        if let Err(err) =
+            do_reshare(&api, &rpc, signer, &x25519_secret_key, data, validators_info, app_state)
+                .await
+        {
+            tracing::error!("Error during reshare: {err}");
+        }
+    });
+    Ok(StatusCode::OK)
+}
+
+async fn do_reshare(
+    api: &OnlineClient<EntropyConfig>,
+    rpc: &LegacyRpcMethods<EntropyConfig>,
+    signer: PairSigner<EntropyConfig, sr25519::Pair>,
+    x25519_secret_key: &StaticSecret,
+    data: OcwMessageReshare,
+    validators_info: Vec<ValidatorInfo>,
+    app_state: AppState,
+) -> Result<(), ValidatorErr> {
     let verifying_key_query = entropy::storage().staking_extension().jump_start_progress();
-    let parent_key_details = query_chain(&api, &rpc, verifying_key_query, None)
+    let parent_key_details = query_chain(api, rpc, verifying_key_query, None)
         .await?
         .ok_or_else(|| ValidatorErr::ChainFetch("Parent verifying key error"))?;
 
@@ -92,16 +121,7 @@ pub async fn new_reshare(
             .map_err(|_| ValidatorErr::Conversion("Verifying key conversion"))?,
     )
     .map_err(|e| ValidatorErr::VerifyingKeyError(e.to_string()))?;
-
-    let is_proper_signer = validators_info
-        .iter()
-        .any(|validator_info| validator_info.tss_account == *signer.account_id());
-
-    if !is_proper_signer {
-        return Ok(StatusCode::MISDIRECTED_REQUEST);
-    }
-
-    let my_stash_address = get_stash_address(&api, &rpc, signer.account_id())
+    let my_stash_address = get_stash_address(api, rpc, signer.account_id())
         .await
         .map_err(|e| ValidatorErr::UserError(e.to_string()))?;
 
@@ -122,7 +142,7 @@ pub async fn new_reshare(
         validators_info.iter().cloned().map(|x| PartyId::new(x.tss_account)).collect();
     // old holders -> next_signers - new_signers (will be at least t)
     let old_holders =
-        &prune_old_holders(&api, &rpc, data.new_signers, validators_info.clone()).await?;
+        &prune_old_holders(api, rpc, data.new_signers, validators_info.clone()).await?;
     let old_holders: BTreeSet<PartyId> =
         old_holders.iter().map(|x| PartyId::new(x.tss_account.clone())).collect();
 
@@ -132,7 +152,7 @@ pub async fn new_reshare(
         old_holders,
     };
     let key_info_query = entropy::storage().parameters().signers_info();
-    let threshold = query_chain(&api, &rpc, key_info_query, None)
+    let threshold = query_chain(api, rpc, key_info_query, None)
         .await?
         .ok_or_else(|| ValidatorErr::ChainFetch("Failed to get signers info"))?
         .threshold;
@@ -158,16 +178,16 @@ pub async fn new_reshare(
         converted_validator_info.push(validator_info.clone());
         tss_accounts.push(validator_info.tss_account.clone());
     }
+
     let channels = get_channels(
         &app_state.listener_state,
         converted_validator_info,
         account_id,
         &session_id,
         &signer,
-        &x25519_secret_key,
+        x25519_secret_key,
     )
     .await?;
-
     let (new_key_share, aux_info) =
         execute_reshare(session_id.clone(), channels, signer.signer(), inputs, &new_holders, None)
             .await?;
@@ -184,8 +204,8 @@ pub async fn new_reshare(
     app_state.kv_store.kv().put(reservation, serialized_key_share.clone()).await?;
 
     // TODO: Error handling really complex needs to be thought about.
-    confirm_key_reshare(&api, &rpc, &signer).await?;
-    Ok(StatusCode::OK)
+    confirm_key_reshare(api, rpc, &signer).await?;
+    Ok(())
 }
 
 /// HTTP POST endpoint called by the off-chain worker (propagation pallet) after a network key reshare.

diff --git a/crates/threshold-signature-server/src/validator/tests.rs b/crates/threshold-signature-server/src/validator/tests.rs
@@ -62,6 +62,7 @@ use parity_scale_codec::Encode;
 use serial_test::serial;
 use sp_core::Pair;
 use sp_keyring::AccountKeyring;
+use std::collections::HashSet;
 use subxt::utils::AccountId32;
 use synedrion::k256::ecdsa::VerifyingKey;
 
@@ -136,6 +137,20 @@ async fn test_reshare() {
         assert_eq!(response_result.unwrap().text().await.unwrap(), "");
     }
 
+    // Now wait until signers have changed
+    let old_signer_ids = HashSet::from_iter(signer_stash_accounts.into_iter().map(|id| id.0));
+    loop {
+        let new_signer_ids: HashSet<[u8; 32]> = {
+            let signer_query = entropy::storage().staking_extension().signers();
+            let signer_ids = query_chain(&api, &rpc, signer_query, None).await.unwrap().unwrap();
+            HashSet::from_iter(signer_ids.into_iter().map(|id| id.0))
+        };
+        if new_signer_ids != old_signer_ids {
+            break;
+        }
+        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+    }
+
     for (tss_account, key_share_and_aux_before) in key_shares_before.iter() {
         let (key_share_before, aux_info_before): KeyShareWithAuxInfo =
             deserialize(key_share_and_aux_before).unwrap();

diff --git a/crates/threshold-signature-server/tests/register_sign_reshare_sign.rs b/crates/threshold-signature-server/tests/register_sign_reshare_sign.rs
@@ -181,37 +181,33 @@ async fn do_reshare(api: &OnlineClient<EntropyConfig>, rpc: &LegacyRpcMethods<En
         assert_eq!(response_result.unwrap().text().await.unwrap(), "");
     }
 
-    let new_signers = {
-        let signer_query = entropy::storage().staking_extension().signers();
-        let signer_ids = query_chain(&api, &rpc, signer_query, None).await.unwrap().unwrap();
-        let mut signers = Vec::new();
-        for signer in signer_ids {
-            let query = entropy::storage().staking_extension().threshold_servers(signer);
-            let server_info = query_chain(&api, &rpc, query, None).await.unwrap().unwrap();
-            signers.push(server_info);
+    // Wait for the reshare protocol to finish
+    let old_signer_ids = HashSet::from_iter(signer_stash_accounts.into_iter().map(|id| id.0));
+    let new_signer_ids = loop {
+        let new_signer_ids: HashSet<[u8; 32]> = {
+            let signer_query = entropy::storage().staking_extension().signers();
+            let signer_ids = query_chain(&api, &rpc, signer_query, None).await.unwrap().unwrap();
+            HashSet::from_iter(signer_ids.into_iter().map(|id| id.0))
+        };
+        if new_signer_ids != old_signer_ids {
+            break new_signer_ids;
         }
-        signers
+        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
     };
 
     // Tell TS servers who do not have an associated chain node to rotate their keyshare.
     // This is called by the chain on getting confirmation of the reshare from all of the new
     // signing group.
-    for signer in new_signers {
+    for signer in new_signer_ids {
+        let query = entropy::storage().staking_extension().threshold_servers(&AccountId32(signer));
+        let server_info = query_chain(&api, &rpc, query, None).await.unwrap().unwrap();
         let _ = client
             .post(format!(
                 "http://{}/rotate_network_key",
-                std::str::from_utf8(&signer.endpoint).unwrap()
+                std::str::from_utf8(&server_info.endpoint).unwrap()
             ))
             .send()
             .await
             .unwrap();
     }
-
-    // Check that the signers have changed since before the reshare
-    let signer_query = entropy::storage().staking_extension().signers();
-    let new_signer_stash_accounts =
-        query_chain(&api, &rpc, signer_query, None).await.unwrap().unwrap();
-    let old: HashSet<[u8; 32]> = signer_stash_accounts.iter().map(|s| s.0).collect();
-    let new: HashSet<[u8; 32]> = new_signer_stash_accounts.iter().map(|s| s.0).collect();
-    assert_ne!(old, new);
 }
diff --git a/node/cli/src/service.rs b/node/cli/src/service.rs
@@ -369,7 +369,7 @@ pub fn new_full_base(
             );
             offchain_db.local_storage_set(
                 sp_core::offchain::StorageKind::PERSISTENT,
-                b"rotate_keyshares",
+                b"rotate_network_key",
                 &format!("{}/rotate_network_key", endpoint).into_bytes(),
             );
             offchain_db.local_storage_set(