diff --git a/fhevm-engine/transaction-sender/src/bin/transaction_sender.rs b/fhevm-engine/transaction-sender/src/bin/transaction_sender.rs index 6c5655cb..e30423f1 100644 --- a/fhevm-engine/transaction-sender/src/bin/transaction_sender.rs +++ b/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, + + #[arg(long, env = "AWS_REGION", requires = "aws_kms_key_id")] + aws_region: Option, + + #[arg(long, env = "AWS_KMS_KEY_ID", group = "signing_method", requires = "aws_region")] + aws_kms_key_id: Option, #[arg(short, long)] database_url: Option, @@ -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 = 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 + .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 { diff --git a/fhevm-engine/transaction-sender/src/lib.rs b/fhevm-engine/transaction-sender/src/lib.rs index 8821e0c5..aa48c2c0 100644 --- a/fhevm-engine/transaction-sender/src/lib.rs +++ b/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 for DynamicSigner { + async fn sign_message(&self, message: &[u8]) -> Result { + 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 { + 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 { + match self { + DynamicSigner::Local(s) => s.chain_id(), + DynamicSigner::Aws(s) => s.chain_id(), + } + } + + fn set_chain_id(&mut self, chain_id: Option) { + match self { + DynamicSigner::Local(s) => s.set_chain_id(chain_id), + DynamicSigner::Aws(s) => s.set_chain_id(chain_id), + } + } + } +} diff --git a/fhevm-engine/transaction-sender/src/ops/verify_proof.rs b/fhevm-engine/transaction-sender/src/ops/verify_proof.rs index 39aff9d9..3fdb8706 100644 --- a/fhevm-engine/transaction-sender/src/ops/verify_proof.rs +++ b/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 + Clone + 'static> { input_verification_address: Address, provider: NonceManagedProvider

, - signer: PrivateKeySigner, + signer: DynamicSigner, // Changed to DynamicSigner conf: crate::ConfigSettings, gas: Option, gw_chain_id: u64, @@ -46,7 +47,7 @@ impl + Clone + 'static> VerifyProofOpera pub(crate) async fn new( input_verification_address: Address, provider: NonceManagedProvider

, - signer: PrivateKeySigner, + signer: DynamicSigner, // Changed to DynamicSigner conf: crate::ConfigSettings, gas: Option, db_pool: Pool, @@ -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 { diff --git a/fhevm-engine/transaction-sender/src/transaction_sender.rs b/fhevm-engine/transaction-sender/src/transaction_sender.rs index 7bf5ca0a..dc600051 100644 --- a/fhevm-engine/transaction-sender/src/transaction_sender.rs +++ b/fhevm-engine/transaction-sender/src/transaction_sender.rs @@ -1,6 +1,5 @@ -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}; @@ -8,7 +7,9 @@ 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 + Clone + 'static> { @@ -27,7 +28,7 @@ impl + Clone + 'static> TransactionSender

{ input_verification_address: Address, ciphertext_commits_address: Address, multichain_acl_address: Address, - signer: PrivateKeySigner, + signer: DynamicSigner, // Changed from PrivateKeySigner provider: NonceManagedProvider

, cancel_token: CancellationToken, conf: ConfigSettings, diff --git a/fhevm-engine/transaction-sender/tests/kms_integration_tests.rs b/fhevm-engine/transaction-sender/tests/kms_integration_tests.rs new file mode 100644 index 00000000..d9584359 --- /dev/null +++ b/fhevm-engine/transaction-sender/tests/kms_integration_tests.rs @@ -0,0 +1,140 @@ +// Suppress common clippy lints that might arise from test setup +#![allow(clippy::invisible_characters)] +#![allow(clippy::redundant_clone)] + +use alloy_primitives::{Address, Signature as AlloySignature, U256}; +use alloy_signer::{Signer, local::PrivateKeySigner}; +use alloy_signer::aws::AwsSigner; // Correct path for AwsSigner +use aws_config::meta::region::RegionProviderChain; +use aws_sdk_kms::{Client as KmsClient, config::Region, types::KeyUsageType}; +use rstest::rstest; +use std::str::FromStr; +use testcontainers::{clients, core::WaitFor, images::generic::GenericImage, Container}; // Added Container import + +// Helper to create a KMS client for LocalStack +async fn get_kms_client(kms_endpoint_url: String) -> KmsClient { + let region_provider = RegionProviderChain::first_try(Region::new("us-east-1")); // LocalStack default region + let sdk_config = aws_config::from_env() + .region(region_provider) + .endpoint_url(kms_endpoint_url) // Crucial: use the dynamic endpoint + .no_credentials_provider() // Important for localstack if no real AWS creds should be used + .load() + .await; + KmsClient::new(&sdk_config) +} + +// Represents the configuration for the signer to be tested +#[derive(Debug, Clone)] +enum SignerTypeConfig { + Local { pk_hex: String }, + Aws { + key_id_placeholder: String, // Placeholder, actual ARN generated in setup + region: String + }, +} + +async fn build_signer_for_test( + signer_type_config: &SignerTypeConfig, + chain_id: u64, + kms_endpoint_url: &str, + actual_kms_key_arn: &str, +) -> Box + Send + Sync> { + match signer_type_config { + SignerTypeConfig::Local { pk_hex } => { + let mut signer = PrivateKeySigner::from_str(pk_hex).expect("Failed to create local signer from hex"); + signer.set_chain_id(Some(chain_id)); + Box::new(signer) + } + SignerTypeConfig::Aws { region, .. } => { + let region_provider = RegionProviderChain::first_try(Region::new(region.clone())) + .or_default_provider(); + + let sdk_config = aws_config::from_env() + .region(region_provider) + .endpoint_url(kms_endpoint_url) + .no_credentials_provider() + .load() + .await; + let kms_client = KmsClient::new(&sdk_config); + let signer = AwsSigner::new(kms_client, actual_kms_key_arn.to_string(), chain_id); + Box::new(signer) + } + } +} + +// TestContext to manage Docker container and related KMS data +struct TestContext<'a> { + _docker_cli: &'a clients::Cli, // Keep reference to Docker client + localstack_container: Option>, + kms_endpoint_url: String, + actual_kms_key_arn: String, +} + +async fn setup_test_context<'a>( + docker_cli: &'a clients::Cli, + needs_kms: bool, +) -> TestContext<'a> { + if needs_kms { + let localstack_image = GenericImage::new("localstack/localstack", "3.0.0") + .with_env_var("SERVICES", "kms") + .with_env_var("DEFAULT_REGION", "us-east-1") + .with_wait_for(WaitFor::message_on_stderr("Ready.")); + + let node: Container<'a, GenericImage> = docker_cli.run(localstack_image); + let kms_port = node.get_host_port_ipv4(4566); + let kms_endpoint_url = format!("http://127.0.0.1:{}", kms_port); + + let kms_client_for_setup = get_kms_client(kms_endpoint_url.clone()).await; + let created_key = kms_client_for_setup.create_key() + .key_usage(KeyUsageType::SignVerify) + .send().await.expect("Failed to create KMS key in LocalStack"); + + let actual_kms_key_arn = created_key.key_metadata().expect("No key metadata") + .arn().expect("No ARN for created key").to_string(); + + TestContext { + _docker_cli: docker_cli, + localstack_container: Some(node), + kms_endpoint_url, + actual_kms_key_arn, + } + } else { + TestContext { + _docker_cli: docker_cli, + localstack_container: None, + kms_endpoint_url: String::new(), + actual_kms_key_arn: String::new(), + } + } +} + + +#[rstest] +#[case::local_signer(SignerTypeConfig::Local { pk_hex: "ac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80".to_string() })] +#[case::aws_signer(SignerTypeConfig::Aws { key_id_placeholder: "arn_will_be_generated".to_string(), region: "us-east-1".to_string() })] +#[tokio::test] +async fn test_sign_message_parameterized(#[case] signer_config: SignerTypeConfig, #[values("Hello World", "Test Message!")] message_str: &str) { + let docker = clients::Cli::default(); + let needs_kms = matches!(signer_config, SignerTypeConfig::Aws {..}); + + let test_ctx = setup_test_context(&docker, needs_kms).await; + + let chain_id = 1337u64; + + let signer = build_signer_for_test( + &signer_config, + chain_id, + &test_ctx.kms_endpoint_url, + &test_ctx.actual_kms_key_arn + ).await; + + let message = message_str.as_bytes(); + let signature = signer.sign_message(message).await.expect("Failed to sign message"); + let recovered_address = signature.recover_address_from_message(message).expect("Failed to recover address from signature"); + + assert_eq!(recovered_address, signer.address(), "Recovered address does not match signer address for config: {:?}", signer_config); + + if let Some(container) = test_ctx.localstack_container { + drop(container); + } +}