feat: Add AWS KMS signing support to transaction-sender

fhevm-engine/transaction-sender/src/bin/transaction_sender.rs

@@ -1,21 +1,30 @@
-use std::str::FromStr;
-
 use alloy::{
-    network::EthereumWallet,
-    primitives::Address,
+    network::{EthereumWallet, TxEnvelope},
+    primitives::{Address, Signature as AlloySignature},
     providers::{ProviderBuilder, WsConnect},
-    signers::local::PrivateKeySigner,
+    signer::{Error as SignerError, Signer},
+    signers::{aws::AwsSigner, local::PrivateKeySigner}, // PrivateKeySigner needed for DynamicSigner::Local variant if constructed here
     transports::http::reqwest::Url,
 };
+// async_trait removed as DynamicSigner definition is moved to lib.rs
+use aws_config::meta::region::RegionProviderChain;
+use aws_sdk_kms::Client as KmsClient;
 use clap::Parser;
+use std::str::FromStr; // For PrivateKeySigner::from_str
 use tokio::signal::unix::{signal, SignalKind};
 use tokio_util::sync::CancellationToken;
 use transaction_sender::{
-    ConfigSettings, FillersWithoutNonceManagement, NonceManagedProvider, TransactionSender,
+    ConfigSettings, DynamicSigner, FillersWithoutNonceManagement, NonceManagedProvider,
+    TransactionSender, // Added DynamicSigner to this import line
 };
 
 #[derive(Parser, Debug, Clone)]
 #[command(version, about, long_about = None)]
+#[clap(group(
+    clap::ArgGroup::new("signing_method")
+        .required(true)
+        .args(&["private_key", "aws_kms_key_id"]),
+))]
 struct Conf {
     #[arg(short, long)]
     input_verification_address: Address,
@@ -29,8 +38,14 @@ struct Conf {
     #[arg(short, long)]
     gateway_url: Url,
 
-    #[arg(short, long)]
-    private_key: String,
+    #[arg(short = 'k', long, env = "PRIVATE_KEY", group = "signing_method")]
+    private_key: Option<String>,
+
+    #[arg(long, env = "AWS_REGION", requires = "aws_kms_key_id")]
+    aws_region: Option<String>,
+
+    #[arg(long, env = "AWS_KMS_KEY_ID", group = "signing_method", requires = "aws_region")]
+    aws_kms_key_id: Option<String>,
 
     #[arg(short, long)]
     database_url: Option<String>,
@@ -87,6 +102,8 @@ struct Conf {
     review_after_transport_retries: u16,
 }
 
+// DynamicSigner enum and its impl Signer removed from here as it's now in lib.rs
+
 fn install_signal_handlers(cancel_token: CancellationToken) -> anyhow::Result<()> {
     let mut sigint = signal(SignalKind::interrupt())?;
     let mut sigterm = signal(SignalKind::terminate())?;
@@ -104,26 +121,72 @@ fn install_signal_handlers(cancel_token: CancellationToken) -> anyhow::Result<()
 async fn main() -> anyhow::Result<()> {
     tracing_subscriber::fmt().json().with_level(true).init();
     let conf = Conf::parse();
-    let signer = PrivateKeySigner::from_str(conf.private_key.trim())?;
-    let wallet = EthereumWallet::new(signer.clone());
+
+    // Declare signer and wallet_chain_id
+    let signer: DynamicSigner;
+    let wallet_chain_id: u64; // AWS Signer needs it, and we'll use it for wallet consistently
+
+    if let Some(pk_str) = conf.private_key.as_deref() {
+        let private_key_signer = PrivateKeySigner::from_str(pk_str.trim())
+            .map_err(|e| anyhow::anyhow!("Failed to parse private key: {}", e))?;
+
+        // Fetch chain_id for the wallet
+        let provider_for_chain_id = ProviderBuilder::default()
+            .on_ws(WsConnect::new(conf.gateway_url.clone()))
+            .await?;
+        wallet_chain_id = provider_for_chain_id.get_chain_id().await?;
+
+        signer = DynamicSigner::Local(private_key_signer);
+    } else if let (Some(key_id), Some(region_str)) = (conf.aws_kms_key_id.clone(), conf.aws_region.clone()) {
+        // Fetch chain_id for AwsSigner and wallet
+        let provider_for_chain_id = ProviderBuilder::default()
+            .on_ws(WsConnect::new(conf.gateway_url.clone()))
+            .await?;
+        wallet_chain_id = provider_for_chain_id.get_chain_id().await?;
+
+        let region_provider = RegionProviderChain::first_try(aws_config::Region::new(region_str))
+            .or_default_provider()
+            .or_else(aws_config::Region::new("us-east-1")); // Default fallback region
+
+        let sdk_config = aws_config::from_env().region(region_provider).load().await;
+        let kms_client = KmsClient::new(&sdk_config);
+
+        // AwsSigner::new is not async and takes chain_id as u64
+        let aws_signer = AwsSigner::new(kms_client, key_id, wallet_chain_id);
+        signer = DynamicSigner::Aws(aws_signer);
+    } else {
+        // This case should ideally be prevented by clap's ArgGroup validation
+        return Err(anyhow::anyhow!("Invalid signer configuration: either private_key or AWS KMS parameters (aws_kms_key_id, aws_region) must be provided."));
+    }
+
+    // Create the wallet, explicitly setting the chain_id obtained
+    let mut wallet: EthereumWallet<DynamicSigner> = EthereumWallet::new(signer.clone());
+    wallet.set_chain_id(Some(wallet_chain_id)); // Set chain_id on the wallet
+
     let database_url = conf
         .database_url
         .clone()
         .unwrap_or_else(|| std::env::var("DATABASE_URL").expect("DATABASE_URL is undefined"));
     let cancel_token = CancellationToken::new();
+
+    // Provider for NonceManagedProvider uses the new wallet
+    let provider_instance = ProviderBuilder::default()
+        .filler(FillersWithoutNonceManagement::default())
+        .wallet(wallet.clone()) // Uses the new wallet: EthereumWallet<DynamicSigner>
+        .on_ws(WsConnect::new(conf.gateway_url)) // gateway_url is still from conf
+        .await?;
+
     let provider = NonceManagedProvider::new(
-        ProviderBuilder::default()
-            .filler(FillersWithoutNonceManagement::default())
-            .wallet(wallet.clone())
-            .on_ws(WsConnect::new(conf.gateway_url))
-            .await?,
-        Some(wallet.default_signer().address()),
+        provider_instance,
+        Some(wallet.address()), // address from DynamicSigner via EthereumWallet
     );
+
+    // TransactionSender::new will have a type error for the signer argument here, which is expected for this step.
     let sender = TransactionSender::new(
         conf.input_verification_address,
         conf.ciphertext_commits_address,
         conf.multichain_acl_address,
-        signer,
+        signer, // This is DynamicSigner, TransactionSender expects PrivateKeySigner
         provider,
         cancel_token.clone(),
         ConfigSettings {

fhevm-engine/transaction-sender/src/lib.rs

@@ -2,6 +2,14 @@ mod nonce_managed_provider;
 mod ops;
 mod transaction_sender;
 
+// IMPORTS FOR DynamicSigner - START
+use alloy_primitives::{Address, Signature as AlloySignature};
+use alloy_network::TxEnvelope;
+use alloy_signer::{Signer, Error as SignerError};
+use alloy_signers::{local::PrivateKeySigner, aws::AwsSigner};
+use async_trait::async_trait;
+// IMPORTS FOR DynamicSigner - END
+
 #[derive(Clone, Debug)]
 pub struct ConfigSettings {
     pub database_url: String,
@@ -62,5 +70,61 @@ impl Default for ConfigSettings {
 pub use nonce_managed_provider::FillersWithoutNonceManagement;
 pub use nonce_managed_provider::NonceManagedProvider;
 pub use transaction_sender::TransactionSender;
+// EXPORT DynamicSigner
+pub use self::dynamic_signer::DynamicSigner;
 
 pub const REVIEW: &str = "review";
+
+// Definition of DynamicSigner moved here
+pub mod dynamic_signer {
+    // Explicit imports for items used within this module
+    use alloy_primitives::{Address, Signature as AlloySignature};
+    use alloy_network::TxEnvelope;
+    use alloy_signer::{Signer, Error as SignerError};
+    use alloy_signers::{local::PrivateKeySigner, aws::AwsSigner};
+    use async_trait::async_trait;
+
+    #[derive(Clone, Debug)]
+    pub enum DynamicSigner {
+        Local(PrivateKeySigner),
+        Aws(AwsSigner),
+    }
+
+    #[async_trait]
+    impl Signer<AlloySignature> for DynamicSigner {
+        async fn sign_message(&self, message: &[u8]) -> Result<AlloySignature, SignerError> {
+            match self {
+                DynamicSigner::Local(s) => s.sign_message(message).await,
+                DynamicSigner::Aws(s) => s.sign_message(message).await,
+            }
+        }
+
+        async fn sign_transaction(&self, tx: &mut TxEnvelope) -> Result<AlloySignature, SignerError> {
+            match self {
+                DynamicSigner::Local(s) => s.sign_transaction(tx).await,
+                DynamicSigner::Aws(s) => s.sign_transaction(tx).await,
+            }
+        }
+
+        fn address(&self) -> Address {
+            match self {
+                DynamicSigner::Local(s) => s.address(),
+                DynamicSigner::Aws(s) => s.address(),
+            }
+        }
+
+        fn chain_id(&self) -> Option<u64> {
+            match self {
+                DynamicSigner::Local(s) => s.chain_id(),
+                DynamicSigner::Aws(s) => s.chain_id(),
+            }
+        }
+
+        fn set_chain_id(&mut self, chain_id: Option<u64>) {
+            match self {
+                DynamicSigner::Local(s) => s.set_chain_id(chain_id),
+                DynamicSigner::Aws(s) => s.set_chain_id(chain_id),
+            }
+        }
+    }
+}

fhevm-engine/transaction-sender/src/ops/verify_proof.rs

@@ -4,10 +4,11 @@ use alloy::network::TransactionBuilder;
 use alloy::primitives::{Address, U256};
 use alloy::providers::Provider;
 use alloy::rpc::types::TransactionRequest;
-use alloy::signers::local::PrivateKeySigner;
-use alloy::signers::SignerSync;
+// PrivateKeySigner removed, DynamicSigner will be used
+// SignerSync removed, async Signer trait's sign_hash will be used
 use alloy::sol;
 use alloy::{network::Ethereum, primitives::FixedBytes, sol_types::SolStruct};
+use crate::DynamicSigner; // Added DynamicSigner
 use async_trait::async_trait;
 use sqlx::{Pool, Postgres};
 use std::convert::TryInto;
@@ -35,7 +36,7 @@ sol!(
 pub(crate) struct VerifyProofOperation<P: Provider<Ethereum> + Clone + 'static> {
     input_verification_address: Address,
     provider: NonceManagedProvider<P>,
-    signer: PrivateKeySigner,
+    signer: DynamicSigner, // Changed to DynamicSigner
     conf: crate::ConfigSettings,
     gas: Option<u64>,
     gw_chain_id: u64,
@@ -46,7 +47,7 @@ impl<P: alloy::providers::Provider<Ethereum> + Clone + 'static> VerifyProofOpera
     pub(crate) async fn new(
         input_verification_address: Address,
         provider: NonceManagedProvider<P>,
-        signer: PrivateKeySigner,
+        signer: DynamicSigner, // Changed to DynamicSigner
         conf: crate::ConfigSettings,
         gas: Option<u64>,
         db_pool: Pool<Postgres>,
@@ -263,9 +264,10 @@ where
                         contractChainId: U256::from(row.chain_id),
                     }
                     .eip712_signing_hash(&domain);
-                    let signature = self
+                    let signature = self // Changed to async call
                         .signer
-                        .sign_hash_sync(&signing_hash)
+                        .sign_hash(&signing_hash)
+                        .await
                         .expect("signing failed");
 
                     if let Some(gas) = self.gas {

fhevm-engine/transaction-sender/src/transaction_sender.rs

@@ -1,14 +1,15 @@
-use alloy::{
-    network::Ethereum, primitives::Address, providers::Provider, signers::local::PrivateKeySigner,
-};
+use alloy::{network::Ethereum, primitives::Address, providers::Provider};
+// PrivateKeySigner import removed as DynamicSigner is used in `new` function signature
 use futures_util::FutureExt;
 use sqlx::{postgres::PgListener, Pool, Postgres};
 use std::{sync::Arc, time::Duration};
 use tokio::task::JoinSet;
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, error, info};
 
-use crate::{nonce_managed_provider::NonceManagedProvider, ops, ConfigSettings};
+use crate::{
+    nonce_managed_provider::NonceManagedProvider, ops, ConfigSettings, DynamicSigner,
+}; // Added DynamicSigner
 
 #[derive(Clone)]
 pub struct TransactionSender<P: Provider<Ethereum> + Clone + 'static> {
@@ -27,7 +28,7 @@ impl<P: Provider<Ethereum> + Clone + 'static> TransactionSender<P> {
         input_verification_address: Address,
         ciphertext_commits_address: Address,
         multichain_acl_address: Address,
-        signer: PrivateKeySigner,
+        signer: DynamicSigner, // Changed from PrivateKeySigner
         provider: NonceManagedProvider<P>,
         cancel_token: CancellationToken,
         conf: ConfigSettings,