fix: smarter getCode validation

crates/cheatcodes/src/config.rs

@@ -8,6 +8,7 @@ use foundry_config::{
     ResolvedRpcEndpoints,
 };
 use foundry_evm_core::opts::EvmOpts;
+use semver::Version;
 use std::{
     collections::HashMap,
     path::{Path, PathBuf},
@@ -47,6 +48,8 @@ pub struct CheatsConfig {
     /// If Some, `vm.getDeployedCode` invocations are validated to be in scope of this list.
     /// If None, no validation is performed.
     pub available_artifacts: Option<Vec<ArtifactId>>,
+    /// Version of the script/test contract which is currently running.
+    pub running_version: Option<Version>,
 }
 
 impl CheatsConfig {
@@ -56,6 +59,7 @@ impl CheatsConfig {
         evm_opts: EvmOpts,
         available_artifacts: Option<Vec<ArtifactId>>,
         script_wallets: Option<ScriptWallets>,
+        running_version: Option<Version>,
     ) -> Self {
         let mut allowed_paths = vec![config.__root.0.clone()];
         allowed_paths.extend(config.libs.clone());
@@ -82,6 +86,7 @@ impl CheatsConfig {
             labels: config.labels.clone(),
             script_wallets,
             available_artifacts,
+            running_version,
         }
     }
 
@@ -200,6 +205,7 @@ impl Default for CheatsConfig {
             labels: Default::default(),
             script_wallets: None,
             available_artifacts: Default::default(),
+            running_version: Default::default(),
         }
     }
 }
@@ -215,6 +221,7 @@ mod tests {
             Default::default(),
             None,
             None,
+            None,
         )
     }
 

crates/cheatcodes/src/fs.rs

@@ -271,14 +271,40 @@ impl Cheatcode for getDeployedCodeCall {
 }
 
 /// Returns the path to the json artifact depending on the input
+/// 
+/// Can parse following input formats:
+/// - `path/to/artifact.json`
+/// - `path/to/contract.sol`
+/// - `path/to/contract.sol:ContractName`
+/// - `path/to/contract.sol:ContractName:0.8.23`
+/// - `path/to/contract.sol:0.8.23`
+/// - `ContractName`
+/// - `ContractName:0.8.23`
 fn get_artifact_path(state: &Cheatcodes, path: &str) -> Result<PathBuf> {
     if path.ends_with(".json") {
         Ok(PathBuf::from(path))
     } else {
         let mut parts = path.split(':');
-        let file = PathBuf::from(parts.next().unwrap());
-        let contract_name = parts.next();
-        let version = parts.next();
+
+        let mut file = None;
+        let mut contract_name = None;
+        let mut version = None;
+
+        let path_or_name = parts.next().unwrap();
+        if path_or_name.ends_with(".sol") {
+            file = Some(PathBuf::from(path_or_name));
+            if let Some(name_or_version) = parts.next() {
+                if name_or_version.contains(".") {
+                    version = Some(name_or_version);
+                } else {
+                    contract_name = Some(name_or_version);
+                    version = parts.next();
+                }
+            }
+        } else {
+            contract_name = Some(path_or_name);
+            version = parts.next();
+        }
 
         let version = if let Some(version) = version {
             Some(Version::parse(version).map_err(|_| fmt_err!("Error parsing version"))?)
@@ -288,44 +314,71 @@ fn get_artifact_path(state: &Cheatcodes, path: &str) -> Result<PathBuf> {
 
         // Use available artifacts list if available
         if let Some(available_ids) = &state.config.available_artifacts {
-            let mut artifact = None;
-
-            for id in available_ids.iter() {
-                // name might be in the form of "Counter.0.8.23"
-                let id_name = id.name.split('.').next().unwrap();
-
-                if !id.source.ends_with(&file) {
-                    continue;
-                }
-                if let Some(name) = contract_name {
-                    if id_name != name {
-                        continue;
+            let filtered = available_ids
+                .iter()
+                .filter(|id| {
+                    // name might be in the form of "Counter.0.8.23"
+                    let id_name = id.name.split('.').next().unwrap();
+
+                    if let Some(path) = &file {
+                        if !id.source.ends_with(path) {
+                            return false;
+                        }
                     }
-                }
-                if let Some(ref version) = version {
-                    if id.version.minor != version.minor ||
-                        id.version.major != version.major ||
-                        id.version.patch != version.patch
-                    {
-                        continue;
+                    if let Some(name) = contract_name {
+                        if id_name != name {
+                            return false;
+                        }
                     }
+                    if let Some(ref version) = version {
+                        if id.version.minor != version.minor ||
+                            id.version.major != version.major ||
+                            id.version.patch != version.patch
+                        {
+                            return false;
+                        }
+                    }
+                    true
+                })
+                .collect::<Vec<_>>();
+
+            let artifact = match filtered.len() {
+                0 => Err(fmt_err!("No matching artifact found")),
+                1 => Ok(filtered[0]),
+                _ => {
+                    // If we know the current script/test contract solc version, try to filter by it
+                    state
+                        .config
+                        .running_version
+                        .as_ref()
+                        .and_then(|version| {
+                            let filtered = filtered
+                                .into_iter()
+                                .filter(|id| id.version == *version)
+                                .collect::<Vec<_>>();
+
+                            (filtered.len() == 1).then_some(filtered[0])
+                        })
+                        .ok_or_else(|| fmt_err!("Multiple matching artifacts found"))
                 }
-                if artifact.is_some() {
-                    return Err(fmt_err!("Multiple matching artifacts found"));
-                }
-                artifact = Some(id);
-            }
+            }?;
 
-            let artifact = artifact.ok_or_else(|| fmt_err!("No matching artifact found"))?;
             Ok(artifact.path.clone())
         } else {
-            let file = file.to_string_lossy();
-            let contract_name = if let Some(contract_name) = contract_name {
-                contract_name.to_owned()
-            } else {
-                file.replace(".sol", "")
-            };
-            Ok(state.config.paths.artifacts.join(format!("{file}/{contract_name}.json")))
+            let path_in_artifacts = match (file.map(|f| f.to_string_lossy().to_string()), contract_name) {
+                (Some(file), Some(contract_name)) => {
+                    Ok(format!("{file}/{contract_name}.json"))
+                }
+                (None, Some(contract_name)) => {
+                    Ok(format!("{contract_name}.sol/{contract_name}.json"))
+                }
+                (Some(file), None) => {
+                    let name = file.replace(".sol", "");
+                    Ok(format!("{file}/{name}.json"))
+                }
+                _ => Err(fmt_err!("Invalid artifact path")),
+            }?;
+            Ok(state.config.paths.artifacts.join(path_in_artifacts))
         }
     }
 }

crates/chisel/src/executor.rs

@@ -308,6 +308,7 @@ impl SessionSource {
                         self.config.evm_opts.clone(),
                         None,
                         None,
+                        self.solc.version().ok(),
                     )
                     .into(),
                 )

crates/forge/bin/cmd/coverage.rs

@@ -317,6 +317,7 @@ impl CoverageArgs {
                 evm_opts.clone(),
                 Some(artifact_ids),
                 None,
+                None,
             ))
             .with_test_options(TestOptions {
                 fuzz: config.fuzz,

crates/forge/bin/cmd/test/mod.rs

@@ -287,6 +287,7 @@ impl TestArgs {
                 evm_opts.clone(),
                 Some(artifact_ids),
                 None,
+                None, // populated separately for each test contract
             ))
             .with_test_options(test_options)
             .enable_isolation(evm_opts.isolate)

crates/forge/src/multi_runner.rs

@@ -9,13 +9,8 @@ use foundry_compilers::{
     artifacts::Libraries, contracts::ArtifactContracts, Artifact, ArtifactId, ProjectCompileOutput,
 };
 use foundry_evm::{
-    backend::Backend,
-    decode::RevertDecoder,
-    executors::{Executor, ExecutorBuilder},
-    fork::CreateFork,
-    inspectors::CheatsConfig,
-    opts::EvmOpts,
-    revm,
+    backend::Backend, decode::RevertDecoder, executors::ExecutorBuilder, fork::CreateFork,
+    inspectors::CheatsConfig, opts::EvmOpts, revm,
 };
 use foundry_linking::{LinkOutput, Linker};
 use rayon::prelude::*;
@@ -158,18 +153,6 @@ impl MultiContractRunner {
 
         // The DB backend that serves all the data.
         let db = Backend::spawn(self.fork.take());
-        let executor = ExecutorBuilder::new()
-            .inspectors(|stack| {
-                stack
-                    .cheatcodes(self.cheats_config.clone())
-                    .trace(self.evm_opts.verbosity >= 3 || self.debug)
-                    .debug(self.debug)
-                    .coverage(self.coverage)
-                    .enable_isolation(self.isolation)
-            })
-            .spec(self.evm_spec)
-            .gas_limit(self.evm_opts.gas_limit())
-            .build(self.env.clone(), db);
 
         let find_timer = Instant::now();
         let contracts = self.matching_contracts(filter).collect::<Vec<_>>();
@@ -182,30 +165,46 @@ impl MultiContractRunner {
         );
 
         contracts.par_iter().for_each_with(tx, |tx, &(id, contract)| {
-            let identifier = id.identifier();
-            let executor = executor.clone();
-            let result = self.run_tests(&identifier, contract, executor, filter);
-            let _ = tx.send((identifier, result));
+            let result = self.run_tests(id, contract, db.clone(), filter);
+            let _ = tx.send((id.identifier(), result));
         })
     }
 
     fn run_tests(
         &self,
-        name: &str,
+        artifact_id: &ArtifactId,
         contract: &TestContract,
-        executor: Executor,
+        db: Backend,
         filter: &dyn TestFilter,
     ) -> SuiteResult {
-        let mut span_name = name;
+        let identifier = artifact_id.identifier();
+        let mut span_name = identifier.as_str();
+
+        let mut cheats_config = self.cheats_config.as_ref().clone();
+        cheats_config.running_version = Some(artifact_id.version.clone());
+
+        let executor = ExecutorBuilder::new()
+            .inspectors(|stack| {
+                stack
+                    .cheatcodes(Arc::new(cheats_config))
+                    .trace(self.evm_opts.verbosity >= 3 || self.debug)
+                    .debug(self.debug)
+                    .coverage(self.coverage)
+                    .enable_isolation(self.isolation)
+            })
+            .spec(self.evm_spec)
+            .gas_limit(self.evm_opts.gas_limit())
+            .build(self.env.clone(), db);
+
         if !enabled!(tracing::Level::TRACE) {
-            span_name = get_contract_name(span_name);
+            span_name = get_contract_name(&identifier);
         }
         let _guard = info_span!("run_tests", name = span_name).entered();
 
         debug!("start executing all tests in contract");
 
         let runner = ContractRunner::new(
-            name,
+            &identifier,
             executor,
             contract,
             self.evm_opts.initial_balance,

crates/forge/tests/it/cheats.rs

@@ -2,14 +2,13 @@
 
 use crate::{
     config::*,
-    test_helpers::{RE_PATH_SEPARATOR, TEST_DATA_DEFAULT},
+    test_helpers::{ForgeTestData, RE_PATH_SEPARATOR, TEST_DATA_DEFAULT, TEST_DATA_MULTI_VERSION},
 };
 use foundry_config::{fs_permissions::PathPermission, FsPermissions};
 use foundry_test_utils::Filter;
 
 /// Executes all cheat code tests but not fork cheat codes
-#[tokio::test(flavor = "multi_thread")]
-async fn test_cheats_local() {
+async fn test_cheats_local(test_data: &ForgeTestData) {
     let mut filter =
         Filter::new(".*", ".*", &format!(".*cheats{RE_PATH_SEPARATOR}*")).exclude_paths("Fork");
 
@@ -18,9 +17,19 @@ async fn test_cheats_local() {
         filter = filter.exclude_tests("(Ffi|File|Line|Root)");
     }
 
-    let mut config = TEST_DATA_DEFAULT.config.clone();
+    let mut config = test_data.config.clone();
     config.fs_permissions = FsPermissions::new(vec![PathPermission::read_write("./")]);
-    let runner = TEST_DATA_DEFAULT.runner_with_config(config);
+    let runner = test_data.runner_with_config(config);
 
     TestConfig::with_filter(runner, filter).run().await;
 }
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_cheats_local_default() {
+    test_cheats_local(&TEST_DATA_DEFAULT).await
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_cheats_local_multi_version() {
+    test_cheats_local(&TEST_DATA_MULTI_VERSION).await
+}

crates/forge/tests/it/test_helpers.rs

@@ -32,13 +32,15 @@ const TESTDATA: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../../testdata");
 pub enum ForgeTestProfile {
     Default,
     Cancun,
+    MultiVersion,
 }
 
 impl fmt::Display for ForgeTestProfile {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
             ForgeTestProfile::Default => write!(f, "default"),
             ForgeTestProfile::Cancun => write!(f, "cancun"),
+            ForgeTestProfile::MultiVersion => write!(f, "multi-version"),
         }
     }
 }
@@ -206,7 +208,13 @@ impl ForgeTestData {
         let output = self.output.clone();
         let artifact_ids = output.artifact_ids().map(|(id, _)| id).collect();
         self.base_runner()
-            .with_cheats_config(CheatsConfig::new(&config, opts.clone(), Some(artifact_ids), None))
+            .with_cheats_config(CheatsConfig::new(
+                &config,
+                opts.clone(),
+                Some(artifact_ids),
+                None,
+                None,
+            ))
             .sender(config.sender)
             .with_test_options(self.test_opts.clone())
             .build(root, output, env, opts.clone())
@@ -274,6 +282,10 @@ pub static TEST_DATA_DEFAULT: Lazy<ForgeTestData> =
 pub static TEST_DATA_CANCUN: Lazy<ForgeTestData> =
     Lazy::new(|| ForgeTestData::new(ForgeTestProfile::Cancun));
 
+/// Data for tests requiring Cancun support on Solc and EVM level.
+pub static TEST_DATA_MULTI_VERSION: Lazy<ForgeTestData> =
+    Lazy::new(|| ForgeTestData::new(ForgeTestProfile::MultiVersion));
+
 pub fn manifest_root() -> &'static Path {
     let mut root = Path::new(env!("CARGO_MANIFEST_DIR"));
     // need to check here where we're executing the test from, if in `forge` we need to also allow