Test CLI command to retrieve quote and change endpoint / TSS account …

…in one command (#1198) * Add get_quote_and_change_endpoint client fn, error handling * Cli command to change endpoint no longer takes quote * Doccomments * Clippy * Add context when calling attest endpoint * Error handle http response in GetTdxQuote command * Improve error handling when getting quotes * Use display not debug for building quote context querystring * Refactor handling of mnemonics in test-cli * Update change TSS endpoints client function and CLI command * get-tdx-quote command now takes quote context as an arguement * Changelog * Use From not Into
entropyxyz · Dec 11, 2024 · 6d3351e · 6d3351e
1 parent cfcb367
commit 6d3351e
Show file tree

Hide file tree

Showing 6 changed files with 166 additions and 101 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -34,6 +34,7 @@ runtime
 - Protocol message versioning ([#1140](https://github.com/entropyxyz/entropy-core/pull/1140))
 - CLI command to get oracle headings ([#1170](https://github.com/entropyxyz/entropy-core/pull/1170))
 - Add TSS endpoint to get TDX quote ([#1173](https://github.com/entropyxyz/entropy-core/pull/1173))
+- Test CLI command to retrieve quote and change endpoint / TSS account in one command ([#1198](https://github.com/entropyxyz/entropy-core/pull/1198))
 
 ### Changed
 - Use correct key rotation endpoint in OCW ([#1104](https://github.com/entropyxyz/entropy-core/pull/1104))

diff --git a/crates/client/src/client.rs b/crates/client/src/client.rs
@@ -17,13 +17,12 @@
 //! Used in integration tests and for the test-cli
 pub use crate::{
     chain_api::{get_api, get_rpc},
-    errors::ClientError,
+    errors::{ClientError, SubstrateError},
 };
-use anyhow::anyhow;
 pub use entropy_protocol::{sign_and_encrypt::EncryptedSignedMessage, KeyParams};
+pub use entropy_shared::{HashingAlgorithm, QuoteContext};
 use parity_scale_codec::Decode;
 use rand::Rng;
-use std::str::FromStr;
 pub use synedrion::KeyShare;
 
 use crate::{
@@ -39,7 +38,7 @@ use crate::{
         EntropyConfig,
     },
     client::entropy::staking_extension::events::{EndpointChanged, ThresholdAccountChanged},
-    substrate::{get_registered_details, submit_transaction_with_pair},
+    substrate::{get_registered_details, query_chain, submit_transaction_with_pair},
     user::{
         self, get_all_signers_from_chain, get_validators_not_signer_for_relay, UserSignatureRequest,
     },
@@ -48,7 +47,6 @@ use crate::{
 
 use base64::prelude::{Engine, BASE64_STANDARD};
 use entropy_protocol::RecoverableSignature;
-use entropy_shared::HashingAlgorithm;
 use futures::stream::StreamExt;
 use sp_core::{
     sr25519::{self, Signature},
@@ -338,50 +336,90 @@ pub async fn put_register_request_on_chain(
     registered_event
 }
 
-/// Changes the endpoint of a validator
+/// Changes the endpoint of a validator, retrieving a TDX quote from the new endpoint internally
+pub async fn get_quote_and_change_endpoint(
+    api: &OnlineClient<EntropyConfig>,
+    rpc: &LegacyRpcMethods<EntropyConfig>,
+    validator_keypair: sr25519::Pair,
+    new_endpoint: String,
+) -> Result<EndpointChanged, ClientError> {
+    let quote = get_tdx_quote(&new_endpoint, QuoteContext::ChangeEndpoint).await?;
+    change_endpoint(api, rpc, validator_keypair, new_endpoint, quote).await
+}
+
+/// Changes the endpoint of a validator, with a TDX quote given as an argument
 pub async fn change_endpoint(
     api: &OnlineClient<EntropyConfig>,
     rpc: &LegacyRpcMethods<EntropyConfig>,
     user_keypair: sr25519::Pair,
     new_endpoint: String,
     quote: Vec<u8>,
-) -> anyhow::Result<EndpointChanged> {
+) -> Result<EndpointChanged, ClientError> {
     let change_endpoint_tx =
         entropy::tx().staking_extension().change_endpoint(new_endpoint.into(), quote);
     let in_block =
         submit_transaction_with_pair(api, rpc, &user_keypair, &change_endpoint_tx, None).await?;
     let result_event = in_block
         .find_first::<entropy::staking_extension::events::EndpointChanged>()?
-        .ok_or(anyhow!("Error with transaction"))?;
+        .ok_or(SubstrateError::NoEvent)?;
     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,
-    new_tss_account: String,
-    new_x25519_public_key: String,
+    validator_keypair: sr25519::Pair,
+    new_tss_account: SubxtAccountId32,
+    new_x25519_public_key: [u8; 32],
     new_pck_certificate_chain: Vec<Vec<u8>>,
     quote: Vec<u8>,
-) -> anyhow::Result<ThresholdAccountChanged> {
-    let tss_account = SubxtAccountId32::from_str(&new_tss_account)?;
-    let x25519_public_key = hex::decode(new_x25519_public_key)?
-        .try_into()
-        .map_err(|_| anyhow!("X25519 pub key needs to be 32 bytes"))?;
+) -> Result<ThresholdAccountChanged, ClientError> {
     let change_threshold_accounts = entropy::tx().staking_extension().change_threshold_accounts(
-        tss_account,
-        x25519_public_key,
+        new_tss_account,
+        new_x25519_public_key,
         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(anyhow!("Error with transaction"))?;
+        .ok_or(SubstrateError::NoEvent)?;
     Ok(result_event)
 }
 
@@ -463,3 +501,31 @@ pub async fn get_oracle_headings(
     }
     Ok(headings)
 }
+
+/// Retrieve a TDX quote using the currently configured endpoint associated with the given validator
+/// ID
+pub async fn get_tdx_quote_with_validator_id(
+    api: &OnlineClient<EntropyConfig>,
+    rpc: &LegacyRpcMethods<EntropyConfig>,
+    validator_stash: &SubxtAccountId32,
+    quote_context: QuoteContext,
+) -> Result<Vec<u8>, ClientError> {
+    let query = entropy::storage().staking_extension().threshold_servers(validator_stash);
+    let server_info = query_chain(api, rpc, query, None).await?.ok_or(ClientError::NoServerInfo)?;
+
+    let tss_endpoint = std::str::from_utf8(&server_info.endpoint)?;
+    get_tdx_quote(tss_endpoint, quote_context).await
+}
+
+/// Retrieve a TDX quote with a given socket address
+pub async fn get_tdx_quote(
+    tss_endpoint: &str,
+    quote_context: QuoteContext,
+) -> Result<Vec<u8>, ClientError> {
+    let response =
+        reqwest::get(format!("http://{}/attest?context={}", tss_endpoint, quote_context)).await?;
+    if response.status() != reqwest::StatusCode::OK {
+        return Err(ClientError::QuoteGet(response.text().await?));
+    }
+    Ok(response.bytes().await?.to_vec())
+}
diff --git a/crates/client/src/errors.rs b/crates/client/src/errors.rs
@@ -123,4 +123,8 @@ pub enum ClientError {
     Codec(#[from] parity_scale_codec::Error),
     #[error("Attestation request: {0}")]
     AttestationRequest(#[from] AttestationRequestError),
+    #[error("Unable to get TDX quote: {0}")]
+    QuoteGet(String),
+    #[error("Unable to get info for TSS server from chain")]
+    NoServerInfo,
 }
diff --git a/crates/client/src/tests.rs b/crates/client/src/tests.rs
@@ -141,8 +141,8 @@ async fn test_change_threshold_accounts() {
         &api,
         &rpc,
         one.into(),
-        tss_public_key.to_string(),
-        hex::encode(*x25519_public_key.as_bytes()),
+        tss_public_key.into(),
+        *x25519_public_key.as_bytes(),
         pck_certificate_chain,
         quote,
     )

diff --git a/crates/shared/src/types.rs b/crates/shared/src/types.rs
@@ -146,6 +146,18 @@ pub enum QuoteContext {
     ChangeThresholdAccounts,
 }
 
+#[cfg(feature = "std")]
+impl std::fmt::Display for QuoteContext {
+    /// Custom display implementation so that it can be used to build a query string
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            QuoteContext::Validate => write!(f, "validate"),
+            QuoteContext::ChangeEndpoint => write!(f, "change_endpoint"),
+            QuoteContext::ChangeThresholdAccounts => write!(f, "change_threshold_accounts"),
+        }
+    }
+}
+
 /// A trait for types which can handle attestation requests.
 #[cfg(not(feature = "wasm"))]
 pub trait AttestationHandler<AccountId> {