Refactor TxSigner

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(

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<String> },
+    MultiSigner { signers: Vec<SingleSignerType> },
+}
+
+#[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<dyn TxSigner> {
-        match self {
-            TxSignerConfig::Local { keypair } => TestTxSigner::new(keypair),
+    pub async fn new_signer(self) -> Arc<TxSigner> {
+        let all_signers: Vec<Arc<dyn TxSignerTrait>> = 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<Arc<dyn TxSignerTrait>> = 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)
     }
 }
 

src/gas_pool/gas_pool_core.rs

@@ -34,7 +34,7 @@ pub struct GasPoolContainer {
 }
 
 pub struct GasPool {
-    signer: Arc<dyn TxSigner>,
+    signer: Arc<TxSigner>,
     gas_pool_store: Arc<dyn Storage>,
     sui_client: SuiClient,
     metrics: Arc<GasPoolCoreMetrics>,
@@ -43,7 +43,7 @@ pub struct GasPool {
 
 impl GasPool {
     pub async fn new(
-        signer: Arc<dyn TxSigner>,
+        signer: Arc<TxSigner>,
         gas_pool_store: Arc<dyn Storage>,
         sui_client: SuiClient,
         metrics: Arc<GasPoolCoreMetrics>,
@@ -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<dyn TxSigner>,
+        signer: Arc<TxSigner>,
         gas_pool_store: Arc<dyn Storage>,
         sui_client: SuiClient,
         gas_usage_daily_cap: u64,

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<dyn TxSigner>,
+    signer: Arc<TxSigner>,
     sponsor_address: SuiAddress,
     sui_client: SuiClient,
     task_queue: Arc<Mutex<VecDeque<JoinHandle<Vec<GasCoin>>>>>,
@@ -177,7 +177,7 @@ impl GasPoolInitializer {
         sui_client: SuiClient,
         storage: Arc<dyn Storage>,
         coin_init_config: CoinInitConfig,
-        signer: Arc<dyn TxSigner>,
+        signer: Arc<TxSigner>,
     ) -> 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<dyn Storage>,
         coin_init_config: CoinInitConfig,
-        signer: Arc<dyn TxSigner>,
+        signer: Arc<TxSigner>,
         mut cancel_receiver: tokio::sync::oneshot::Receiver<()>,
     ) {
         loop {
@@ -243,7 +243,7 @@ impl GasPoolInitializer {
         storage: &Arc<dyn Storage>,
         mode: RunMode,
         target_init_coin_balance: u64,
-        signer: &Arc<dyn TxSigner>,
+        signer: &Arc<TxSigner>,
     ) {
         if storage
             .acquire_init_lock(sponsor_address, MAX_INIT_DURATION_SEC)

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<u64>) -> (TestCluster, Arc<dyn TxSigner>) {
+pub async fn start_sui_cluster(init_gas_amounts: Vec<u64>) -> (TestCluster, Arc<TxSigner>) {
     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<u64>) -> (TestCluster, Arc<
         ])
         .build()
         .await;
-    (cluster, TestTxSigner::new(keypair.into()))
+    (
+        cluster,
+        TxSigner::new(vec![InMemoryTxSigner::new(keypair.into())]),
+    )
 }
 
 pub async fn start_gas_station(

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<Self> {
+        Arc::new(Self { keypair })
+    }
+}
+
+#[async_trait::async_trait]
+impl TxSignerTrait for InMemoryTxSigner {
+    async fn sign_transaction(
+        &self,
+        tx_data: &TransactionData,
+    ) -> anyhow::Result<GenericSignature> {
+        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()
+    }
+}