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, 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(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.call_kind = 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                    self.transaction.function_abi = Some(function.full_signature());
72                    self.transaction.display_function = Some(function.name.clone());
73
74                    let values = function.abi_decode_input(data).inspect_err(|_| {
75                        error!(
76                            contract=?self.transaction.contract_name,
77                            signature=?function,
78                            data=hex::encode(data),
79                            "Failed to decode function arguments",
80                        );
81                    })?;
82                    self.transaction.arguments =
83                        Some(values.iter().map(format_token_raw).collect());
84                }
85            }
86        }
87
88        Ok(())
89    }
90
91    /// Populate the transaction as CREATE tx
92    ///
93    /// If this is a CREATE2 transaction this attempt to decode the arguments from the CREATE2
94    /// deployer's function
95    pub fn set_create(
96        &mut self,
97        is_create2: bool,
98        address: Address,
99        contracts: &BTreeMap<Address, &ContractData>,
100    ) -> Result<()> {
101        if is_create2 {
102            self.transaction.call_kind = CallKind::Create2;
103        } else {
104            self.transaction.call_kind = CallKind::Create;
105        }
106
107        let info = contracts.get(&address);
108        self.transaction.contract_name = info.map(|info| info.name.clone());
109        self.transaction.contract_address = Some(address);
110
111        let Some(data) = self.transaction.transaction.input() else { return Ok(()) };
112        let Some(info) = info else { return Ok(()) };
113        let Some(bytecode) = info.bytecode() else { return Ok(()) };
114
115        // `create2` transactions are prefixed by a 32 byte salt.
116        let creation_code = if is_create2 {
117            if data.len() < 32 {
118                return Ok(());
119            }
120            &data[32..]
121        } else {
122            data
123        };
124
125        // The constructor args start after bytecode.
126        let contains_constructor_args = creation_code.len() > bytecode.len();
127        if !contains_constructor_args {
128            return Ok(());
129        }
130        let constructor_args = &creation_code[bytecode.len()..];
131
132        let Some(constructor) = info.abi.constructor() else { return Ok(()) };
133        let values = constructor.abi_decode_input(constructor_args).inspect_err(|_| {
134                error!(
135                    contract=?self.transaction.contract_name,
136                    signature=%format!("constructor({})", constructor.inputs.iter().map(|p| &p.ty).format(",")),
137                    is_create2,
138                    constructor_args=%hex::encode(constructor_args),
139                    "Failed to decode constructor arguments",
140                );
141                debug!(full_data=%hex::encode(data), bytecode=%hex::encode(creation_code));
142            })?;
143        self.transaction.arguments = Some(values.iter().map(format_token_raw).collect());
144
145        Ok(())
146    }
147
148    /// Populates additional data from the transaction execution result.
149    pub fn with_execution_result(
150        mut self,
151        result: &ScriptResult<N>,
152        gas_estimate_multiplier: u64,
153        linked_build_data: &LinkedBuildData,
154    ) -> Self {
155        let mut created_contracts =
156            result.get_created_contracts(&linked_build_data.known_contracts);
157
158        // Add the additional contracts created in this transaction, so we can verify them later.
159        created_contracts.retain(|contract| {
160            // Filter out the contract that was created by the transaction itself.
161            self.transaction.contract_address != Some(contract.address)
162        });
163
164        self.transaction.additional_contracts = created_contracts;
165
166        if !self.transaction.is_fixed_gas_limit
167            && let Some(unsigned) = self.transaction.transaction.as_unsigned_mut()
168        {
169            // We inflate the gas used by the user specified percentage
170            unsigned.set_gas_limit(result.gas_used * gas_estimate_multiplier / 100);
171        }
172
173        self
174    }
175
176    pub fn build(self) -> TransactionWithMetadata<N> {
177        self.transaction
178    }
179}
180
181impl<N: Network> From<TransactionWithMetadata<N>> for ScriptTransactionBuilder<N> {
182    fn from(transaction: TransactionWithMetadata<N>) -> Self {
183        Self { transaction }
184    }
185}