Update test CLI for new registration and signing flows (#1008)

* First pass at new registration flow * Hack jumpstart during registration flow * Split regstration into two functions depending on old or new flow * Only do jumpstart if using the on-chain flow * Support signing for both old and new flows * Make `mnemonic_option` consistent across registration and signing * Clean up a few things in the register flow * Panic if we didn't get any successful registrations * Remove hacked jumpstart from registration flow * Move comment around * Add `CHANGELOG` entry * Only return single registration event * Move `get_registered_details` fully into client * Update error used in signing test
entropyxyz · Aug 20, 2024 · bba31a3 · bba31a3
1 parent b60e821
commit bba31a3
Show file tree

Hide file tree

Showing 8 changed files with 146 additions and 87 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -26,6 +26,7 @@ At the moment this project **does not** adhere to
 - Signing flow with derived accounts ([#990](https://github.com/entropyxyz/entropy-core/pull/990))
 - TSS attestation endpoint ([#1001](https://github.com/entropyxyz/entropy-core/pull/1001))
 - Add `network-jumpstart` command to `entropy-test-cli` ([#1004](https://github.com/entropyxyz/entropy-core/pull/1004))
+- Update test CLI for new registration and signing flows ([#1008](https://github.com/entropyxyz/entropy-core/pull/1008))
 
 ### Changed
 - Fix TSS `AccountId` keys in chainspec ([#993](https://github.com/entropyxyz/entropy-core/pull/993))

diff --git a/crates/client/src/client.rs b/crates/client/src/client.rs
@@ -37,7 +37,7 @@ use crate::{
         EntropyConfig,
     },
     client::entropy::staking_extension::events::{EndpointChanged, ThresholdAccountChanged},
-    substrate::{query_chain, submit_transaction_with_pair},
+    substrate::{get_registered_details, query_chain, submit_transaction_with_pair},
     user::{get_signers_from_chain, UserSignatureRequest},
     Hasher,
 };
@@ -76,39 +76,33 @@ pub async fn register(
     signature_request_keypair: sr25519::Pair,
     program_account: SubxtAccountId32,
     programs_data: BoundedVec<ProgramInstance>,
+    on_chain: bool,
 ) -> Result<([u8; VERIFYING_KEY_LENGTH], RegisteredInfo), ClientError> {
-    // Send register transaction
-    put_register_request_on_chain(
-        api,
-        rpc,
-        signature_request_keypair.clone(),
-        program_account,
-        programs_data,
-    )
-    .await?;
+    let registration_event = if on_chain {
+        put_register_request_on_chain(
+            api,
+            rpc,
+            signature_request_keypair.clone(),
+            program_account,
+            programs_data,
+        )
+        .await?
+    } else {
+        put_old_register_request_on_chain(
+            api,
+            rpc,
+            signature_request_keypair.clone(),
+            program_account,
+            programs_data,
+        )
+        .await?
+    };
 
-    let account_id: SubxtAccountId32 = signature_request_keypair.public().into();
+    let verifying_key = registration_event.1 .0;
+    let registered_info = get_registered_details(api, rpc, verifying_key.clone()).await?;
+    let verifying_key = verifying_key.try_into().map_err(|_| ClientError::BadVerifyingKeyLength)?;
 
-    for _ in 0..50 {
-        let block_hash = rpc.chain_get_block_hash(None).await?;
-        let events =
-            EventsClient::new(api.clone()).at(block_hash.ok_or(ClientError::BlockHash)?).await?;
-        let registered_event = events.find::<entropy::registry::events::AccountRegistered>();
-        for event in registered_event.flatten() {
-            // check if the event belongs to this user
-            if event.0 == account_id {
-                let registered_query = entropy::storage().registry().registered(&event.1);
-                let registered_status = query_chain(api, rpc, registered_query, block_hash).await?;
-                if let Some(status) = registered_status {
-                    let verifying_key =
-                        event.1 .0.try_into().map_err(|_| ClientError::BadVerifyingKeyLength)?;
-                    return Ok((verifying_key, status));
-                }
-            }
-        }
-        std::thread::sleep(std::time::Duration::from_millis(1000));
-    }
-    Err(ClientError::RegistrationTimeout)
+    Ok((verifying_key, registered_info))
 }
 
 /// Request to sign a message
@@ -131,7 +125,12 @@ pub async fn sign(
     auxilary_data: Option<Vec<u8>>,
 ) -> Result<RecoverableSignature, ClientError> {
     let message_hash = Hasher::keccak(&message);
-    let validators_info = get_signers_from_chain(api, rpc, false).await?;
+
+    let registered_info =
+        get_registered_details(api, rpc, signature_verifying_key.to_vec()).await?;
+    let with_parent_key = registered_info.derivation_path.is_some();
+    let validators_info = get_signers_from_chain(api, rpc, with_parent_key).await?;
+
     tracing::debug!("Validators info {:?}", validators_info);
     let block_number = rpc.chain_get_header(None).await?.ok_or(ClientError::BlockNumber)?.number;
     let signature_request = UserSignatureRequest {
@@ -289,19 +288,75 @@ pub async fn get_programs(
     Ok(programs)
 }
 
-/// Submit a register transaction
+/// Submits a transaction registering an account on-chain.
+#[tracing::instrument(
+    skip_all,
+    fields(
+        user_account = ?signature_request_keypair.public(),
+    )
+)]
 pub async fn put_register_request_on_chain(
     api: &OnlineClient<EntropyConfig>,
     rpc: &LegacyRpcMethods<EntropyConfig>,
     signature_request_keypair: sr25519::Pair,
     deployer: SubxtAccountId32,
     program_instances: BoundedVec<ProgramInstance>,
-) -> Result<(), ClientError> {
-    let registering_tx = entropy::tx().registry().register(deployer, program_instances);
+) -> Result<entropy::registry::events::AccountRegistered, ClientError> {
+    tracing::debug!("Registering an account using on-chain flow.");
+
+    let registering_tx = entropy::tx().registry().register_on_chain(deployer, program_instances);
+    let registered_events =
+        submit_transaction_with_pair(api, rpc, &signature_request_keypair, &registering_tx, None)
+            .await?;
+
+    // Note: In the case of the new registration flow we can have many registration events for a
+    // single signature request account. We only care about the first one we find.
+    let registered_event = registered_events
+        .find::<entropy::registry::events::AccountRegistered>()
+        .flatten()
+        .find_map(|event| (event.0 == signature_request_keypair.public().into()).then_some(event))
+        .ok_or(ClientError::NotRegistered);
+
+    registered_event
+}
 
+/// Submits a transaction registering an account on-chain using the old off-chain flow.
+#[tracing::instrument(
+    skip_all,
+    fields(
+        user_account = ?signature_request_keypair.public(),
+    )
+)]
+pub async fn put_old_register_request_on_chain(
+    api: &OnlineClient<EntropyConfig>,
+    rpc: &LegacyRpcMethods<EntropyConfig>,
+    signature_request_keypair: sr25519::Pair,
+    deployer: SubxtAccountId32,
+    program_instances: BoundedVec<ProgramInstance>,
+) -> Result<entropy::registry::events::AccountRegistered, ClientError> {
+    tracing::debug!("Registering an account using old off-chain flow.");
+
+    let registering_tx = entropy::tx().registry().register(deployer, program_instances);
     submit_transaction_with_pair(api, rpc, &signature_request_keypair, &registering_tx, None)
         .await?;
-    Ok(())
+
+    let account_id: SubxtAccountId32 = signature_request_keypair.public().into();
+
+    for _ in 0..50 {
+        let block_hash = rpc.chain_get_block_hash(None).await?;
+        let events =
+            EventsClient::new(api.clone()).at(block_hash.ok_or(ClientError::BlockHash)?).await?;
+        let registered_event = events.find::<entropy::registry::events::AccountRegistered>();
+        for event in registered_event.flatten() {
+            // check if the event belongs to this user
+            if event.0 == account_id {
+                return Ok(event);
+            }
+        }
+        std::thread::sleep(std::time::Duration::from_millis(1000));
+    }
+
+    Err(ClientError::RegistrationTimeout)
 }
 
 /// Check that the verfiying key from a new signature matches that in the from the
@@ -318,6 +373,7 @@ pub async fn check_verifying_key(
         entropy::storage().registry().registered(BoundedVec(verifying_key_serialized));
     let query_registered_status = query_chain(api, rpc, registered_query, None).await;
     query_registered_status?.ok_or(ClientError::NotRegistered)?;
+
     Ok(())
 }
 

diff --git a/crates/client/src/errors.rs b/crates/client/src/errors.rs
@@ -26,6 +26,8 @@ pub enum SubstrateError {
     GenericSubstrate(#[from] subxt::error::Error),
     #[error("Could not sumbit transaction {0}")]
     BadEvent(String),
+    #[error("User is not registered on-chain")]
+    NotRegistered,
 }
 
 /// An error on getting the current subgroup signers
@@ -92,7 +94,7 @@ pub enum ClientError {
     BadRecoveryId,
     #[error("Cannot parse chain query response: {0}")]
     TryFromSlice(#[from] std::array::TryFromSliceError),
-    #[error("User not registered")]
+    #[error("User is not registered on-chain")]
     NotRegistered,
     #[error("No synced validators")]
     NoSyncedValidators,

diff --git a/crates/client/src/substrate.rs b/crates/client/src/substrate.rs
@@ -13,7 +13,10 @@
 // You should have received a copy of the GNU Affero General Public License
 // along with this program.  If not, see <https://www.gnu.org/licenses/>.
 //! For interacting with the substrate chain node
+use crate::chain_api::entropy::runtime_types::bounded_collections::bounded_vec::BoundedVec;
+use crate::chain_api::entropy::runtime_types::pallet_registry::pallet::RegisteredInfo;
 use crate::chain_api::{entropy, EntropyConfig};
+
 use entropy_shared::MORTALITY_BLOCKS;
 use sp_core::{sr25519, Pair};
 use subxt::{
@@ -107,6 +110,38 @@ where
     Ok(result)
 }
 
+/// Returns a registered user's key visibility
+#[tracing::instrument(skip_all, fields(verifying_key))]
+pub async fn get_registered_details(
+    api: &OnlineClient<EntropyConfig>,
+    rpc: &LegacyRpcMethods<EntropyConfig>,
+    verifying_key: Vec<u8>,
+) -> Result<RegisteredInfo, SubstrateError> {
+    tracing::info!("Querying chain for registration info.");
+
+    let registered_info_query =
+        entropy::storage().registry().registered(BoundedVec(verifying_key.clone()));
+    let registered_result = query_chain(api, rpc, registered_info_query, None).await?;
+
+    let registration_info = if let Some(old_registration_info) = registered_result {
+        tracing::debug!("Found user in old `Registered` struct.");
+
+        old_registration_info
+    } else {
+        // We failed with the old registration path, let's try the new one
+        tracing::warn!("Didn't find user in old `Registered` struct, trying new one.");
+
+        let registered_info_query =
+            entropy::storage().registry().registered_on_chain(BoundedVec(verifying_key));
+
+        query_chain(api, rpc, registered_info_query, None)
+            .await?
+            .ok_or_else(|| SubstrateError::NotRegistered)?
+    };
+
+    Ok(registration_info)
+}
+
 /// A wrapper around [sr25519::Pair] which implements [Signer]
 /// This is needed because on wasm we cannot use the generic `subxt::tx::PairSigner`
 #[derive(Clone)]

diff --git a/crates/test-cli/src/lib.rs b/crates/test-cli/src/lib.rs
@@ -79,6 +79,9 @@ enum CliCommand {
         /// If giving a mnemonic it must be enclosed in quotes, eg: "--mnemonic-option "alarm mutual concert...""
         #[arg(short, long)]
         mnemonic_option: Option<String>,
+        /// Indicates that a user wants to register using the fully on-chain registration flow.
+        #[arg(long)]
+        on_chain: bool,
     },
     /// Ask the network to sign a given message
     Sign {
@@ -89,6 +92,7 @@ enum CliCommand {
         /// Optional auxiliary data passed to the program, given as hex
         auxilary_data: Option<String>,
         /// The mnemonic to use for the call
+        #[arg(short, long)]
         mnemonic_option: Option<String>,
     },
     /// Update the program for a particular account
@@ -173,7 +177,7 @@ pub async fn run_command(
     let rpc = get_rpc(&endpoint_addr).await?;
 
     match cli.command {
-        CliCommand::Register { mnemonic_option, programs } => {
+        CliCommand::Register { mnemonic_option, programs, on_chain } => {
             let mnemonic = if let Some(mnemonic_option) = mnemonic_option {
                 mnemonic_option
             } else {
@@ -198,10 +202,11 @@ pub async fn run_command(
                 program_keypair.clone(),
                 program_account,
                 BoundedVec(programs_info),
+                on_chain,
             )
             .await?;
 
-            Ok(format!("Verfiying key: {},\n{:?}", hex::encode(verifying_key), registered_info))
+            Ok(format!("Verifying key: {},\n{:?}", hex::encode(verifying_key), registered_info))
         },
         CliCommand::Sign { signature_verifying_key, message, auxilary_data, mnemonic_option } => {
             let mnemonic = if let Some(mnemonic_option) = mnemonic_option {

diff --git a/crates/threshold-signature-server/src/helpers/substrate.rs b/crates/threshold-signature-server/src/helpers/substrate.rs
@@ -20,7 +20,6 @@ use crate::{
             self,
             runtime_types::{
                 bounded_collections::bounded_vec::BoundedVec, pallet_programs::pallet::ProgramInfo,
-                pallet_registry::pallet::RegisteredInfo,
             },
         },
         EntropyConfig,
@@ -72,38 +71,6 @@ pub async fn get_oracle_data(
     Ok(oracle_info.0)
 }
 
-/// Returns a registered user's key visibility
-#[tracing::instrument(skip_all, fields(verifying_key))]
-pub async fn get_registered_details(
-    api: &OnlineClient<EntropyConfig>,
-    rpc: &LegacyRpcMethods<EntropyConfig>,
-    verifying_key: Vec<u8>,
-) -> Result<RegisteredInfo, UserErr> {
-    tracing::info!("Querying chain for registration info.");
-
-    let registered_info_query =
-        entropy::storage().registry().registered(BoundedVec(verifying_key.clone()));
-    let registered_result = query_chain(api, rpc, registered_info_query, None).await?;
-
-    let registration_info = if let Some(old_registration_info) = registered_result {
-        tracing::debug!("Found user in old `Registered` struct.");
-
-        old_registration_info
-    } else {
-        // We failed with the old registration path, let's try the new one
-        tracing::warn!("Didn't find user in old `Registered` struct, trying new one.");
-
-        let registered_info_query =
-            entropy::storage().registry().registered_on_chain(BoundedVec(verifying_key));
-
-        query_chain(api, rpc, registered_info_query, None)
-            .await?
-            .ok_or_else(|| UserErr::ChainFetch("Not Registering error: Register Onchain first"))?
-    };
-
-    Ok(registration_info)
-}
-
 /// Takes Stash keys and returns validator info from chain
 pub async fn get_validators_info(
     api: &OnlineClient<EntropyConfig>,

diff --git a/crates/threshold-signature-server/src/user/api.rs b/crates/threshold-signature-server/src/user/api.rs
@@ -26,6 +26,7 @@ use axum::{
 use base64::prelude::{Engine, BASE64_STANDARD};
 use bip39::{Language, Mnemonic};
 use blake2::{Blake2s256, Digest};
+use entropy_client::substrate::get_registered_details;
 use entropy_kvdb::kv_manager::{
     error::{InnerKvError, KvError},
     helpers::serialize as key_serialize,
@@ -68,8 +69,7 @@ use crate::{
         launch::LATEST_BLOCK_NUMBER_NEW_USER,
         signing::{do_signing, Hasher},
         substrate::{
-            get_oracle_data, get_program, get_registered_details, get_stash_address, query_chain,
-            submit_transaction,
+            get_oracle_data, get_program, get_stash_address, query_chain, submit_transaction,
         },
         user::{check_in_registration_group, compute_hash, do_dkg},
         validator::{get_signer, get_signer_and_x25519_secret},