From c04fc3a272a0656d7479bb22c995c6720361b91f Mon Sep 17 00:00:00 2001 From: Xun Li Date: Mon, 16 Dec 2024 11:36:11 -0800 Subject: [PATCH] Refactor TxSigner --- src/benchmarks/kms_stress.rs | 4 +- src/config.rs | 40 ++++++-- src/gas_pool/gas_pool_core.rs | 17 ++-- src/gas_pool_initializer.rs | 8 +- src/test_env.rs | 10 +- src/tx_signer.rs | 150 ------------------------------ src/tx_signer/in_memory_signer.rs | 38 ++++++++ src/tx_signer/mod.rs | 112 ++++++++++++++++++++++ src/tx_signer/sidecar_signer.rs | 79 ++++++++++++++++ 9 files changed, 280 insertions(+), 178 deletions(-) delete mode 100644 src/tx_signer.rs create mode 100644 src/tx_signer/in_memory_signer.rs create mode 100644 src/tx_signer/mod.rs create mode 100644 src/tx_signer/sidecar_signer.rs diff --git a/src/benchmarks/kms_stress.rs b/src/benchmarks/kms_stress.rs index 28d8705..509d6b2 100644 --- a/src/benchmarks/kms_stress.rs +++ b/src/benchmarks/kms_stress.rs @@ -1,13 +1,15 @@ // Copyright (c) Mysten Labs, Inc. // SPDX-License-Identifier: Apache-2.0 -use crate::tx_signer::{SidecarTxSigner, TxSigner}; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; use sui_types::base_types::{random_object_ref, SuiAddress}; use sui_types::transaction::{ProgrammableTransaction, TransactionData, TransactionKind}; +use crate::tx_signer::sidecar_signer::SidecarTxSigner; +use crate::tx_signer::TxSignerTrait; + pub async fn run_kms_stress_test(kms_url: String, num_tasks: usize) { let signer = SidecarTxSigner::new(kms_url).await; let test_tx_data = TransactionData::new( diff --git a/src/config.rs b/src/config.rs index c159e84..46a0168 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,7 +1,9 @@ // Copyright (c) Mysten Labs, Inc. // SPDX-License-Identifier: Apache-2.0 -use crate::tx_signer::{SidecarTxSigner, TestTxSigner, TxSigner}; +use crate::tx_signer::in_memory_signer::InMemoryTxSigner; +use crate::tx_signer::sidecar_signer::SidecarTxSigner; +use crate::tx_signer::{TxSigner, TxSignerTrait}; use serde::{Deserialize, Serialize}; use serde_with::serde_as; use std::net::Ipv4Addr; @@ -82,7 +84,15 @@ impl Default for GasPoolStorageConfig { pub enum TxSignerConfig { Local { keypair: SuiKeyPair }, Sidecar { sidecar_url: String }, - MultiSidecar { sidecar_urls: Vec }, + MultiSigner { signers: Vec }, +} + +#[serde_as] +#[derive(Debug, Deserialize, Serialize)] +#[serde(rename_all = "kebab-case")] +pub enum SingleSignerType { + Local { keypair: SuiKeyPair }, + Sidecar { sidecar_url: String }, } impl Default for TxSignerConfig { @@ -95,16 +105,28 @@ impl Default for TxSignerConfig { } impl TxSignerConfig { - pub async fn new_signer(self) -> Arc { - match self { - TxSignerConfig::Local { keypair } => TestTxSigner::new(keypair), + pub async fn new_signer(self) -> Arc { + let all_signers: Vec> = match self { + TxSignerConfig::Local { keypair } => vec![InMemoryTxSigner::new(keypair)], TxSignerConfig::Sidecar { sidecar_url } => { - SidecarTxSigner::new(vec![sidecar_url]).await + vec![SidecarTxSigner::new(sidecar_url).await] } - TxSignerConfig::MultiSidecar { sidecar_urls } => { - SidecarTxSigner::new(sidecar_urls).await + TxSignerConfig::MultiSigner { signers } => { + let mut all_signers: Vec> = Vec::new(); + for signer_config in signers { + match signer_config { + SingleSignerType::Local { keypair } => { + all_signers.push(InMemoryTxSigner::new(keypair)) + } + SingleSignerType::Sidecar { sidecar_url } => { + all_signers.push(SidecarTxSigner::new(sidecar_url).await) + } + } + } + all_signers } - } + }; + TxSigner::new(all_signers) } } diff --git a/src/gas_pool/gas_pool_core.rs b/src/gas_pool/gas_pool_core.rs index 2a46dd5..df21438 100644 --- a/src/gas_pool/gas_pool_core.rs +++ b/src/gas_pool/gas_pool_core.rs @@ -34,7 +34,7 @@ pub struct GasPoolContainer { } pub struct GasPool { - signer: Arc, + signer: Arc, gas_pool_store: Arc, sui_client: SuiClient, metrics: Arc, @@ -43,7 +43,7 @@ pub struct GasPool { impl GasPool { pub async fn new( - signer: Arc, + signer: Arc, gas_pool_store: Arc, sui_client: SuiClient, metrics: Arc, @@ -274,19 +274,14 @@ impl GasPool { /// Performs an end-to-end flow of reserving gas, signing a transaction, and releasing the gas coins. pub async fn debug_check_health(&self) -> anyhow::Result<()> { let gas_budget = MIST_PER_SUI / 10; - let (_address, _reservation_id, gas_coins) = + let (sender, _reservation_id, gas_coins) = self.reserve_gas(gas_budget, Duration::from_secs(3)).await?; let tx_kind = TransactionKind::ProgrammableTransaction( ProgrammableTransactionBuilder::new().finish(), ); // Since we just want to check the health of the signer, we don't need to actually execute the transaction. - let tx_data = TransactionData::new_with_gas_coins( - tx_kind, - SuiAddress::default(), - gas_coins, - gas_budget, - 0, - ); + let tx_data = + TransactionData::new_with_gas_coins(tx_kind, sender, gas_coins, gas_budget, 0); self.signer.sign_transaction(&tx_data).await?; Ok(()) } @@ -337,7 +332,7 @@ impl GasPool { impl GasPoolContainer { pub async fn new( - signer: Arc, + signer: Arc, gas_pool_store: Arc, sui_client: SuiClient, gas_usage_daily_cap: u64, diff --git a/src/gas_pool_initializer.rs b/src/gas_pool_initializer.rs index 3f6eb08..a9124f6 100644 --- a/src/gas_pool_initializer.rs +++ b/src/gas_pool_initializer.rs @@ -36,7 +36,7 @@ const MAX_INIT_DURATION_SEC: u64 = 60 * 60 * 12; struct CoinSplitEnv { target_init_coin_balance: u64, gas_cost_per_object: u64, - signer: Arc, + signer: Arc, sponsor_address: SuiAddress, sui_client: SuiClient, task_queue: Arc>>>>, @@ -177,7 +177,7 @@ impl GasPoolInitializer { sui_client: SuiClient, storage: Arc, coin_init_config: CoinInitConfig, - signer: Arc, + signer: Arc, ) -> Self { for address in signer.get_all_addresses() { if !storage.is_initialized(address).await.unwrap() { @@ -211,7 +211,7 @@ impl GasPoolInitializer { sui_client: SuiClient, storage: Arc, coin_init_config: CoinInitConfig, - signer: Arc, + signer: Arc, mut cancel_receiver: tokio::sync::oneshot::Receiver<()>, ) { loop { @@ -243,7 +243,7 @@ impl GasPoolInitializer { storage: &Arc, mode: RunMode, target_init_coin_balance: u64, - signer: &Arc, + signer: &Arc, ) { if storage .acquire_init_lock(sponsor_address, MAX_INIT_DURATION_SEC) diff --git a/src/test_env.rs b/src/test_env.rs index d713858..7cab63e 100644 --- a/src/test_env.rs +++ b/src/test_env.rs @@ -8,7 +8,8 @@ use crate::metrics::{GasPoolCoreMetrics, GasPoolRpcMetrics}; use crate::rpc::GasPoolServer; use crate::storage::connect_storage_for_testing; use crate::sui_client::SuiClient; -use crate::tx_signer::{TestTxSigner, TxSigner}; +use crate::tx_signer::in_memory_signer::InMemoryTxSigner; +use crate::tx_signer::TxSigner; use crate::AUTH_ENV_NAME; use std::sync::Arc; use sui_config::local_ip_utils::{get_available_port, localhost_for_testing}; @@ -21,7 +22,7 @@ use sui_types::transaction::{TransactionData, TransactionDataAPI}; use test_cluster::{TestCluster, TestClusterBuilder}; use tracing::debug; -pub async fn start_sui_cluster(init_gas_amounts: Vec) -> (TestCluster, Arc) { +pub async fn start_sui_cluster(init_gas_amounts: Vec) -> (TestCluster, Arc) { let (sponsor, keypair) = get_account_key_pair(); let cluster = TestClusterBuilder::new() .with_accounts(vec![ @@ -37,7 +38,10 @@ pub async fn start_sui_cluster(init_gas_amounts: Vec) -> (TestCluster, Arc< ]) .build() .await; - (cluster, TestTxSigner::new(keypair.into())) + ( + cluster, + TxSigner::new(vec![InMemoryTxSigner::new(keypair.into())]), + ) } pub async fn start_gas_station( diff --git a/src/tx_signer.rs b/src/tx_signer.rs deleted file mode 100644 index dadf418..0000000 --- a/src/tx_signer.rs +++ /dev/null @@ -1,150 +0,0 @@ -// Copyright (c) Mysten Labs, Inc. -// SPDX-License-Identifier: Apache-2.0 - -use anyhow::anyhow; -use fastcrypto::encoding::{Base64, Encoding}; -use reqwest::Client; -use serde::Deserialize; -use serde_json::json; -use shared_crypto::intent::{Intent, IntentMessage}; -use std::collections::HashMap; -use std::str::FromStr; -use std::sync::atomic::AtomicUsize; -use std::sync::{atomic, Arc}; -use sui_types::base_types::SuiAddress; -use sui_types::crypto::{Signature, SuiKeyPair}; -use sui_types::signature::GenericSignature; -use sui_types::transaction::{TransactionData, TransactionDataAPI}; - -#[async_trait::async_trait] -pub trait TxSigner: Send + Sync { - async fn sign_transaction(&self, tx_data: &TransactionData) - -> anyhow::Result; - fn get_one_address(&self) -> SuiAddress; - fn get_all_addresses(&self) -> Vec; - fn is_valid_address(&self, address: &SuiAddress) -> bool; -} - -#[derive(Deserialize)] -#[serde(rename_all = "camelCase")] -struct SignatureResponse { - signature: String, -} - -#[derive(Deserialize)] -#[serde(rename_all = "camelCase")] -struct SuiAddressResponse { - sui_pubkey_address: SuiAddress, -} - -// TODO: Add a mock side car server with tests for multi-address support. -pub struct SidecarTxSigner { - client: Client, - sidecar_url_map: HashMap, - sui_addresses: Vec, - next_address_idx: AtomicUsize, -} - -impl SidecarTxSigner { - pub async fn new(sidecar_urls: Vec) -> Arc { - let client = Client::new(); - let mut sidecar_url_map = HashMap::new(); - let mut sui_addresses = vec![]; - for sidecar_url in sidecar_urls { - let resp = client - .get(format!("{}/{}", &sidecar_url, "get-pubkey-address")) - .send() - .await - .unwrap_or_else(|err| panic!("Failed to get pubkey address: {}", err)); - let sui_address = resp - .json::() - .await - .unwrap_or_else(|err| panic!("Failed to parse address response: {}", err)) - .sui_pubkey_address; - sui_addresses.push(sui_address); - sidecar_url_map.insert(sui_address, sidecar_url); - } - Arc::new(Self { - client, - sidecar_url_map, - sui_addresses, - next_address_idx: AtomicUsize::new(0), - }) - } -} - -#[async_trait::async_trait] -impl TxSigner for SidecarTxSigner { - async fn sign_transaction( - &self, - tx_data: &TransactionData, - ) -> anyhow::Result { - let sponsor_address = tx_data.gas_data().owner; - let sidecar_url = self - .sidecar_url_map - .get(&sponsor_address) - .ok_or_else(|| anyhow!("Address is not a valid sponsor: {:?}", sponsor_address))?; - let bytes = Base64::encode(bcs::to_bytes(&tx_data)?); - let resp = self - .client - .post(format!("{}/{}", sidecar_url, "sign-transaction")) - .header("Content-Type", "application/json") - .json(&json!({"txBytes": bytes})) - .send() - .await?; - let sig_bytes = resp.json::().await?; - let sig = GenericSignature::from_str(&sig_bytes.signature) - .map_err(|err| anyhow!(err.to_string()))?; - Ok(sig) - } - - fn get_one_address(&self) -> SuiAddress { - // Round robin the address we are using. - let idx = self - .next_address_idx - .fetch_add(1, atomic::Ordering::Relaxed); - self.sui_addresses[idx % self.sui_addresses.len()] - } - - fn get_all_addresses(&self) -> Vec { - self.sui_addresses.clone() - } - - fn is_valid_address(&self, address: &SuiAddress) -> bool { - self.sidecar_url_map.contains_key(address) - } -} - -pub struct TestTxSigner { - keypair: SuiKeyPair, -} - -impl TestTxSigner { - pub fn new(keypair: SuiKeyPair) -> Arc { - Arc::new(Self { keypair }) - } -} - -#[async_trait::async_trait] -impl TxSigner for TestTxSigner { - async fn sign_transaction( - &self, - tx_data: &TransactionData, - ) -> anyhow::Result { - let intent_msg = IntentMessage::new(Intent::sui_transaction(), tx_data); - let sponsor_sig = Signature::new_secure(&intent_msg, &self.keypair).into(); - Ok(sponsor_sig) - } - - fn get_one_address(&self) -> SuiAddress { - (&self.keypair.public()).into() - } - - fn get_all_addresses(&self) -> Vec { - vec![self.get_one_address()] - } - - fn is_valid_address(&self, address: &SuiAddress) -> bool { - address == &self.get_one_address() - } -} diff --git a/src/tx_signer/in_memory_signer.rs b/src/tx_signer/in_memory_signer.rs new file mode 100644 index 0000000..160a32f --- /dev/null +++ b/src/tx_signer/in_memory_signer.rs @@ -0,0 +1,38 @@ +// Copyright (c) Mysten Labs, Inc. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::Arc; + +use shared_crypto::intent::{Intent, IntentMessage}; +use sui_types::base_types::SuiAddress; +use sui_types::crypto::{Signature, SuiKeyPair}; +use sui_types::signature::GenericSignature; +use sui_types::transaction::TransactionData; + +use super::TxSignerTrait; + +pub struct InMemoryTxSigner { + keypair: SuiKeyPair, +} + +impl InMemoryTxSigner { + pub fn new(keypair: SuiKeyPair) -> Arc { + Arc::new(Self { keypair }) + } +} + +#[async_trait::async_trait] +impl TxSignerTrait for InMemoryTxSigner { + async fn sign_transaction( + &self, + tx_data: &TransactionData, + ) -> anyhow::Result { + let intent_msg = IntentMessage::new(Intent::sui_transaction(), tx_data); + let sponsor_sig = Signature::new_secure(&intent_msg, &self.keypair).into(); + Ok(sponsor_sig) + } + + fn sui_address(&self) -> SuiAddress { + (&self.keypair.public()).into() + } +} diff --git a/src/tx_signer/mod.rs b/src/tx_signer/mod.rs new file mode 100644 index 0000000..055bf1a --- /dev/null +++ b/src/tx_signer/mod.rs @@ -0,0 +1,112 @@ +// Copyright (c) Mysten Labs, Inc. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::HashMap; +use std::sync::atomic::{self, AtomicUsize}; +use std::sync::Arc; +use sui_types::base_types::SuiAddress; +use sui_types::signature::GenericSignature; +use sui_types::transaction::{TransactionData, TransactionDataAPI}; + +pub mod in_memory_signer; +pub mod sidecar_signer; + +#[async_trait::async_trait] +pub trait TxSignerTrait: Send + Sync { + async fn sign_transaction(&self, tx_data: &TransactionData) + -> anyhow::Result; + fn sui_address(&self) -> SuiAddress; +} + +pub struct TxSigner { + signers: Vec>, + next_signer_idx: AtomicUsize, + address_index_map: HashMap, +} + +impl TxSigner { + pub fn new(signers: Vec>) -> Arc { + let address_index_map: HashMap<_, _> = signers + .iter() + .enumerate() + .map(|(i, s)| (s.sui_address(), i)) + .collect(); + Arc::new(Self { + signers, + next_signer_idx: AtomicUsize::new(0), + address_index_map, + }) + } + + pub fn get_all_addresses(&self) -> Vec { + self.signers.iter().map(|s| s.sui_address()).collect() + } + + pub fn is_valid_address(&self, address: &SuiAddress) -> bool { + self.address_index_map.contains_key(address) + } + + pub fn get_one_address(&self) -> SuiAddress { + let idx = self.next_signer_idx.fetch_add(1, atomic::Ordering::Relaxed); + self.signers[idx % self.signers.len()].sui_address() + } + + pub async fn sign_transaction( + &self, + tx_data: &TransactionData, + ) -> anyhow::Result { + let sponsor_address = tx_data.gas_data().owner; + let idx = *self + .address_index_map + .get(&sponsor_address) + .ok_or_else(|| anyhow::anyhow!("No signer found for address: {}", sponsor_address))?; + self.signers[idx].sign_transaction(tx_data).await + } +} + +#[cfg(test)] +mod tests { + use in_memory_signer::InMemoryTxSigner; + use sui_types::{ + crypto::get_account_key_pair, + programmable_transaction_builder::ProgrammableTransactionBuilder, + transaction::TransactionKind, + }; + + use super::*; + + #[tokio::test] + async fn test_multi_tx_signer() { + let (sender1, key1) = get_account_key_pair(); + let (sender2, key2) = get_account_key_pair(); + let (sender3, key3) = get_account_key_pair(); + let senders = vec![sender1, sender2, sender3]; + let signer1 = InMemoryTxSigner::new(key1.into()); + let signer2 = InMemoryTxSigner::new(key2.into()); + let signer3 = InMemoryTxSigner::new(key3.into()); + let tx_signer = TxSigner::new(vec![signer1, signer2, signer3]); + for sender in senders { + let tx_data = TransactionData::new_with_gas_coins( + TransactionKind::ProgrammableTransaction( + ProgrammableTransactionBuilder::new().finish(), + ), + sender, + vec![], + 0, + 0, + ); + tx_signer.sign_transaction(&tx_data).await.unwrap(); + } + let (sender4, _) = get_account_key_pair(); + let tx_data = TransactionData::new_with_gas_coins( + TransactionKind::ProgrammableTransaction( + ProgrammableTransactionBuilder::new().finish(), + ), + sender4, + vec![], + 0, + 0, + ); + assert!(tx_signer.sign_transaction(&tx_data).await.is_err()); + } +} diff --git a/src/tx_signer/sidecar_signer.rs b/src/tx_signer/sidecar_signer.rs new file mode 100644 index 0000000..d5b08e1 --- /dev/null +++ b/src/tx_signer/sidecar_signer.rs @@ -0,0 +1,79 @@ +// Copyright (c) Mysten Labs, Inc. +// SPDX-License-Identifier: Apache-2.0 + +use anyhow::anyhow; +use fastcrypto::encoding::{Base64, Encoding}; +use reqwest::Client; +use serde::Deserialize; +use serde_json::json; +use std::str::FromStr; +use std::sync::Arc; +use sui_types::base_types::SuiAddress; +use sui_types::signature::GenericSignature; +use sui_types::transaction::TransactionData; + +use super::TxSignerTrait; + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct SignatureResponse { + signature: String, +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct SuiAddressResponse { + sui_pubkey_address: SuiAddress, +} + +pub struct SidecarTxSigner { + client: Client, + sidecar_url: String, + sui_address: SuiAddress, +} + +impl SidecarTxSigner { + pub async fn new(sidecar_url: String) -> Arc { + let client = Client::new(); + let resp = client + .get(format!("{}/{}", &sidecar_url, "get-pubkey-address")) + .send() + .await + .unwrap_or_else(|err| panic!("Failed to get pubkey address: {}", err)); + let sui_address = resp + .json::() + .await + .unwrap_or_else(|err| panic!("Failed to parse address response: {}", err)) + .sui_pubkey_address; + Arc::new(Self { + client, + sidecar_url, + sui_address, + }) + } +} + +#[async_trait::async_trait] +impl TxSignerTrait for SidecarTxSigner { + async fn sign_transaction( + &self, + tx_data: &TransactionData, + ) -> anyhow::Result { + let bytes = Base64::encode(bcs::to_bytes(&tx_data)?); + let resp = self + .client + .post(format!("{}/{}", self.sidecar_url, "sign-transaction")) + .header("Content-Type", "application/json") + .json(&json!({"txBytes": bytes})) + .send() + .await?; + let sig_bytes = resp.json::().await?; + let sig = GenericSignature::from_str(&sig_bytes.signature) + .map_err(|err| anyhow!(err.to_string()))?; + Ok(sig) + } + + fn sui_address(&self) -> SuiAddress { + self.sui_address + } +}