Skip to main content

forge_script/
transaction.rs

1use super::ScriptResult;
2use crate::build::LinkedBuildData;
3use alloy_dyn_abi::JsonAbiExt;
4use alloy_network::{Network, TransactionBuilder};
5use alloy_primitives::{Address, B256, TxKind, hex};
6use eyre::Result;
7use forge_script_sequence::TransactionWithMetadata;
8use foundry_common::{ContractData, SELECTOR_LEN, TransactionMaybeSigned, fmt::format_token_raw};
9use foundry_evm::traces::CallTraceDecoder;
10use itertools::Itertools;
11use revm_inspectors::tracing::types::CallKind;
12use std::collections::BTreeMap;
13
14#[derive(Debug)]
15pub struct ScriptTransactionBuilder<N: Network> {
16    transaction: TransactionWithMetadata<N>,
17}
18
19impl<N: Network> ScriptTransactionBuilder<N> {
20    pub fn new(transaction: TransactionMaybeSigned<N>, rpc: String) -> Self {
21        let mut transaction = TransactionWithMetadata::from_tx_request(transaction);
22        transaction.rpc = rpc;
23        // If tx.gas is already set that means it was specified in script
24        transaction.is_fixed_gas_limit = transaction.tx().gas().is_some();
25
26        Self { transaction }
27    }
28
29    /// Populate the transaction as CALL tx
30    pub fn set_call(
31        &mut self,
32        local_contracts: &BTreeMap<Address, &ContractData>,
33        decoder: &CallTraceDecoder,
34        create2_deployer: Address,
35    ) -> Result<()> {
36        if let Some(TxKind::Call(to)) = self.transaction.transaction.to() {
37            if to == create2_deployer {
38                if let Some(input) = self.transaction.transaction.input() {
39                    let (salt, init_code) = input.split_at(32);
40
41                    self.set_create(
42                        true,
43                        create2_deployer.create2_from_code(B256::from_slice(salt), init_code),
44                        local_contracts,
45                    )?;
46                }
47            } else {
48                self.transaction.opcode = CallKind::Call;
49                self.transaction.contract_address = Some(to);
50
51                let Some(data) = self.transaction.transaction.input() else { return Ok(()) };
52
53                if data.len() < SELECTOR_LEN {
54                    return Ok(());
55                }
56
57                let (selector, data) = data.split_at(SELECTOR_LEN);
58
59                let function = if let Some(info) = local_contracts.get(&to) {
60                    // This CALL is made to a local contract.
61                    self.transaction.contract_name = Some(info.name.clone());
62                    info.abi.functions().find(|function| function.selector() == selector)
63                } else {
64                    // This CALL is made to an external contract; try to decode it from the given
65                    // decoder.
66                    decoder.functions.get(selector).and_then(|v| v.first())
67                };
68
69                if let Some(function) = function {
70                    self.transaction.function = Some(function.signature());
71
72                    let values = function.abi_decode_input(data).inspect_err(|_| {
73                        error!(
74                            contract=?self.transaction.contract_name,
75                            signature=?function,
76                            data=hex::encode(data),
77                            "Failed to decode function arguments",
78                        );
79                    })?;
80                    self.transaction.arguments =
81                        Some(values.iter().map(format_token_raw).collect());
82                }
83            }
84        }
85
86        Ok(())
87    }
88
89    /// Populate the transaction as CREATE tx
90    ///
91    /// If this is a CREATE2 transaction this attempt to decode the arguments from the CREATE2
92    /// deployer's function
93    pub fn set_create(
94        &mut self,
95        is_create2: bool,
96        address: Address,
97        contracts: &BTreeMap<Address, &ContractData>,
98    ) -> Result<()> {
99        if is_create2 {
100            self.transaction.opcode = CallKind::Create2;
101        } else {
102            self.transaction.opcode = CallKind::Create;
103        }
104
105        let info = contracts.get(&address);
106        self.transaction.contract_name = info.map(|info| info.name.clone());
107        self.transaction.contract_address = Some(address);
108
109        let Some(data) = self.transaction.transaction.input() else { return Ok(()) };
110        let Some(info) = info else { return Ok(()) };
111        let Some(bytecode) = info.bytecode() else { return Ok(()) };
112
113        // `create2` transactions are prefixed by a 32 byte salt.
114        let creation_code = if is_create2 {
115            if data.len() < 32 {
116                return Ok(());
117            }
118            &data[32..]
119        } else {
120            data
121        };
122
123        // The constructor args start after bytecode.
124        let contains_constructor_args = creation_code.len() > bytecode.len();
125        if !contains_constructor_args {
126            return Ok(());
127        }
128        let constructor_args = &creation_code[bytecode.len()..];
129
130        let Some(constructor) = info.abi.constructor() else { return Ok(()) };
131        let values = constructor.abi_decode_input(constructor_args).inspect_err(|_| {
132                error!(
133                    contract=?self.transaction.contract_name,
134                    signature=%format!("constructor({})", constructor.inputs.iter().map(|p| &p.ty).format(",")),
135                    is_create2,
136                    constructor_args=%hex::encode(constructor_args),
137                    "Failed to decode constructor arguments",
138                );
139                debug!(full_data=%hex::encode(data), bytecode=%hex::encode(creation_code));
140            })?;
141        self.transaction.arguments = Some(values.iter().map(format_token_raw).collect());
142
143        Ok(())
144    }
145
146    /// Populates additional data from the transaction execution result.
147    pub fn with_execution_result(
148        mut self,
149        result: &ScriptResult,
150        gas_estimate_multiplier: u64,
151        linked_build_data: &LinkedBuildData,
152    ) -> Self {
153        let mut created_contracts =
154            result.get_created_contracts(&linked_build_data.known_contracts);
155
156        // Add the additional contracts created in this transaction, so we can verify them later.
157        created_contracts.retain(|contract| {
158            // Filter out the contract that was created by the transaction itself.
159            self.transaction.contract_address != Some(contract.address)
160        });
161
162        self.transaction.additional_contracts = created_contracts;
163
164        if !self.transaction.is_fixed_gas_limit
165            && let Some(unsigned) = self.transaction.transaction.as_unsigned_mut()
166        {
167            // We inflate the gas used by the user specified percentage
168            unsigned.set_gas_limit(result.gas_used * gas_estimate_multiplier / 100);
169        }
170
171        self
172    }
173
174    pub fn build(self) -> TransactionWithMetadata<N> {
175        self.transaction
176    }
177}
178
179impl<N: Network> From<TransactionWithMetadata<N>> for ScriptTransactionBuilder<N> {
180    fn from(transaction: TransactionWithMetadata<N>) -> Self {
181        Self { transaction }
182    }
183}