Merge pull request #39 from flashbots/brock/pr30-mods

brock/pr30 mods

cli/src/main.rs

@@ -1,6 +1,7 @@
 mod commands;
 
 use alloy::{
+    network::AnyNetwork,
     primitives::{
         utils::{format_ether, parse_ether},
         Address, U256,
@@ -13,7 +14,7 @@ use commands::{ContenderCli, ContenderSubcommand};
 use contender_core::{
     db::{DbOps, RunTx},
     generator::{
-        types::{FunctionCallDefinition, RpcProvider},
+        types::{AnyProvider, FunctionCallDefinition},
         RandSeed,
     },
     spammer::{BlockwiseSpammer, LogCallback, NilCallback, TimedSpammer},
@@ -63,7 +64,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             min_balance,
         } => {
             let url = Url::parse(rpc_url.as_ref()).expect("Invalid RPC URL");
-            let rpc_client = ProviderBuilder::new().on_http(url.to_owned());
+            let rpc_client = ProviderBuilder::new()
+                .network::<AnyNetwork>()
+                .on_http(url.to_owned());
             let testconfig: TestConfig = TestConfig::from_file(&testfile)?;
             let min_balance = parse_ether(&min_balance)?;
 
@@ -101,7 +104,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             let testconfig = TestConfig::from_file(&testfile)?;
             let rand_seed = seed.map(|s| RandSeed::from_str(&s)).unwrap_or_default();
             let url = Url::parse(rpc_url.as_ref()).expect("Invalid RPC URL");
-            let rpc_client = ProviderBuilder::new().on_http(url.to_owned());
+            let rpc_client = ProviderBuilder::new()
+                .network::<AnyNetwork>()
+                .on_http(url.to_owned());
             let duration = duration.unwrap_or_default();
             let min_balance = parse_ether(&min_balance)?;
 
@@ -223,7 +228,7 @@ const DEFAULT_PRV_KEYS: [&str; 10] = [
 
 async fn spam_callback_default(
     log_txs: bool,
-    rpc_client: Option<Arc<RpcProvider>>,
+    rpc_client: Option<Arc<AnyProvider>>,
 ) -> SpamCallbackType {
     if let Some(rpc_client) = rpc_client {
         if log_txs {
@@ -236,7 +241,7 @@ async fn spam_callback_default(
 async fn check_balances(
     prv_keys: &[PrivateKeySigner],
     min_balance: U256,
-    rpc_client: &RpcProvider,
+    rpc_client: &AnyProvider,
 ) {
     for prv_key in prv_keys {
         let address = prv_key.address();

cli/univ2ConfigTest.toml

@@ -141,6 +141,7 @@ args = [
 
 ### the spam step will be repeated
 [[spam]]
+kind = "uniswap_v2"
 to = "{uniRouterV2}"
 from = "0x70997970C51812dc3A010C7d01b50e0d17dc79C8"
 signature = "swapExactTokensForTokens(uint256 amountIn, uint256 amountOutMin, address[] path, address to, uint256 deadline) external returns (uint256[] memory)"
@@ -158,6 +159,7 @@ min = "1"
 max = "100000000000000000"
 
 [[spam]]
+kind = "uniswap_v2"
 to = "{uniRouterV2}"
 from = "0x3C44CdDdB6a900fa2b585dd299e03d12FA4293BC"
 signature = "swapExactTokensForTokens(uint256 amountIn, uint256 amountOutMin, address[] path, address to, uint256 deadline) external returns (uint256[] memory)"
@@ -172,4 +174,4 @@ args = [
 [[spam.fuzz]]
 param = "amountIn"
 min = "100000"
-max = "10000000000"
+max = "10000000000"

sqlite_db/src/lib.rs

@@ -95,6 +95,7 @@ struct RunTxRow {
     end_timestamp: usize,
     block_number: u64,
     gas_used: String,
+    kind: Option<String>,
 }
 
 impl RunTxRow {
@@ -106,6 +107,7 @@ impl RunTxRow {
             end_timestamp: row.get(3)?,
             block_number: row.get(4)?,
             gas_used: row.get(5)?,
+            kind: row.get(6)?,
         })
     }
 }
@@ -119,6 +121,7 @@ impl From<RunTxRow> for RunTx {
             end_timestamp: row.end_timestamp,
             block_number: row.block_number,
             gas_used: row.gas_used.parse().expect("invalid gas_used parameter"),
+            kind: row.kind,
         }
     }
 }
@@ -151,6 +154,7 @@ impl DbOps for SqliteDb {
                 end_timestamp INTEGER NOT NULL,
                 block_number INTEGER NOT NULL,
                 gas_used TEXT NOT NULL,
+                kind TEXT,
                 FOREIGN KEY(run_id) REFERENCES runs(runid)
             )",
             params![],
@@ -178,7 +182,7 @@ impl DbOps for SqliteDb {
     fn get_run_txs(&self, run_id: u64) -> Result<Vec<RunTx>> {
         let pool = self.get_pool()?;
         let mut stmt = pool
-            .prepare("SELECT run_id, tx_hash, start_timestamp, end_timestamp, block_number, gas_used FROM run_txs WHERE run_id = ?1")
+            .prepare("SELECT run_id, tx_hash, start_timestamp, end_timestamp, block_number, gas_used, kind FROM run_txs WHERE run_id = ?1")
             .map_err(|e| ContenderError::with_err(e, "failed to prepare statement"))?;
 
         let rows = stmt
@@ -233,17 +237,39 @@ impl DbOps for SqliteDb {
         Ok(res.into())
     }
 
+    fn get_named_tx_by_address(&self, address: &Address) -> Result<NamedTx> {
+        let pool = self.get_pool()?;
+        let mut stmt = pool
+            .prepare(
+                "SELECT name, tx_hash, contract_address FROM named_txs WHERE contract_address = ?1 ORDER BY id DESC LIMIT 1",
+            )
+            .map_err(|e| ContenderError::with_err(e, "failed to prepare statement"))?;
+
+        let row = stmt
+            .query_map(params![address.encode_hex()], |row| {
+                NamedTxRow::from_row(row)
+            })
+            .map_err(|e| ContenderError::with_err(e, "failed to map row"))?;
+        let res = row
+            .last()
+            .transpose()
+            .map_err(|e| ContenderError::with_err(e, "no row found"))?
+            .ok_or(ContenderError::DbError("no existing row", None))?;
+        Ok(res.into())
+    }
+
     fn insert_run_txs(&self, run_id: u64, run_txs: Vec<RunTx>) -> Result<()> {
         let pool = self.get_pool()?;
         let stmts = run_txs.iter().map(|tx| {
             format!(
-                "INSERT INTO run_txs (run_id, tx_hash, start_timestamp, end_timestamp, block_number, gas_used) VALUES ({}, '{}', {}, {}, {}, '{}');",
+                "INSERT INTO run_txs (run_id, tx_hash, start_timestamp, end_timestamp, block_number, gas_used, kind) VALUES ({}, '{}', {}, {}, {}, '{}', '{}');",
                 run_id,
                 tx.tx_hash.encode_hex(),
                 tx.start_timestamp,
                 tx.end_timestamp,
                 tx.block_number,
-                tx.gas_used.to_string()
+                tx.gas_used.to_string(),
+                tx.kind.to_owned().unwrap_or("NULL".to_string()),
             )
         });
         pool.execute_batch(&format!(
@@ -322,13 +348,15 @@ mod tests {
                 end_timestamp: 200,
                 block_number: 1,
                 gas_used: 100,
+                kind: Some("test".to_string()),
             },
             RunTx {
                 tx_hash: TxHash::from_slice(&[1u8; 32]),
                 start_timestamp: 200,
                 end_timestamp: 300,
                 block_number: 2,
                 gas_used: 200,
+                kind: Some("test".to_string()),
             },
         ];
         db.insert_run_txs(run_id as u64, run_txs).unwrap();

src/db/mod.rs

@@ -12,6 +12,7 @@ pub struct RunTx {
     pub end_timestamp: usize,
     pub block_number: u64,
     pub gas_used: u128,
+    pub kind: Option<String>,
 }
 
 #[derive(Debug, Serialize, Clone)]
@@ -49,6 +50,8 @@ pub trait DbOps {
 
     fn get_named_tx(&self, name: &str) -> Result<NamedTx>;
 
+    fn get_named_tx_by_address(&self, address: &Address) -> Result<NamedTx>;
+
     fn insert_run_txs(&self, run_id: u64, run_txs: Vec<RunTx>) -> Result<()>;
 
     fn get_run_txs(&self, run_id: u64) -> Result<Vec<RunTx>>;
@@ -77,6 +80,14 @@ impl DbOps for MockDb {
         Ok(NamedTx::new(String::default(), TxHash::default(), None))
     }
 
+    fn get_named_tx_by_address(&self, address: &Address) -> Result<NamedTx> {
+        Ok(NamedTx::new(
+            String::default(),
+            TxHash::default(),
+            Some(*address),
+        ))
+    }
+
     fn insert_run_txs(&self, _run_id: u64, _run_txs: Vec<RunTx>) -> Result<()> {
         Ok(())
     }

src/generator/mod.rs

@@ -27,18 +27,71 @@ pub mod types;
 /// Utility functions used in the generator module.
 pub mod util;
 
-impl NamedTxRequest {
-    pub fn with_name(name: &str, tx: TransactionRequest) -> Self {
+/// Syntactical sugar for creating a [`NamedTxRequest`].
+///
+/// This is useful for imperatively assigning optional fields to a tx.
+/// It is _not_ useful when you're dynamically assigning these fields (i.e. you have an Option to check first).
+///
+/// ### Example:
+/// ```
+/// use alloy::rpc::types::TransactionRequest;
+/// # use contender_core::generator::NamedTxRequestBuilder;
+///
+/// let tx_req = TransactionRequest::default();
+/// let named_tx_req = NamedTxRequestBuilder::new(tx_req)
+///     .with_name("unique_tx_name")
+///     .with_kind("tx_kind")
+///     .build();
+/// assert_eq!(named_tx_req.name, Some("unique_tx_name".to_owned()));
+/// assert_eq!(named_tx_req.kind, Some("tx_kind".to_owned()));
+/// ```
+pub struct NamedTxRequestBuilder {
+    name: Option<String>,
+    kind: Option<String>,
+    tx: TransactionRequest,
+}
+
+impl NamedTxRequestBuilder {
+    pub fn new(tx: TransactionRequest) -> Self {
         Self {
-            name: Some(name.to_string()),
+            name: None,
+            kind: None,
             tx,
         }
     }
+
+    pub fn with_name(&mut self, name: &str) -> &mut Self {
+        self.name = Some(name.to_owned());
+        self
+    }
+
+    pub fn with_kind(&mut self, kind: &str) -> &mut Self {
+        self.kind = Some(kind.to_owned());
+        self
+    }
+
+    pub fn build(&self) -> NamedTxRequest {
+        NamedTxRequest::new(
+            self.tx.to_owned(),
+            self.name.to_owned(),
+            self.kind.to_owned(),
+        )
+    }
+}
+
+impl NamedTxRequest {
+    pub fn new(tx: TransactionRequest, name: Option<String>, kind: Option<String>) -> Self {
+        Self { name, kind, tx }
+    }
 }
 
 impl From<TransactionRequest> for NamedTxRequest {
     fn from(tx: TransactionRequest) -> Self {
-        Self { name: None, tx }
+        Self {
+            name: None,
+            kind: None,
+            tx,
+        }
     }
 }
 
@@ -103,11 +156,13 @@ where
                     // lookup placeholder values in DB & update map before templating
                     templater.find_placeholder_values(&step.bytecode, &mut placeholder_map, db)?;
 
-                    // create txs with template values
-                    let tx = NamedTxRequest::with_name(
-                        &step.name,
+                    // create tx with template values
+                    let tx = NamedTxRequestBuilder::new(
                         templater.template_contract_deploy(step, &placeholder_map)?,
-                    );
+                    )
+                    .with_name(&step.name)
+                    .build();
+
                     let handle = on_create_step(tx.to_owned())?;
                     if let Some(handle) = handle {
                         handle.await.map_err(|e| {
@@ -123,10 +178,11 @@ where
                     // lookup placeholders in DB & update map before templating
                     templater.find_fncall_placeholders(step, db, &mut placeholder_map)?;
 
-                    // create txs with template values
+                    // setup tx with template values
                     let tx: NamedTxRequest = templater
                         .template_function_call(step, &placeholder_map)?
                         .into();
+
                     let handle = on_setup_step(tx.to_owned())?;
                     if let Some(handle) = handle {
                         handle.await.map_err(|e| {
@@ -188,9 +244,12 @@ where
                         let mut step = step.to_owned();
                         step.args = Some(args);
 
-                        let tx: NamedTxRequest = templater
-                            .template_function_call(&step, &placeholder_map)?
-                            .into();
+                        let tx = NamedTxRequest::new(
+                            templater.template_function_call(&step, &placeholder_map)?,
+                            None,
+                            step.kind,
+                        );
+
                         let handle = on_spam_setup(tx.to_owned())?;
                         if let Some(handle) = handle {
                             handle

src/generator/types.rs

@@ -1,4 +1,5 @@
 use alloy::{
+    network::AnyNetwork,
     primitives::U256,
     providers::RootProvider,
     rpc::types::TransactionRequest,
@@ -8,11 +9,13 @@ use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use tokio::task::JoinHandle;
 
-pub type RpcProvider = RootProvider<Http<Client>>;
+pub type EthProvider = RootProvider<Http<Client>>;
+pub type AnyProvider = RootProvider<Http<Client>, AnyNetwork>;
 
 #[derive(Clone, Debug)]
 pub struct NamedTxRequest {
     pub name: Option<String>,
+    pub kind: Option<String>,
     pub tx: TransactionRequest,
 }
 
@@ -30,6 +33,8 @@ pub struct FunctionCallDefinition {
     pub value: Option<String>,
     /// Parameters to fuzz during the test.
     pub fuzz: Option<Vec<FuzzParam>>,
+    /// Optional type of the spam transaction for categorization.
+    pub kind: Option<String>
 }
 
 #[derive(Clone, Deserialize, Debug, Serialize)]