feat: smarter verification

crates/cheatcodes/src/fs.rs

@@ -356,9 +356,9 @@ fn get_artifact_code(state: &Cheatcodes, path: &str, deployed: bool) -> Result<B
             }?;
 
             let maybe_bytecode = if deployed {
-                artifact.1.deployed_bytecode.clone()
+                artifact.1.deployed_bytecode().cloned()
             } else {
-                artifact.1.bytecode.clone()
+                artifact.1.bytecode().cloned()
             };
 
             return maybe_bytecode

crates/common/src/contracts.rs

@@ -1,17 +1,28 @@
 //! Commonly used contract types and functions.
 
 use alloy_json_abi::{Event, Function, JsonAbi};
-use alloy_primitives::{Address, Bytes, Selector, B256};
+use alloy_primitives::{hex, Address, Bytes, Selector, B256};
 use eyre::Result;
 use foundry_compilers::{
-    artifacts::{CompactContractBytecode, ContractBytecodeSome},
+    artifacts::{
+        BytecodeObject, CompactBytecode, CompactContractBytecode, CompactDeployedBytecode,
+        ContractBytecodeSome, Offsets,
+    },
     ArtifactId,
 };
 use std::{
     collections::BTreeMap,
     ops::{Deref, DerefMut},
+    str::FromStr,
 };
 
+/// Libraries' runtime code always starts with the following instruction:
+/// `PUSH20 0x0000000000000000000000000000000000000000`
+///
+/// See: <https://docs.soliditylang.org/en/latest/contracts.html#call-protection-for-libraries>
+const CALL_PROTECTION_BYTECODE_PREFIX: [u8; 21] =
+    hex!("730000000000000000000000000000000000000000");
+
 /// Container for commonly used contract data.
 #[derive(Debug, Clone)]
 pub struct ContractData {
@@ -20,9 +31,21 @@ pub struct ContractData {
     /// Contract ABI.
     pub abi: JsonAbi,
     /// Contract creation code.
-    pub bytecode: Option<Bytes>,
+    pub bytecode: Option<CompactBytecode>,
     /// Contract runtime code.
-    pub deployed_bytecode: Option<Bytes>,
+    pub deployed_bytecode: Option<CompactDeployedBytecode>,
+}
+
+impl ContractData {
+    /// Returns reference to bytes of contract creation code, if present.
+    pub fn bytecode(&self) -> Option<&Bytes> {
+        self.bytecode.as_ref()?.bytes().filter(|b| !b.is_empty())
+    }
+
+    /// Returns reference to bytes of contract deployed code, if present.
+    pub fn deployed_bytecode(&self) -> Option<&Bytes> {
+        self.deployed_bytecode.as_ref()?.bytes().filter(|b| !b.is_empty())
+    }
 }
 
 type ArtifactWithContractRef<'a> = (&'a ArtifactId, &'a ContractData);
@@ -42,18 +65,10 @@ impl ContractsByArtifact {
                 .into_iter()
                 .filter_map(|(id, artifact)| {
                     let name = id.name.clone();
-                    let bytecode = artifact.bytecode.and_then(|b| b.into_bytes())?;
-                    let deployed_bytecode =
-                        artifact.deployed_bytecode.and_then(|b| b.into_bytes())?;
-
-                    // Exclude artifacts with present but empty bytecode. Such artifacts are usually
-                    // interfaces and abstract contracts.
-                    let bytecode = (bytecode.len() > 0).then_some(bytecode);
-                    let deployed_bytecode =
-                        (deployed_bytecode.len() > 0).then_some(deployed_bytecode);
-                    let abi = artifact.abi?;
-
-                    Some((id, ContractData { name, abi, bytecode, deployed_bytecode }))
+
+                    let CompactContractBytecode { abi, bytecode, deployed_bytecode } = artifact;
+
+                    Some((id, ContractData { name, abi: abi?, bytecode, deployed_bytecode }))
                 })
                 .collect(),
         )
@@ -62,7 +77,7 @@ impl ContractsByArtifact {
     /// Finds a contract which has a similar bytecode as `code`.
     pub fn find_by_creation_code(&self, code: &[u8]) -> Option<ArtifactWithContractRef> {
         self.iter().find(|(_, contract)| {
-            if let Some(bytecode) = &contract.bytecode {
+            if let Some(bytecode) = contract.bytecode() {
                 bytecode_diff_score(bytecode.as_ref(), code) <= 0.1
             } else {
                 false
@@ -73,14 +88,107 @@ impl ContractsByArtifact {
     /// Finds a contract which has a similar deployed bytecode as `code`.
     pub fn find_by_deployed_code(&self, code: &[u8]) -> Option<ArtifactWithContractRef> {
         self.iter().find(|(_, contract)| {
-            if let Some(deployed_bytecode) = &contract.deployed_bytecode {
+            if let Some(deployed_bytecode) = contract.deployed_bytecode() {
                 bytecode_diff_score(deployed_bytecode.as_ref(), code) <= 0.1
             } else {
                 false
             }
         })
     }
 
+    /// Finds a contract which deployed bytecode exactly matches the given code. Accounts for link
+    /// references and immutables.
+    pub fn find_by_deployed_code_exact(&self, code: &[u8]) -> Option<ArtifactWithContractRef> {
+        self.iter().find(|(_, contract)| {
+            let Some(ref deployed_bytecode) = contract.deployed_bytecode else {
+                return false;
+            };
+            let Some(ref deployed_code) = deployed_bytecode.bytecode else {
+                return false;
+            };
+
+            let len = match deployed_code.object {
+                BytecodeObject::Bytecode(ref bytes) => bytes.len(),
+                BytecodeObject::Unlinked(ref bytes) => bytes.len() / 2,
+            };
+
+            if len != code.len() {
+                return false;
+            }
+
+            // Collect ignored offsets by chaining link and immutable references.
+            let mut ignored = deployed_bytecode
+                .immutable_references
+                .values()
+                .chain(deployed_code.link_references.values().flat_map(|v| v.values()))
+                .flatten()
+                .cloned()
+                .collect::<Vec<_>>();
+
+            // For libraries solidity adds a call protection prefix to the bytecode. We need to
+            // ignore it as it includes library address determined at runtime.
+            // See https://docs.soliditylang.org/en/latest/contracts.html#call-protection-for-libraries and
+            // https://github.com/NomicFoundation/hardhat/blob/af7807cf38842a4f56e7f4b966b806e39631568a/packages/hardhat-verify/src/internal/solc/bytecode.ts#L172
+            let has_call_protection = match deployed_code.object {
+                BytecodeObject::Bytecode(ref bytes) => {
+                    bytes.starts_with(&CALL_PROTECTION_BYTECODE_PREFIX)
+                }
+                BytecodeObject::Unlinked(ref bytes) => {
+                    if let Ok(bytes) =
+                        Bytes::from_str(&bytes[..CALL_PROTECTION_BYTECODE_PREFIX.len() * 2])
+                    {
+                        bytes.starts_with(&CALL_PROTECTION_BYTECODE_PREFIX)
+                    } else {
+                        false
+                    }
+                }
+            };
+
+            if has_call_protection {
+                ignored.push(Offsets { start: 1, length: 20 });
+            }
+
+            ignored.sort_by_key(|o| o.start);
+
+            let mut left = 0;
+            for offset in ignored {
+                let right = offset.start as usize;
+
+                let matched = match deployed_code.object {
+                    BytecodeObject::Bytecode(ref bytes) => bytes[left..right] == code[left..right],
+                    BytecodeObject::Unlinked(ref bytes) => {
+                        if let Ok(bytes) = Bytes::from_str(&bytes[left * 2..right * 2]) {
+                            bytes == code[left..right]
+                        } else {
+                            false
+                        }
+                    }
+                };
+
+                if !matched {
+                    return false;
+                }
+
+                left = right + offset.length as usize;
+            }
+
+            if left < code.len() {
+                match deployed_code.object {
+                    BytecodeObject::Bytecode(ref bytes) => bytes[left..] == code[left..],
+                    BytecodeObject::Unlinked(ref bytes) => {
+                        if let Ok(bytes) = Bytes::from_str(&bytes[left * 2..]) {
+                            bytes == code[left..]
+                        } else {
+                            false
+                        }
+                    }
+                }
+            } else {
+                true
+            }
+        })
+    }
+
     /// Finds a contract which has the same contract name or identifier as `id`. If more than one is
     /// found, return error.
     pub fn find_by_name_or_identifier(&self, id: &str) -> Result<Option<ArtifactWithContractRef>> {

crates/evm/traces/src/identifier/local.rs

@@ -19,7 +19,7 @@ impl<'a> LocalTraceIdentifier<'a> {
     pub fn new(known_contracts: &'a ContractsByArtifact) -> Self {
         let mut ordered_ids = known_contracts
             .iter()
-            .filter_map(|(id, contract)| Some((id, contract.deployed_bytecode.as_ref()?)))
+            .filter_map(|(id, contract)| Some((id, contract.deployed_bytecode()?)))
             .map(|(id, bytecode)| (id, bytecode.len()))
             .collect::<Vec<_>>();
         ordered_ids.sort_by_key(|(_, len)| *len);
@@ -41,7 +41,7 @@ impl<'a> LocalTraceIdentifier<'a> {
 
         let mut check = |id| {
             let contract = self.known_contracts.get(id)?;
-            if let Some(deployed_bytecode) = &contract.deployed_bytecode {
+            if let Some(deployed_bytecode) = contract.deployed_bytecode() {
                 let score = bytecode_diff_score(deployed_bytecode, code);
                 if score == 0.0 {
                     trace!(target: "evm::traces", "found exact match");

crates/forge/bin/cmd/create.rs

@@ -169,7 +169,7 @@ impl CreateArgs {
         // since we don't know the address yet.
         let mut verify = forge_verify::VerifyArgs {
             address: Default::default(),
-            contract: self.contract.clone(),
+            contract: Some(self.contract.clone()),
             compiler_version: None,
             constructor_args,
             constructor_args_path: None,
@@ -199,7 +199,9 @@ impl CreateArgs {
         verify.etherscan.key =
             config.get_etherscan_config_with_chain(Some(chain.into()))?.map(|c| c.key);
 
-        verify.verification_provider()?.preflight_check(verify).await?;
+        let context = verify.resolve_context().await?;
+
+        verify.verification_provider()?.preflight_check(verify, context).await?;
         Ok(())
     }
 
@@ -318,7 +320,7 @@ impl CreateArgs {
             if self.opts.compiler.optimize { self.opts.compiler.optimizer_runs } else { None };
         let verify = forge_verify::VerifyArgs {
             address,
-            contract: self.contract,
+            contract: Some(self.contract),
             compiler_version: None,
             constructor_args,
             constructor_args_path: None,

crates/script/src/execute.rs

@@ -16,7 +16,7 @@ use foundry_cli::utils::{ensure_clean_constructor, needs_setup};
 use foundry_common::{
     fmt::{format_token, format_token_raw},
     provider::get_http_provider,
-    shell, ContractData, ContractsByArtifact,
+    shell, ContractsByArtifact,
 };
 use foundry_config::{Config, NamedChain};
 use foundry_debugger::Debugger;
@@ -61,20 +61,25 @@ impl LinkedState {
     pub async fn prepare_execution(self) -> Result<PreExecutionState> {
         let Self { args, script_config, script_wallets, build_data } = self;
 
-        let ContractData { abi, bytecode, .. } = build_data.get_target_contract()?;
+        let target_contract = build_data.get_target_contract()?;
 
-        let bytecode = bytecode.ok_or_eyre("target contract has no bytecode")?;
+        let bytecode = target_contract.bytecode().ok_or_eyre("target contract has no bytecode")?;
 
-        let (func, calldata) = args.get_method_and_calldata(&abi)?;
+        let (func, calldata) = args.get_method_and_calldata(&target_contract.abi)?;
 
-        ensure_clean_constructor(&abi)?;
+        ensure_clean_constructor(&target_contract.abi)?;
 
         Ok(PreExecutionState {
             args,
             script_config,
             script_wallets,
             build_data,
-            execution_data: ExecutionData { func, calldata, bytecode, abi },
+            execution_data: ExecutionData {
+                func,
+                calldata,
+                bytecode: bytecode.clone(),
+                abi: target_contract.abi,
+            },
         })
     }
 }

crates/script/src/lib.rs

@@ -361,8 +361,8 @@ impl ScriptArgs {
 
         // From artifacts
         for (artifact, contract) in known_contracts.iter() {
-            let Some(bytecode) = &contract.bytecode else { continue };
-            let Some(deployed_bytecode) = &contract.deployed_bytecode else { continue };
+            let Some(bytecode) = contract.bytecode() else { continue };
+            let Some(deployed_bytecode) = contract.deployed_bytecode() else { continue };
             bytecodes.push((artifact.name.clone(), bytecode, deployed_bytecode));
         }
 

crates/script/src/transaction.rs

@@ -131,7 +131,7 @@ impl TransactionWithMetadata {
 
         let Some(data) = self.transaction.input.input() else { return Ok(()) };
         let Some(info) = info else { return Ok(()) };
-        let Some(bytecode) = info.bytecode.as_ref() else { return Ok(()) };
+        let Some(bytecode) = info.bytecode() else { return Ok(()) };
 
         // `create2` transactions are prefixed by a 32 byte salt.
         let creation_code = if is_create2 {

crates/script/src/verify.rs

@@ -106,7 +106,7 @@ impl VerifyBundle {
         libraries: &[String],
     ) -> Option<VerifyArgs> {
         for (artifact, contract) in self.known_contracts.iter() {
-            let Some(bytecode) = contract.bytecode.as_ref() else { continue };
+            let Some(bytecode) = contract.bytecode() else { continue };
             // If it's a CREATE2, the tx.data comes with a 32-byte salt in the beginning
             // of the transaction
             if data.split_at(create2_offset).1.starts_with(bytecode) {
@@ -130,7 +130,7 @@ impl VerifyBundle {
 
                 let verify = VerifyArgs {
                     address: contract_address,
-                    contract,
+                    contract: Some(contract),
                     compiler_version: Some(version.to_string()),
                     constructor_args: Some(hex::encode(constructor_args)),
                     constructor_args_path: None,

crates/verify/Cargo.toml

@@ -36,6 +36,7 @@ semver = "1"
 regex = { version = "1", default-features = false }
 once_cell = "1"
 yansi.workspace = true
+itertools.workspace = true
 
 [dev-dependencies]
 tokio = { workspace = true, features = ["macros"] }