clean

entropyxyz · Aug 12, 2024 · 397cf09 · 397cf09
1 parent e1c8ed3
commit 397cf09
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 80 deletions.
diff --git a/crates/protocol/src/execute_protocol.rs b/crates/protocol/src/execute_protocol.rs
@@ -73,8 +73,7 @@ pub async fn execute_protocol_generic<Res: synedrion::ProtocolResult>(
     mut chans: Channels,
     session: Session<Res, sr25519::Signature, PairWrapper, PartyId>,
     session_id_hash: [u8; 32],
-) -> Result<(Res::Success, Broadcaster, mpsc::Receiver<ProtocolMessage>), GenericProtocolError<Res>>
-{
+) -> Result<(Res::Success, Channels), GenericProtocolError<Res>> {
     let session_id = synedrion::SessionId::from_seed(&session_id_hash);
     let tx = &chans.0;
     let rx = &mut chans.1;
@@ -144,7 +143,7 @@ pub async fn execute_protocol_generic<Res: synedrion::ProtocolResult>(
         }
 
         match session.finalize_round(&mut OsRng, accum)? {
-            FinalizeOutcome::Success(res) => break Ok((res, chans.0, chans.1)),
+            FinalizeOutcome::Success(res) => break Ok((res, chans)),
             FinalizeOutcome::AnotherRound {
                 session: new_session,
                 cached_messages: new_cached_messages,
@@ -231,12 +230,10 @@ pub async fn execute_dkg(
         )
         .map_err(ProtocolExecutionErr::SessionCreation)?;
 
-        let (init_keyshare, broadcaster, rx) =
+        let (init_keyshare, chans) =
             execute_protocol_generic(chans, session, session_id_hash).await?;
 
         tracing::info!("Finished key init protocol");
-        // Setup channels for the next session
-        let chans = Channels(broadcaster.clone(), rx);
 
         // Send verifying key
         let verifying_key = init_keyshare.verifying_key();
@@ -304,15 +301,12 @@ pub async fn execute_dkg(
         inputs,
     )
     .map_err(ProtocolExecutionErr::SessionCreation)?;
-    let (new_key_share_option, broadcaster, rx) =
+    let (new_key_share_option, chans) =
         execute_protocol_generic(chans, session, session_id_hash).await?;
     let new_key_share =
         new_key_share_option.ok_or(ProtocolExecutionErr::NoOutputFromReshareProtocol)?;
     tracing::info!("Finished reshare protocol");
 
-    // Setup channels for the next session
-    let chans = Channels(broadcaster.clone(), rx);
-
     // Now run the aux gen protocol to get AuxInfo
     let session_id_hash = session_id.blake2(Some(Subsession::AuxGen))?;
     let session = make_aux_gen_session(
@@ -335,14 +329,15 @@ pub async fn execute_dkg(
     fields(threshold_accounts, my_idx),
     level = tracing::Level::DEBUG
 )]
-pub async fn execute_proactive_refresh(
+pub async fn execute_reshare(
     session_id: SessionId,
     chans: Channels,
     threshold_pair: &sr25519::Pair,
     threshold_accounts: Vec<AccountId32>,
     inputs: KeyResharingInputs<KeyParams, PartyId>,
+    aux_info_option: Option<AuxInfo<KeyParams, PartyId>>,
 ) -> Result<
-    (ThresholdKeyShare<KeyParams, PartyId>, Broadcaster, mpsc::Receiver<ProtocolMessage>),
+    (ThresholdKeyShare<KeyParams, PartyId>, AuxInfo<KeyParams, PartyId>),
     ProtocolExecutionErr,
 > {
     tracing::debug!("Executing proactive refresh");
@@ -363,10 +358,25 @@ pub async fn execute_proactive_refresh(
     )
     .map_err(ProtocolExecutionErr::SessionCreation)?;
 
-    let (new_key_share, brodcaster, rx) =
-        execute_protocol_generic(chans, session, session_id_hash).await?;
+    let (new_key_share, chans) = execute_protocol_generic(chans, session, session_id_hash).await?;
+
+    let aux_info = if let Some(aux_info) = aux_info_option {
+        aux_info
+    } else {
+        // Now run an aux gen session
+        let session_id_hash_aux_data = session_id.blake2(Some(Subsession::AuxGen))?;
+        let session = make_aux_gen_session(
+            &mut OsRng,
+            SynedrionSessionId::from_seed(session_id_hash_aux_data.as_slice()),
+            PairWrapper(threshold_pair.clone()),
+            &party_ids,
+        )
+        .map_err(ProtocolExecutionErr::SessionCreation)?;
+
+        execute_protocol_generic(chans, session, session_id_hash).await?.0
+    };
 
-    Ok((new_key_share.ok_or(ProtocolExecutionErr::NoOutputFromReshareProtocol)?, brodcaster, rx))
+    Ok((new_key_share.ok_or(ProtocolExecutionErr::NoOutputFromReshareProtocol)?, aux_info))
 }
 
 /// Psuedo-randomly select a subset of the parties of size `threshold`

diff --git a/crates/protocol/tests/helpers/mod.rs b/crates/protocol/tests/helpers/mod.rs
@@ -16,9 +16,7 @@
 //! A simple protocol server, like a mini version of entropy-tss, for benchmarking
 use anyhow::{anyhow, ensure};
 use entropy_protocol::{
-    execute_protocol::{
-        execute_dkg, execute_proactive_refresh, execute_signing_protocol, Channels,
-    },
+    execute_protocol::{execute_dkg, execute_reshare, execute_signing_protocol, Channels},
     protocol_transport::{
         errors::WsError,
         noise::{noise_handshake_initiator, noise_handshake_responder},
@@ -147,8 +145,7 @@ pub async fn server(
             };
 
             let new_keyshare =
-                execute_proactive_refresh(session_id, channels, &pair, tss_accounts, inputs)
-                    .await?;
+                execute_reshare(session_id, channels, &pair, tss_accounts, inputs, None).await?;
             Ok(ProtocolOutput::Reshare(new_keyshare.0))
         },
         SessionId::Dkg { .. } => {

diff --git a/crates/threshold-signature-server/src/signing_client/api.rs b/crates/threshold-signature-server/src/signing_client/api.rs
@@ -24,13 +24,11 @@ use axum::{
 };
 use blake2::{Blake2s256, Digest};
 use entropy_protocol::{
-    execute_protocol::{execute_proactive_refresh, Channels},
-    protocol_transport::Broadcaster,
-    KeyParams, KeyShareWithAuxInfo, Listener, PartyId, ProtocolMessage, SessionId, ValidatorInfo,
+    execute_protocol::{execute_reshare, Channels},
+    KeyParams, Listener, PartyId, SessionId, ValidatorInfo,
 };
 use parity_scale_codec::Encode;
 use std::{collections::BTreeSet, time::Duration};
-use tokio::sync::mpsc;
 
 use entropy_kvdb::kv_manager::{
     helpers::{deserialize, serialize as key_serialize},
@@ -97,33 +95,24 @@ pub async fn proactive_refresh(
         let exists_result = app_state.kv_store.kv().exists(&key).await?;
         if exists_result {
             let old_key_share = app_state.kv_store.kv().get(&key).await?;
-            let (deserialized_old_key, _aux_info): (
+            let (deserialized_old_key, aux_info): (
                 ThresholdKeyShare<KeyParams, PartyId>,
                 AuxInfo<KeyParams, PartyId>,
             ) = deserialize(&old_key_share)
                 .ok_or_else(|| ProtocolErr::Deserialization("Failed to load KeyShare".into()))?;
 
-            let (new_key_share, _broadcaster, _rx) = do_proactive_refresh(
+            let (new_key_share, aux_info) = do_proactive_refresh(
                 &ocw_data.validators_info,
                 &signer,
                 &x25519_secret_key,
                 &app_state.listener_state,
                 encoded_key,
                 deserialized_old_key,
                 ocw_data.block_number,
+                aux_info,
             )
             .await?;
 
-            // Get aux info from existing entry
-            let aux_info = {
-                let existing_entry = app_state.kv_store.kv().get(&key).await?;
-                let (_old_key_share, aux_info): KeyShareWithAuxInfo = deserialize(&existing_entry)
-                    .ok_or_else(|| {
-                        ProtocolErr::Deserialization("Failed to load KeyShare".into())
-                    })?;
-                aux_info
-            };
-
             // Since this is a refresh with the parties not changing, store the old aux_info
             let serialized_key_share = key_serialize(&(new_key_share, aux_info))
                 .map_err(|_| ProtocolErr::KvSerialize("Kv Serialize Error".to_string()))?;
@@ -153,7 +142,7 @@ async fn handle_socket_result(socket: WebSocket, app_state: AppState) {
     };
 }
 
-#[allow(clippy::type_complexity)]
+#[allow(clippy::type_complexity, clippy::too_many_arguments)]
 #[tracing::instrument(
     skip_all,
     fields(validators_info, verifying_key, my_subgroup),
@@ -167,10 +156,8 @@ pub async fn do_proactive_refresh(
     verifying_key: Vec<u8>,
     old_key: ThresholdKeyShare<KeyParams, PartyId>,
     block_number: u32,
-) -> Result<
-    (ThresholdKeyShare<KeyParams, PartyId>, Broadcaster, mpsc::Receiver<ProtocolMessage>),
-    ProtocolErr,
-> {
+    aux_info: AuxInfo<KeyParams, PartyId>,
+) -> Result<(ThresholdKeyShare<KeyParams, PartyId>, AuxInfo<KeyParams, PartyId>), ProtocolErr> {
     tracing::debug!("Preparing to perform proactive refresh");
     tracing::debug!("Signing with {:?}", &signer.signer().public());
 
@@ -217,9 +204,15 @@ pub async fn do_proactive_refresh(
     )
     .await?;
 
-    let result =
-        execute_proactive_refresh(session_id, channels, signer.signer(), tss_accounts, inputs)
-            .await?;
+    let result = execute_reshare(
+        session_id,
+        channels,
+        signer.signer(),
+        tss_accounts,
+        inputs,
+        Some(aux_info),
+    )
+    .await?;
     Ok(result)
 }
 

diff --git a/crates/threshold-signature-server/src/validator/api.rs b/crates/threshold-signature-server/src/validator/api.rs
@@ -29,28 +29,21 @@ use crate::{
 };
 use axum::{body::Bytes, extract::State, http::StatusCode};
 use entropy_kvdb::kv_manager::{helpers::serialize as key_serialize, KvManager};
-use entropy_protocol::Subsession;
 pub use entropy_protocol::{
     decode_verifying_key,
     errors::ProtocolExecutionErr,
-    execute_protocol::{
-        execute_proactive_refresh, execute_protocol_generic, Channels, PairWrapper,
-    },
+    execute_protocol::{execute_protocol_generic, execute_reshare, Channels, PairWrapper},
     KeyParams, KeyShareWithAuxInfo, Listener, PartyId, SessionId, ValidatorInfo,
 };
 use entropy_shared::{OcwMessageReshare, NETWORK_PARENT_KEY};
 use parity_scale_codec::{Decode, Encode};
-use rand_core::OsRng;
 use sp_core::Pair;
 use std::{collections::BTreeSet, str::FromStr};
 use subxt::{
     backend::legacy::LegacyRpcMethods, ext::sp_core::sr25519, tx::PairSigner, utils::AccountId32,
     OnlineClient,
 };
-use synedrion::{
-    make_aux_gen_session, sessions::SessionId as SynedrionSessionId, AuxInfo, KeyResharingInputs,
-    NewHolder, OldHolder,
-};
+use synedrion::{KeyResharingInputs, NewHolder, OldHolder};
 
 /// HTTP POST endpoint called by the off-chain worker (propagation pallet) during network reshare.
 ///
@@ -175,33 +168,9 @@ pub async fn new_reshare(
     )
     .await?;
 
-    let (new_key_share, brodcaster, rx) = execute_proactive_refresh(
-        session_id.clone(),
-        channels,
-        signer.signer(),
-        tss_accounts,
-        inputs,
-    )
-    .await?;
-
-    // Setup channels for the next session
-    let channels = Channels(brodcaster, rx);
-
-    // Now run an aux gen session
-    let session_id_hash = session_id.blake2(Some(Subsession::AuxGen))?;
-    let session = make_aux_gen_session(
-        &mut OsRng,
-        SynedrionSessionId::from_seed(session_id_hash.as_slice()),
-        PairWrapper(signer.signer().clone()),
-        &party_ids,
-    )
-    .map_err(ProtocolExecutionErr::SessionCreation)?;
-
-    let aux_info: AuxInfo<KeyParams, PartyId> =
-        execute_protocol_generic(channels, session, session_id_hash)
-            .await
-            .map_err(|_| ValidatorErr::ProtocolError("Error executing protocol".to_string()))?
-            .0;
+    let (new_key_share, aux_info) =
+        execute_reshare(session_id.clone(), channels, signer.signer(), tss_accounts, inputs, None)
+            .await?;
 
     let serialized_key_share = key_serialize(&(new_key_share, aux_info))
         .map_err(|_| ProtocolErr::KvSerialize("Kv Serialize Error".to_string()))?;