fix(forge/script): deployment size check (#3202)

* fix(forge/script): deployment size check

* chore: clippy

* chore: fmt

* add comments

* fix: prompt user out of the loop

* fix: bail instead of panicking

cli/src/cmd/forge/script/cmd.rs

@@ -165,10 +165,7 @@ impl ScriptArgs {
 
                 verify.known_contracts = unwrap_contracts(&highlevel_known_contracts, false);
 
-                self.check_contract_sizes(
-                    result.transactions.as_ref(),
-                    &highlevel_known_contracts,
-                )?;
+                self.check_contract_sizes(&result, &highlevel_known_contracts)?;
 
                 self.handle_broadcastable_transactions(
                     &target,

cli/src/cmd/forge/script/mod.rs

@@ -26,8 +26,10 @@ use forge::{
     executor::{opts::EvmOpts, Backend},
     trace::{
         identifier::{EtherscanIdentifier, LocalTraceIdentifier, SignaturesIdentifier},
-        CallTraceArena, CallTraceDecoder, CallTraceDecoderBuilder, TraceKind,
+        CallTraceArena, CallTraceDecoder, CallTraceDecoderBuilder, RawOrDecodedCall,
+        RawOrDecodedReturnData, TraceKind,
     },
+    CallKind,
 };
 use foundry_common::{evm::EvmArgs, ContractsByArtifact, CONTRACT_MAX_SIZE, SELECTOR_LEN};
 use foundry_config::{figment, Config};
@@ -451,10 +453,50 @@ impl ScriptArgs {
     /// the user.
     fn check_contract_sizes(
         &self,
-        transactions: Option<&VecDeque<TypedTransaction>>,
+        result: &ScriptResult,
         known_contracts: &BTreeMap<ArtifactId, ContractBytecodeSome>,
     ) -> eyre::Result<()> {
-        for (data, to) in transactions.iter().flat_map(|txes| {
+        // (name, &init, &deployed)[]
+        let mut bytecodes: Vec<(String, &[u8], &[u8])> = vec![];
+
+        // From artifacts
+        known_contracts.iter().for_each(|(artifact, bytecode)| {
+            const MSG: &str = "Code should have been linked before.";
+            let init_code = bytecode.bytecode.object.as_bytes().expect(MSG);
+            // Ignore abstract contracts
+            if let Some(ref deployed_code) = bytecode.deployed_bytecode.bytecode {
+                let deployed_code = deployed_code.object.as_bytes().expect(MSG);
+                bytecodes.push((artifact.name.clone(), init_code, deployed_code));
+            }
+        });
+
+        // From traces
+        let create_nodes = result.traces.iter().flat_map(|(_, traces)| {
+            traces
+                .arena
+                .iter()
+                .filter(|node| matches!(node.kind(), CallKind::Create | CallKind::Create2))
+        });
+        let mut unknown_c = 0usize;
+        for node in create_nodes {
+            // Calldata == init code
+            if let RawOrDecodedCall::Raw(ref init_code) = node.trace.data {
+                // Output is the runtime code
+                if let RawOrDecodedReturnData::Raw(ref deployed_code) = node.trace.output {
+                    // Only push if it was not present already
+                    if !bytecodes.iter().any(|(_, b, _)| b == init_code) {
+                        bytecodes.push((format!("Unknown{}", unknown_c), init_code, deployed_code));
+                        unknown_c += 1;
+                    }
+                    continue
+                }
+            }
+            // Both should be raw and not decoded since it's just bytecode
+            eyre::bail!("Create node returned decoded data: {:?}", node);
+        }
+
+        let mut prompt_user = false;
+        for (data, to) in result.transactions.iter().flat_map(|txes| {
             txes.iter().filter_map(|tx| {
                 tx.data().filter(|data| data.len() > CONTRACT_MAX_SIZE).map(|data| (data, tx.to()))
             })
@@ -472,40 +514,30 @@ impl ScriptArgs {
             }
 
             // Find artifact with a deployment code same as the data.
-            if let Some((artifact, bytecode)) =
-                known_contracts.iter().find(|(_, bytecode)| {
-                    bytecode
-                        .bytecode
-                        .object
-                        .as_bytes()
-                        .expect("Code should have been linked before.") ==
-                        &data[offset..]
-                })
+            if let Some((name, _, deployed_code)) =
+                bytecodes.iter().find(|(_, init_code, _)| *init_code == &data[offset..])
             {
-                // Find the deployed code size of the artifact.
-                if let Some(deployed_bytecode) = &bytecode.deployed_bytecode.bytecode {
-                    let deployment_size = deployed_bytecode.object.bytes_len();
-
-                    if deployment_size > CONTRACT_MAX_SIZE {
-                        println!(
-                            "{}",
-                            Paint::red(format!(
-                                "`{}` is above the contract size limit ({} vs {}).",
-                                artifact.name, deployment_size, CONTRACT_MAX_SIZE
-                            ))
-                        );
-
-                        if self.broadcast &&
-                            !Confirm::new()
-                                .with_prompt("Do you wish to continue?".to_string())
-                                .interact()?
-                        {
-                            eyre::bail!("User canceled the script.");
-                        }
-                    }
+                let deployment_size = deployed_code.len();
+
+                if deployment_size > CONTRACT_MAX_SIZE {
+                    prompt_user = self.broadcast;
+                    println!(
+                        "{}",
+                        Paint::red(format!(
+                            "`{}` is above the EIP-170 contract size limit ({} > {}).",
+                            name, deployment_size, CONTRACT_MAX_SIZE
+                        ))
+                    );
                 }
             }
         }
+
+        if prompt_user &&
+            !Confirm::new().with_prompt("Do you wish to continue?".to_string()).interact()?
+        {
+            eyre::bail!("User canceled the script.");
+        }
+
         Ok(())
     }
 }

cli/tests/it/create.rs

@@ -148,7 +148,7 @@ forgetest_async!(
         let (_api, handle) = spawn(NodeConfig::test()).await;
         let rpc = handle.http_endpoint();
         let wallet = handle.dev_wallets().next().unwrap();
-        let pk = hex::encode(&wallet.signer().to_bytes());
+        let pk = hex::encode(wallet.signer().to_bytes());
         cmd.args(["init", "--force"]);
         cmd.assert_non_empty_stdout();
 

cli/tests/it/script.rs

@@ -685,8 +685,7 @@ contract Script0 is Script {
 
         let run_latest = foundry_common::fs::json_files(prj.root().join("broadcast"))
             .into_iter()
-            .filter(|file| file.ends_with("run-latest.json"))
-            .next()
+            .find(|file| file.ends_with("run-latest.json"))
             .expect("No broadcast artifacts");
 
         let content = foundry_common::fs::read_to_string(run_latest).unwrap();
@@ -775,8 +774,7 @@ contract Script0 is Script {
 
         let run_latest = foundry_common::fs::json_files(prj.root().join("broadcast"))
             .into_iter()
-            .filter(|file| file.ends_with("run-latest.json"))
-            .next()
+            .find(|file| file.ends_with("run-latest.json"))
             .expect("No broadcast artifacts");
 
         let content = foundry_common::fs::read_to_string(run_latest).unwrap();

cli/tests/it/verify.rs

@@ -11,7 +11,7 @@ use foundry_config::Config;
 use foundry_utils::Retry;
 use std::{fs, path::PathBuf};
 
-const VERIFICATION_PROVIDERS: &'static [&'static str] = &["etherscan", "sourcify"];
+const VERIFICATION_PROVIDERS: &[&str] = &["etherscan", "sourcify"];
 
 /// Adds a `Unique` contract to the source directory of the project that can be imported as
 /// `import {Unique} from "./unique.sol";`
@@ -366,7 +366,7 @@ forgetest_async!(
         assert!(stdout.contains("Sending transactions [0 - 4]"), "{}", err);
 
         // ensure all transactions are successful
-        assert_eq!(5, stdout.matches("✅").count(), "{}", err);
+        assert_eq!(5, stdout.matches('✅').count(), "{}", err);
 
         // ensure verified all deployments
         // Note: the 5th tx creates contracts internally, which are little flaky at times because