feat(cheatcodes): forge execution context check (#7377)

* feat(cheatcodes): forge execution context check

* Add tests for test and snapshot contexts

* Add isTestContext, isScriptContext cheatcodes

* Add script dry run and broadcast tests

* Proper enum in cheatcodes schema, alphabetical order

* Single isContext cheatcode in env group, taking enum as param. remove context group

* Changes after review: remove discriminant calls, use OnceLock

* Review changes: tests should not be async

* Review changes: implement PartialEq for ForgeContext, remove is_forge_context fn

* Properly add new ForgeContext enum

Author: grandizzy

crates/cheatcodes/assets/cheatcodes.json

@@ -88,6 +88,7 @@ impl Cheatcodes<'static> {
             enums: Cow::Owned(vec![
                 Vm::CallerMode::ENUM.clone(),
                 Vm::AccountAccessKind::ENUM.clone(),
+                Vm::ForgeContext::ENUM.clone(),
             ]),
             errors: Vm::VM_ERRORS.iter().map(|&x| x.clone()).collect(),
             events: Cow::Borrowed(&[]),

crates/cheatcodes/spec/src/lib.rs

@@ -3,6 +3,7 @@
 #![allow(missing_docs)]
 
 use super::*;
+use crate::Vm::ForgeContext;
 use alloy_sol_types::sol;
 use foundry_macros::Cheatcode;
 
@@ -63,6 +64,28 @@ interface Vm {
         Extcodecopy,
     }
 
+    /// Forge execution contexts.
+    enum ForgeContext {
+        /// Test group execution context (test, coverage or snapshot).
+        TestGroup,
+        /// `forge test` execution context.
+        Test,
+        /// `forge coverage` execution context.
+        Coverage,
+        /// `forge snapshot` execution context.
+        Snapshot,
+        /// Script group execution context (dry run, broadcast or resume).
+        ScriptGroup,
+        /// `forge script` execution context.
+        ScriptDryRun,
+        /// `forge script --broadcast` execution context.
+        ScriptBroadcast,
+        /// `forge script --resume` execution context.
+        ScriptResume,
+        /// Unknown `forge` execution context.
+        Unknown,
+    }
+
     /// An Ethereum log. Returned by `getRecordedLogs`.
     struct Log {
         /// The topics of the log, including the signature, if any.
@@ -1598,6 +1621,10 @@ interface Vm {
         external view
         returns (bytes[] memory value);
 
+    /// Returns true if `forge` command was executed in given context.
+    #[cheatcode(group = Environment)]
+    function isContext(ForgeContext context) external view returns (bool isContext);
+
     // ======== Scripts ========
 
     // -------- Broadcasting Transactions --------
@@ -2065,3 +2092,30 @@ interface Vm {
     function toBase64URL(string calldata data) external pure returns (string memory);
 }
 }
+
+impl PartialEq for ForgeContext {
+    // Handles test group case (any of test, coverage or snapshot)
+    // and script group case (any of dry run, broadcast or resume).
+    fn eq(&self, other: &Self) -> bool {
+        match (self, other) {
+            (_, &ForgeContext::TestGroup) => {
+                self == &ForgeContext::Test ||
+                    self == &ForgeContext::Snapshot ||
+                    self == &ForgeContext::Coverage
+            }
+            (_, &ForgeContext::ScriptGroup) => {
+                self == &ForgeContext::ScriptDryRun ||
+                    self == &ForgeContext::ScriptBroadcast ||
+                    self == &ForgeContext::ScriptResume
+            }
+            (&ForgeContext::Test, &ForgeContext::Test) |
+            (&ForgeContext::Snapshot, &ForgeContext::Snapshot) |
+            (&ForgeContext::Coverage, &ForgeContext::Coverage) |
+            (&ForgeContext::ScriptDryRun, &ForgeContext::ScriptDryRun) |
+            (&ForgeContext::ScriptBroadcast, &ForgeContext::ScriptBroadcast) |
+            (&ForgeContext::ScriptResume, &ForgeContext::ScriptResume) |
+            (&ForgeContext::Unknown, &ForgeContext::Unknown) => true,
+            _ => false,
+        }
+    }
+}

crates/cheatcodes/spec/src/vm.rs

@@ -4,7 +4,10 @@ use crate::{string, Cheatcode, Cheatcodes, Error, Result, Vm::*};
 use alloy_dyn_abi::DynSolType;
 use alloy_primitives::Bytes;
 use alloy_sol_types::SolValue;
-use std::env;
+use std::{env, sync::OnceLock};
+
+/// Stores the forge execution context for the duration of the program.
+static FORGE_CONTEXT: OnceLock<ForgeContext> = OnceLock::new();
 
 impl Cheatcode for setEnvCall {
     fn apply(&self, _state: &mut Cheatcodes) -> Result {
@@ -235,6 +238,19 @@ impl Cheatcode for envOr_13Call {
     }
 }
 
+impl Cheatcode for isContextCall {
+    fn apply(&self, _state: &mut Cheatcodes) -> Result {
+        let Self { context } = self;
+        Ok((FORGE_CONTEXT.get() == Some(context)).abi_encode())
+    }
+}
+
+/// Set `forge` command current execution context for the duration of the program.
+/// Execution context is immutable, subsequent calls of this function won't change the context.
+pub fn set_execution_context(context: ForgeContext) {
+    let _ = FORGE_CONTEXT.set(context);
+}
+
 fn env(key: &str, ty: &DynSolType) -> Result {
     get_env(key).and_then(|val| string::parse(&val, ty).map_err(map_env_err(key, &val)))
 }

crates/cheatcodes/src/env.rs

@@ -34,8 +34,10 @@ mod test;
 mod toml;
 mod utils;
 
+pub use env::set_execution_context;
 pub use script::ScriptWallets;
 pub use test::expect::ExpectedCallTracker;
+pub use Vm::ForgeContext;
 
 /// Cheatcode implementation.
 pub(crate) trait Cheatcode: CheatcodeDef + DynCheatcode {

crates/cheatcodes/src/lib.rs

@@ -5,6 +5,7 @@ use clap::{CommandFactory, Parser};
 use clap_complete::generate;
 use eyre::Result;
 use foundry_cli::{handler, utils};
+use foundry_evm::inspectors::cheatcodes::{set_execution_context, ForgeContext};
 
 mod cmd;
 use cmd::{cache::CacheSubcommands, generate::GenerateSubcommands, watch};
@@ -23,6 +24,8 @@ fn main() -> Result<()> {
     utils::enable_paint();
 
     let opts = Forge::parse();
+    init_execution_context(&opts.cmd);
+
     match opts.cmd {
         ForgeSubcommand::Test(cmd) => {
             if cmd.is_watch() {
@@ -107,3 +110,25 @@ fn main() -> Result<()> {
         },
     }
 }
+
+/// Set the program execution context based on `forge` subcommand used.
+/// The execution context can be set only once per program, and it can be checked by using
+/// cheatcodes.
+fn init_execution_context(subcommand: &ForgeSubcommand) {
+    let context = match subcommand {
+        ForgeSubcommand::Test(_) => ForgeContext::Test,
+        ForgeSubcommand::Coverage(_) => ForgeContext::Coverage,
+        ForgeSubcommand::Snapshot(_) => ForgeContext::Snapshot,
+        ForgeSubcommand::Script(cmd) => {
+            if cmd.broadcast {
+                ForgeContext::ScriptBroadcast
+            } else if cmd.resume {
+                ForgeContext::ScriptResume
+            } else {
+                ForgeContext::ScriptDryRun
+            }
+        }
+        _ => ForgeContext::Unknown,
+    };
+    set_execution_context(context);
+}

crates/forge/bin/main.rs

@@ -0,0 +1,81 @@
+//! Contains tests for checking forge execution context cheatcodes
+const FORGE_TEST_CONTEXT_CONTRACT: &str = r#"
+import "./test.sol";
+interface Vm {
+    enum ForgeContext { TestGroup, Test, Coverage, Snapshot, ScriptGroup, ScriptDryRun, ScriptBroadcast, ScriptResume, Unknown }
+    function isContext(ForgeContext context) external view returns (bool isContext);
+}
+
+contract ForgeContextTest is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+
+    function testForgeTestContext() external view {
+        require(vm.isContext(Vm.ForgeContext.TestGroup) && !vm.isContext(Vm.ForgeContext.ScriptGroup), "wrong context");
+        require(vm.isContext(Vm.ForgeContext.Test), "wrong context");
+        require(!vm.isContext(Vm.ForgeContext.Coverage), "wrong context");
+        require(!vm.isContext(Vm.ForgeContext.Snapshot), "wrong context");
+    }
+    function testForgeSnapshotContext() external view {
+        require(vm.isContext(Vm.ForgeContext.TestGroup) && !vm.isContext(Vm.ForgeContext.ScriptGroup), "wrong context");
+        require(vm.isContext(Vm.ForgeContext.Snapshot), "wrong context");
+        require(!vm.isContext(Vm.ForgeContext.Test), "wrong context");
+        require(!vm.isContext(Vm.ForgeContext.Coverage), "wrong context");
+    }
+    function testForgeCoverageContext() external view {
+        require(vm.isContext(Vm.ForgeContext.TestGroup) && !vm.isContext(Vm.ForgeContext.ScriptGroup), "wrong context");
+        require(vm.isContext(Vm.ForgeContext.Coverage), "wrong context");
+        require(!vm.isContext(Vm.ForgeContext.Test), "wrong context");
+        require(!vm.isContext(Vm.ForgeContext.Snapshot), "wrong context");
+    }
+
+    function runDryRun() external view {
+        require(vm.isContext(Vm.ForgeContext.ScriptGroup) && !vm.isContext(Vm.ForgeContext.TestGroup), "wrong context");
+        require(vm.isContext(Vm.ForgeContext.ScriptDryRun), "wrong context");
+        require(!vm.isContext(Vm.ForgeContext.ScriptBroadcast), "wrong context");
+        require(!vm.isContext(Vm.ForgeContext.ScriptResume), "wrong context");
+    }
+    function runBroadcast() external view {
+        require(vm.isContext(Vm.ForgeContext.ScriptGroup) && !vm.isContext(Vm.ForgeContext.TestGroup), "wrong context");
+        require(vm.isContext(Vm.ForgeContext.ScriptBroadcast), "wrong context");
+        require(!vm.isContext(Vm.ForgeContext.ScriptDryRun), "wrong context");
+        require(!vm.isContext(Vm.ForgeContext.ScriptResume), "wrong context");
+    }
+}
+   "#;
+
+// tests that context properly set for `forge test` command
+forgetest!(can_set_forge_test_standard_context, |prj, cmd| {
+    prj.insert_ds_test();
+    prj.add_source("ForgeContextTest.t.sol", FORGE_TEST_CONTEXT_CONTRACT).unwrap();
+    cmd.args(["test", "--match-test", "testForgeTestContext"]).assert_success();
+});
+
+// tests that context properly set for `forge snapshot` command
+forgetest!(can_set_forge_test_snapshot_context, |prj, cmd| {
+    prj.insert_ds_test();
+    prj.add_source("ForgeContextTest.t.sol", FORGE_TEST_CONTEXT_CONTRACT).unwrap();
+    cmd.args(["snapshot", "--match-test", "testForgeSnapshotContext"]).assert_success();
+});
+
+// tests that context properly set for `forge coverage` command
+forgetest!(can_set_forge_test_coverage_context, |prj, cmd| {
+    prj.insert_ds_test();
+    prj.add_source("ForgeContextTest.t.sol", FORGE_TEST_CONTEXT_CONTRACT).unwrap();
+    cmd.args(["coverage", "--match-test", "testForgeCoverageContext"]).assert_success();
+});
+
+// tests that context properly set for `forge script` command
+forgetest!(can_set_forge_script_dry_run_context, |prj, cmd| {
+    prj.insert_ds_test();
+    let script =
+        prj.add_source("ForgeScriptContextTest.s.sol", FORGE_TEST_CONTEXT_CONTRACT).unwrap();
+    cmd.arg("script").arg(script).args(["--sig", "runDryRun()"]).assert_success();
+});
+
+// tests that context properly set for `forge script --broadcast` command
+forgetest!(can_set_forge_script_broadcast_context, |prj, cmd| {
+    prj.insert_ds_test();
+    let script =
+        prj.add_source("ForgeScriptContextTest.s.sol", FORGE_TEST_CONTEXT_CONTRACT).unwrap();
+    cmd.arg("script").arg(script).args(["--broadcast", "--sig", "runBroadcast()"]).assert_success();
+});

crates/forge/tests/cli/context.rs

@@ -8,6 +8,7 @@ mod build;
 mod cache;
 mod cmd;
 mod config;
+mod context;
 mod coverage;
 mod create;
 mod debug;