From 568ca3384ad76e941030ca7e2aa7bbba11170c34 Mon Sep 17 00:00:00 2001
From: peg <ameba23@systemli.org>
Date: Mon, 9 Dec 2024 11:09:18 +0100
Subject: [PATCH] Update change TSS endpoints client function and CLI command

---
 crates/client/src/client.rs | 41 +++++++++++++++++++++++++++++++++----
 crates/test-cli/src/lib.rs  | 15 ++++----------
 2 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/crates/client/src/client.rs b/crates/client/src/client.rs
index 44ee4a881..93cf1a2a7 100644
--- a/crates/client/src/client.rs
+++ b/crates/client/src/client.rs
@@ -365,11 +365,39 @@ pub async fn change_endpoint(
     Ok(result_event)
 }
 
+/// Changes the threshold account info of a validator, retrieving a TDX quote from the new endpoint internally
+pub async fn get_quote_and_change_threshold_accounts(
+    api: &OnlineClient<EntropyConfig>,
+    rpc: &LegacyRpcMethods<EntropyConfig>,
+    validator_keypair: sr25519::Pair,
+    new_tss_account: SubxtAccountId32,
+    new_x25519_public_key: [u8; 32],
+    new_pck_certificate_chain: Vec<Vec<u8>>,
+) -> Result<ThresholdAccountChanged, ClientError> {
+    let quote = get_tdx_quote_with_validator_id(
+        api,
+        rpc,
+        &SubxtAccountId32(validator_keypair.public().0),
+        QuoteContext::ChangeThresholdAccounts,
+    )
+    .await?;
+    change_threshold_accounts(
+        api,
+        rpc,
+        validator_keypair,
+        new_tss_account,
+        new_x25519_public_key,
+        new_pck_certificate_chain,
+        quote,
+    )
+    .await
+}
+
 /// Changes the threshold account info of a validator
 pub async fn change_threshold_accounts(
     api: &OnlineClient<EntropyConfig>,
     rpc: &LegacyRpcMethods<EntropyConfig>,
-    user_keypair: sr25519::Pair,
+    validator_keypair: sr25519::Pair,
     new_tss_account: SubxtAccountId32,
     new_x25519_public_key: [u8; 32],
     new_pck_certificate_chain: Vec<Vec<u8>>,
@@ -381,9 +409,14 @@ pub async fn change_threshold_accounts(
         new_pck_certificate_chain,
         quote,
     );
-    let in_block =
-        submit_transaction_with_pair(api, rpc, &user_keypair, &change_threshold_accounts, None)
-            .await?;
+    let in_block = submit_transaction_with_pair(
+        api,
+        rpc,
+        &validator_keypair,
+        &change_threshold_accounts,
+        None,
+    )
+    .await?;
     let result_event = in_block
         .find_first::<entropy::staking_extension::events::ThresholdAccountChanged>()?
         .ok_or(SubstrateError::NoEvent)?;
diff --git a/crates/test-cli/src/lib.rs b/crates/test-cli/src/lib.rs
index 184b53dad..b972fc9dc 100644
--- a/crates/test-cli/src/lib.rs
+++ b/crates/test-cli/src/lib.rs
@@ -27,9 +27,9 @@ use entropy_client::{
         EntropyConfig,
     },
     client::{
-        change_threshold_accounts, get_accounts, get_api, get_oracle_headings, get_programs,
-        get_quote_and_change_endpoint, get_rpc, get_tdx_quote, jumpstart_network, register,
-        remove_program, sign, store_program, update_programs, VERIFYING_KEY_LENGTH,
+        get_accounts, get_api, get_oracle_headings, get_programs, get_quote_and_change_endpoint,
+        get_quote_and_change_threshold_accounts, get_rpc, get_tdx_quote, jumpstart_network,
+        register, remove_program, sign, store_program, update_programs, VERIFYING_KEY_LENGTH,
     },
 };
 pub use entropy_shared::{QuoteContext, PROGRAM_VERSION_NUMBER};
@@ -162,11 +162,6 @@ enum CliCommand {
         new_x25519_public_key: String,
         /// The new Provisioning Certification Key (PCK) certificate chain to be used for the TSS.
         new_pck_certificate_chain: Vec<String>,
-        /// The Intel TDX quote used to prove that this TSS is running on TDX hardware.
-        ///
-        /// The quote format is specified in:
-        /// https://download.01.org/intel-sgx/latest/dcap-latest/linux/docs/Intel_TDX_DCAP_Quoting_Library_API.pdf
-        quote: String,
         /// The mnemonic for the validator stash account to use for the call, should be stash address
         #[arg(short, long)]
         mnemonic_option: Option<String>,
@@ -485,7 +480,6 @@ pub async fn run_command(
             new_tss_account,
             new_x25519_public_key,
             new_pck_certificate_chain,
-            quote,
             mnemonic_option,
         } => {
             let user_keypair = handle_mnemonic(mnemonic_option)?;
@@ -497,14 +491,13 @@ pub async fn run_command(
                 .map_err(|_| anyhow!("X25519 pub key needs to be 32 bytes"))?;
             let new_pck_certificate_chain =
                 new_pck_certificate_chain.iter().cloned().map(|i| i.into()).collect::<_>();
-            let result_event = change_threshold_accounts(
+            let result_event = get_quote_and_change_threshold_accounts(
                 &api,
                 &rpc,
                 user_keypair,
                 new_tss_account,
                 new_x25519_public_key,
                 new_pck_certificate_chain,
-                quote.into(),
             )
             .await?;
             cli.log(format!("Event result: {:?}", result_event));