Update test CLI for new registration and signing flows

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))

crates/client/src/client.rs

@@ -76,39 +76,40 @@ pub async fn register(
     signature_request_keypair: sr25519::Pair,
     program_account: SubxtAccountId32,
     programs_data: BoundedVec<ProgramInstance>,
-) -> 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?;
+    on_chain: bool,
+) -> Result<Vec<([u8; VERIFYING_KEY_LENGTH], RegisteredInfo)>, ClientError> {
+    let account_registration_events = 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 mut registration_info = vec![];
+    for event in account_registration_events {
+        let verifying_key = event.1 .0;
+        let registered_info = get_registered_details(api, rpc, verifying_key.clone()).await?;
 
-    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));
+        registration_info.push((
+            verifying_key.try_into().map_err(|_| ClientError::BadVerifyingKeyLength)?,
+            registered_info,
+        ))
     }
-    Err(ClientError::RegistrationTimeout)
+
+    Ok(registration_info)
 }
 
 /// Request to sign a message
@@ -131,7 +132,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 +295,112 @@ 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<Vec<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.
+    let registered_events: Vec<_> = registered_events
+        .find::<entropy::registry::events::AccountRegistered>()
+        .flatten()
+        .filter(|event| event.0 == signature_request_keypair.public().into())
+        .collect();
+
+    Ok(registered_events)
+}
+
+/// 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<Vec<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(vec![event]);
+            }
+        }
+        std::thread::sleep(std::time::Duration::from_millis(1000));
+    }
+
+    Err(ClientError::RegistrationTimeout)
+}
+
+/// Returns a registered user's key visibility
+///
+/// TODO (Nando): This was copied from `entropy-tss::helpers::substrate`
+///
+/// What's the best place for this? Ideally here, and then we import this into the entropy-tss, but
+/// we need the `full-client` feature enabled...
+#[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, ClientError> {
+    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(|| ClientError::NotRegistered)?
+    };
+
+    Ok(registration_info)
 }
 
 /// Check that the verfiying key from a new signature matches that in the from the

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 {
@@ -192,16 +196,25 @@ pub async fn run_command(
                 );
             }
 
-            let (verifying_key, registered_info) = register(
+            let registrations = register(
                 &api,
                 &rpc,
                 program_keypair.clone(),
                 program_account,
                 BoundedVec(programs_info),
+                on_chain,
             )
             .await?;
 
-            Ok(format!("Verfiying key: {},\n{:?}", hex::encode(verifying_key), registered_info))
+            if registrations.is_empty() {
+                panic!("Failed to register an account!")
+            }
+
+            // The CLI currently doesn't support sending multiple registration requests in a single
+            // go, so we just grab one event (presumably the only one) and log that.
+            let (verifying_key, registered_info) = &registrations[0];
+
+            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 {