diff --git a/crates/cli/src/commands/types.rs b/crates/cli/src/commands/contender_subcommand.rs similarity index 88% rename from crates/cli/src/commands/types.rs rename to crates/cli/src/commands/contender_subcommand.rs index 6dbd98b..ff94160 100644 --- a/crates/cli/src/commands/types.rs +++ b/crates/cli/src/commands/contender_subcommand.rs @@ -47,13 +47,14 @@ Requires --priv-key to be set for each 'from' address in the given testfile.", )] duration: Option, - /// The seed to use for generating spam transactions. If not provided, one is generated. + /// The seed to use for generating spam transactions & accounts. #[arg( short, long, - long_help = "The seed to use for generating spam transactions" + long_help = "The seed to use for generating spam transactions", + default_value = "0xffffffffffffffffffffffffffffffff13131313131313131313131313131313" )] - seed: Option, + seed: String, /// The private keys to use for blockwise spamming. /// Required if `txs_per_block` is set. @@ -109,6 +110,23 @@ May be specified multiple times." default_value = "1.0" )] min_balance: String, + + /// The seed used to generate pool accounts. + #[arg( + short, + long, + long_help = "The seed used to generate pool accounts.", + default_value = "0xffffffffffffffffffffffffffffffff13131313131313131313131313131313" + )] + seed: String, + + /// The number of signers to generate for each pool. + #[arg( + short, + long = "signers-per-pool", + long_help = "The number of signers to generate for each pool." + )] + num_signers_per_pool: Option, }, #[command( diff --git a/crates/cli/src/commands/mod.rs b/crates/cli/src/commands/mod.rs index f41e287..d53de23 100644 --- a/crates/cli/src/commands/mod.rs +++ b/crates/cli/src/commands/mod.rs @@ -1,16 +1,16 @@ +mod contender_subcommand; mod report; mod run; mod setup; mod spam; -mod types; use clap::Parser; +pub use contender_subcommand::ContenderSubcommand; pub use report::report; pub use run::run; pub use setup::setup; pub use spam::{spam, SpamCommandArgs}; -pub use types::ContenderSubcommand; #[derive(Parser, Debug)] pub struct ContenderCli { diff --git a/crates/cli/src/commands/setup.rs b/crates/cli/src/commands/setup.rs index 05161e7..483260e 100644 --- a/crates/cli/src/commands/setup.rs +++ b/crates/cli/src/commands/setup.rs @@ -2,12 +2,17 @@ use alloy::{ network::AnyNetwork, primitives::utils::parse_ether, providers::ProviderBuilder, signers::local::PrivateKeySigner, transports::http::reqwest::Url, }; -use contender_core::{generator::RandSeed, test_scenario::TestScenario}; +use contender_core::{ + agent_controller::{AgentStore, SignerStore}, + generator::RandSeed, + test_scenario::TestScenario, +}; use contender_testfile::TestConfig; use std::str::FromStr; use crate::util::{ - check_private_keys_fns, find_insufficient_balance_addrs, get_signers_with_defaults, + check_private_keys_fns, find_insufficient_balance_addrs, fund_accounts, get_setup_pools, + get_signers_with_defaults, }; pub async fn setup( @@ -16,11 +21,14 @@ pub async fn setup( rpc_url: impl AsRef, private_keys: Option>, min_balance: String, + seed: RandSeed, + signers_per_period: usize, ) -> Result<(), Box> { let url = Url::parse(rpc_url.as_ref()).expect("Invalid RPC URL"); let rpc_client = ProviderBuilder::new() .network::() .on_http(url.to_owned()); + let eth_client = ProviderBuilder::new().on_http(url.to_owned()); let testconfig: TestConfig = TestConfig::from_file(testfile.as_ref())?; let min_balance = parse_ether(&min_balance)?; @@ -30,10 +38,10 @@ pub async fn setup( .iter() .map(|key| PrivateKeySigner::from_str(key).expect("invalid private key")) .collect::>(); - let signers = get_signers_with_defaults(private_keys); + let default_signers = get_signers_with_defaults(private_keys); check_private_keys_fns( &testconfig.setup.to_owned().unwrap_or_default(), - signers.as_slice(), + &default_signers, ); let broke_accounts = find_insufficient_balance_addrs( &user_signers.iter().map(|s| s.address()).collect::>(), @@ -42,17 +50,41 @@ pub async fn setup( ) .await?; if !broke_accounts.is_empty() { - panic!("Some accounts do not have sufficient balance"); + panic!("Insufficient balance in provided user account(s)"); + } + + let mut agents = AgentStore::new(); + let from_pools = get_setup_pools(&testconfig); + let mut all_signers = vec![]; + all_signers.extend_from_slice(&user_signers); + + for from_pool in &from_pools { + if agents.has_agent(from_pool) { + continue; + } + + let agent = SignerStore::new_random(signers_per_period, &seed, from_pool); + all_signers.extend_from_slice(&agent.signers); + agents.add_agent(from_pool, agent); } + fund_accounts( + &rpc_client, + ð_client, + min_balance, + &all_signers, + &default_signers[0], + ) + .await?; + let mut scenario = TestScenario::new( testconfig.to_owned(), db.clone().into(), url, None, - RandSeed::new(), - &signers, - Default::default(), + seed, + &default_signers, + agents, ) .await?; diff --git a/crates/cli/src/commands/spam.rs b/crates/cli/src/commands/spam.rs index a982902..8387965 100644 --- a/crates/cli/src/commands/spam.rs +++ b/crates/cli/src/commands/spam.rs @@ -14,7 +14,7 @@ use contender_core::{ use contender_testfile::TestConfig; use crate::util::{ - check_private_keys, fund_accounts, get_from_pools, get_signers_with_defaults, + check_private_keys, fund_accounts, get_signers_with_defaults, get_spam_pools, spam_callback_default, SpamCallbackType, }; @@ -26,7 +26,7 @@ pub struct SpamCommandArgs { pub txs_per_block: Option, pub txs_per_second: Option, pub duration: Option, - pub seed: Option, + pub seed: String, pub private_keys: Option>, pub disable_reports: bool, pub min_balance: String, @@ -37,10 +37,7 @@ pub async fn spam( args: SpamCommandArgs, ) -> Result<(), Box> { let testconfig = TestConfig::from_file(&args.testfile)?; - let rand_seed = args - .seed - .map(|s| RandSeed::seed_from_str(s.as_ref())) - .unwrap_or_default(); + let rand_seed = RandSeed::seed_from_str(&args.seed); let url = Url::parse(&args.rpc_url).expect("Invalid RPC URL"); let rpc_client = ProviderBuilder::new() .network::() @@ -57,7 +54,7 @@ pub async fn spam( .expect("No spam function calls found in testfile"); // distill all from_pool arguments from the spam requests - let from_pools = get_from_pools(&testconfig); + let from_pools = get_spam_pools(&testconfig); let mut agents = AgentStore::new(); let signers_per_period = args diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index e25fa11..5ac088d 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -5,7 +5,7 @@ mod util; use std::sync::LazyLock; use commands::{ContenderCli, ContenderSubcommand, SpamCommandArgs}; -use contender_core::db::DbOps; +use contender_core::{db::DbOps, generator::RandSeed}; use contender_sqlite::SqliteDb; static DB: LazyLock = std::sync::LazyLock::new(|| { @@ -16,6 +16,7 @@ static DB: LazyLock = std::sync::LazyLock::new(|| { async fn main() -> Result<(), Box> { let args = ContenderCli::parse_args(); let _ = DB.create_tables(); // ignore error; tables already exist + let db = DB.clone(); match args.command { ContenderSubcommand::Setup { @@ -23,7 +24,20 @@ async fn main() -> Result<(), Box> { rpc_url, private_keys, min_balance, - } => commands::setup(&DB.clone(), testfile, rpc_url, private_keys, min_balance).await?, + seed, + num_signers_per_pool, + } => { + commands::setup( + &db, + testfile, + rpc_url, + private_keys, + min_balance, + RandSeed::seed_from_str(&seed), + num_signers_per_pool.unwrap_or(1), + ) + .await? + } ContenderSubcommand::Spam { testfile, @@ -38,7 +52,7 @@ async fn main() -> Result<(), Box> { min_balance, } => { commands::spam( - &DB.clone(), + &db, SpamCommandArgs { testfile, rpc_url, @@ -55,9 +69,7 @@ async fn main() -> Result<(), Box> { .await? } - ContenderSubcommand::Report { id, out_file } => { - commands::report(&DB.clone(), id, out_file)? - } + ContenderSubcommand::Report { id, out_file } => commands::report(&db, id, out_file)?, ContenderSubcommand::Run { scenario, @@ -68,7 +80,7 @@ async fn main() -> Result<(), Box> { txs_per_duration, } => { commands::run( - &DB.clone(), + &db, scenario, rpc_url, private_key, diff --git a/crates/cli/src/util.rs b/crates/cli/src/util.rs index 5691b51..564e4c4 100644 --- a/crates/cli/src/util.rs +++ b/crates/cli/src/util.rs @@ -33,7 +33,17 @@ pub const DEFAULT_PRV_KEYS: [&str; 10] = [ "0x2a871d0798f97d79848a013d4936a73bf4cc922c825d33c1cf7073dff6d409c6", ]; -pub fn get_from_pools(testconfig: &TestConfig) -> Vec { +pub fn get_setup_pools(testconfig: &TestConfig) -> Vec { + testconfig + .setup + .to_owned() + .unwrap_or_default() + .into_iter() + .filter_map(|s| s.from_pool) + .collect() +} + +pub fn get_spam_pools(testconfig: &TestConfig) -> Vec { let mut from_pools = vec![]; let spam = testconfig .spam diff --git a/crates/core/src/generator/mod.rs b/crates/core/src/generator/mod.rs index b8a4eb1..e0d0200 100644 --- a/crates/core/src/generator/mod.rs +++ b/crates/core/src/generator/mod.rs @@ -203,27 +203,41 @@ where } PlanType::Setup(on_setup_step) => { let setup_steps = conf.get_setup_steps()?; - for step in setup_steps.iter() { - // lookup placeholders in DB & update map before templating - templater.find_fncall_placeholders(step, db, &mut placeholder_map)?; - - // setup tx with template values - let tx = NamedTxRequest::new( - templater.template_function_call( - &self.make_strict_call(step, 0)?, - &placeholder_map, - )?, - None, - step.kind.to_owned(), - ); - let handle = on_setup_step(tx.to_owned())?; - if let Some(handle) = handle { - handle.await.map_err(|e| { - ContenderError::with_err(e, "join error; callback crashed") - })?; + let agents = self.get_agent_store(); + let num_accts = agents + .all_agents() + .next() + .map(|(_, store)| store.signers.len()) + .unwrap_or(1); + + for i in 0..(num_accts) { + for step in setup_steps.iter() { + if i > 0 && step.from_pool.is_none() { + // only loop on from_pool steps; single-account steps can't be repeated + continue; + } + // lookup placeholders in DB & update map before templating + templater.find_fncall_placeholders(step, db, &mut placeholder_map)?; + + // setup tx with template values + let tx = NamedTxRequest::new( + templater.template_function_call( + &self.make_strict_call(step, i)?, // 'from' address injected here + &placeholder_map, + )?, + None, + step.kind.to_owned(), + ); + + let handle = on_setup_step(tx.to_owned())?; + if let Some(handle) = handle { + handle.await.map_err(|e| { + ContenderError::with_err(e, "join error; callback crashed") + })?; + } + txs.push(tx.into()); } - txs.push(tx.into()); } } PlanType::Spam(num_txs, on_spam_setup) => { @@ -281,7 +295,7 @@ where for step in spam_steps.iter() { // converts a FunctionCallDefinition to a NamedTxRequest (filling in fuzzable args), // returns a callback handle and the processed tx request - let process_tx = |req| { + let prepare_tx = |req| { let args = get_fuzzed_args(req, &canonical_fuzz_map, i); let fuzz_tx_value = get_fuzzed_tx_value(req, &canonical_fuzz_map, i); let mut req = req.to_owned(); @@ -304,7 +318,7 @@ where match step { SpamRequest::Tx(req) => { - let (handle, tx) = process_tx(req)?; + let (handle, tx) = prepare_tx(req)?; if let Some(handle) = handle { handle.await.map_err(|e| { ContenderError::with_err(e, "error from callback") @@ -315,7 +329,7 @@ where SpamRequest::Bundle(req) => { let mut bundle_txs = vec![]; for tx in req.txs.iter() { - let (handle, txr) = process_tx(tx)?; + let (handle, txr) = prepare_tx(tx)?; if let Some(handle) = handle { handle.await.map_err(|e| { ContenderError::with_err(e, "error from callback") diff --git a/crates/core/src/test_scenario.rs b/crates/core/src/test_scenario.rs index 49c8810..fd32a36 100644 --- a/crates/core/src/test_scenario.rs +++ b/crates/core/src/test_scenario.rs @@ -583,8 +583,9 @@ where #[cfg(test)] pub mod tests { - use crate::agent_controller::AgentStore; + use crate::agent_controller::{AgentStore, SignerStore}; use crate::db::MockDb; + use crate::generator::named_txs::ExecutionRequest; use crate::generator::templater::Templater; use crate::generator::types::{ CreateDefinition, FunctionCallDefinition, FuzzParam, SpamRequest, @@ -595,8 +596,11 @@ pub mod tests { use crate::test_scenario::TestScenario; use crate::Result; use alloy::hex::ToHexExt; + use alloy::network::{Ethereum, EthereumWallet, TransactionBuilder}; use alloy::node_bindings::AnvilInstance; - use alloy::primitives::Address; + use alloy::primitives::{Address, U256}; + use alloy::providers::{Provider, ProviderBuilder}; + use alloy::rpc::types::TransactionRequest; use std::collections::HashMap; pub struct MockConfig; @@ -654,6 +658,16 @@ pub mod tests { fuzz: None, kind: None, }, + FunctionCallDefinition { + to: "0x7a250d5630B4cF539739dF2C5dAcb4c659F2488D".to_owned(), + from: None, + from_pool: Some("pool1".to_owned()), + value: None, + signature: "increment()".to_owned(), + args: vec![].into(), + fuzz: None, + kind: None, + }, ]) } @@ -723,6 +737,39 @@ pub mod tests { ) -> TestScenario { let seed = RandSeed::seed_from_bytes(&[0x01; 32]); let signers = &get_test_signers(); + let provider = ProviderBuilder::new() + .network::() + .on_http(anvil.endpoint_url()); + let mut agents = AgentStore::new(); + let pool1 = + SignerStore::new_random(10, &RandSeed::seed_from_str("0x0defa117"), "0x0defa117"); + let pool_signers = pool1.signers.to_vec(); + let admin = &signers[0]; + agents.add_agent("pool1", pool1); + let mut nonce = provider + .get_transaction_count(admin.address()) + .await + .unwrap(); + let chain_id = anvil.chain_id(); + for signer in &pool_signers { + let tx = TransactionRequest::default() + .with_from(signers[0].address()) + .with_to(signer.address()) + .with_value(U256::from(100_000_000_000_000_000u128)) + .with_nonce(nonce) + .with_max_fee_per_gas(10_000_000_000_u128) + .with_max_priority_fee_per_gas(1_000_000_000_u128) + .with_gas_limit(21000) + .with_chain_id(chain_id); + nonce += 1; + let signed_tx = tx + .build::(&admin.to_owned().into()) + .await + .unwrap(); + + let res = provider.send_tx_envelope(signed_tx).await.unwrap(); + res.get_receipt().await.unwrap(); + } TestScenario::new( MockConfig, @@ -731,7 +778,7 @@ pub mod tests { None, seed.to_owned(), signers, - AgentStore::new(), + agents, ) .await .unwrap() @@ -758,7 +805,7 @@ pub mod tests { })) .await .unwrap(); - assert_eq!(setup_txs.len(), 2); + assert_eq!(setup_txs.len(), 12); let spam_txs = scenario .load_txs(PlanType::Spam(20, |tx| { @@ -770,6 +817,41 @@ pub mod tests { assert!(spam_txs.len() >= 20); } + #[tokio::test] + async fn setup_steps_use_agent_signers() { + let anvil = spawn_anvil(); + let mut scenario = get_test_scenario(&anvil).await; + scenario.deploy_contracts().await.unwrap(); + let setup_steps = scenario + .load_txs(PlanType::Setup(|_| Ok(None))) + .await + .unwrap(); + scenario.run_setup().await.unwrap(); + let mut used_agent_keys = 0; + for step in setup_steps { + let tx = match step { + ExecutionRequest::Tx(tx) => tx, + _ => continue, + }; + let from = tx.tx.from.unwrap(); + assert!(scenario.wallet_map.contains_key(&from)); + assert!(scenario.agent_store.has_agent("pool1")); + if scenario + .agent_store + .get_agent("pool1") + .unwrap() + .signers + .iter() + .map(|s| s.address()) + .collect::>() + .contains(&from) + { + used_agent_keys += 1; + } + } + assert!(used_agent_keys > 1); + } + #[tokio::test] async fn scenario_creates_contracts() { let anvil = spawn_anvil(); diff --git a/scenarios/univ2ConfigTest.toml b/scenarios/univ2ConfigTest.toml index 977d899..217a5b3 100644 --- a/scenarios/univ2ConfigTest.toml +++ b/scenarios/univ2ConfigTest.toml @@ -41,19 +41,18 @@ value = "10000000000000000000" [[setup]] to = "{weth}" -from = "0x70997970C51812dc3A010C7d01b50e0d17dc79C8" +from_pool = "pool1" signature = "function deposit() public payable" value = "10000000000000000000" [[setup]] to = "{weth}" -from = "0x3C44CdDdB6a900fa2b585dd299e03d12FA4293BC" +from_pool = "pool2" signature = "function deposit() public payable" value = "10000000000000000000" ## uniV2 pair: weth/token ###################################################### - [[setup]] to = "{uniV2Factory}" from = "0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266" @@ -87,7 +86,7 @@ args = [ # contender1 [[setup]] to = "{weth}" -from = "0x3C44CdDdB6a900fa2b585dd299e03d12FA4293BC" +from_pool = "pool1" signature = "approve(address spender, uint256 amount) returns (bool)" args = [ "{uniRouterV2}", @@ -96,7 +95,7 @@ args = [ [[setup]] to = "{testToken}" -from = "0x3C44CdDdB6a900fa2b585dd299e03d12FA4293BC" +from_pool = "pool1" signature = "approve(address spender, uint256 amount) returns (bool)" args = [ "{uniRouterV2}", @@ -106,7 +105,7 @@ args = [ # contender2 [[setup]] to = "{weth}" -from = "0x70997970C51812dc3A010C7d01b50e0d17dc79C8" +from_pool = "pool2" signature = "approve(address spender, uint256 amount) returns (bool)" args = [ "{uniRouterV2}", @@ -115,7 +114,7 @@ args = [ [[setup]] to = "{testToken}" -from = "0x70997970C51812dc3A010C7d01b50e0d17dc79C8" +from_pool = "pool2" signature = "approve(address spender, uint256 amount) returns (bool)" args = [ "{uniRouterV2}", @@ -147,7 +146,7 @@ args = [ [spam.tx] kind = "uniswap_v2" to = "{uniRouterV2}" -from = "0x70997970C51812dc3A010C7d01b50e0d17dc79C8" +from_pool = "pool1" signature = "swapExactTokensForTokens(uint256 amountIn, uint256 amountOutMin, address[] path, address to, uint256 deadline) external returns (uint256[] memory)" args = [ "1000000000000000000", @@ -163,7 +162,7 @@ fuzz = [{param = "amountIn", min = "1", max = "100000000000000000"}] [spam.tx] kind = "uniswap_v2" to = "{uniRouterV2}" -from = "0x3C44CdDdB6a900fa2b585dd299e03d12FA4293BC" +from_pool = "pool2" signature = "swapExactTokensForTokens(uint256 amountIn, uint256 amountOutMin, address[] path, address to, uint256 deadline) external returns (uint256[] memory)" args = [ "1000000000000000000",