feat(script): revert if address(this) used (#10295)

* Update stack.rs

* Update lib.rs

* Update lib.rs

* Update stack.rs

* Create ScriptAddressWarn.t.sol

* Update mod.rs

* Update stack.rs

* Update lib.rs

* Create script.rs

* Fix compilation, cleanup, add new test in script tests

* Set and check current address is script address

* Update stack.rs

* Allow calls to external libraries

* changes after review: use sh_err

---------

Co-authored-by: grandizzy <[email protected]>
Co-authored-by: grandizzy <[email protected]>

Author: gap-editor

crates/evm/evm/src/executors/mod.rs

@@ -257,6 +257,11 @@ impl Executor {
         self
     }
 
+    #[inline]
+    pub fn set_script(&mut self, script_address: Address) {
+        self.inspector_mut().script(script_address);
+    }
+
     #[inline]
     pub fn set_trace_printer(&mut self, trace_printer: bool) -> &mut Self {
         self.inspector_mut().print(trace_printer);

crates/evm/evm/src/inspectors/mod.rs

@@ -13,5 +13,8 @@ pub use chisel_state::ChiselState;
 mod logs;
 pub use logs::LogCollector;
 
+mod script;
+pub use script::ScriptExecutionInspector;
+
 mod stack;
 pub use stack::{InspectorData, InspectorStack, InspectorStackBuilder};

crates/evm/evm/src/inspectors/script.rs

@@ -0,0 +1,37 @@
+use alloy_primitives::Address;
+use foundry_common::sh_err;
+use revm::{
+    interpreter::{opcode::ADDRESS, InstructionResult, Interpreter},
+    Database, EvmContext, Inspector,
+};
+
+/// An inspector that enforces certain rules during script execution.
+///
+/// Currently, it only warns if the `ADDRESS` opcode is used within the script's main contract.
+#[derive(Clone, Debug, Default)]
+pub struct ScriptExecutionInspector {
+    /// The address of the script contract being executed.
+    pub script_address: Address,
+}
+
+impl<DB: Database> Inspector<DB> for ScriptExecutionInspector {
+    #[inline]
+    fn step(&mut self, interpreter: &mut Interpreter, _ecx: &mut EvmContext<DB>) {
+        // Check if both target and bytecode address are the same as script contract address
+        // (allow calling external libraries when bytecode address is different).
+        if interpreter.current_opcode() == ADDRESS &&
+            interpreter.contract.target_address == self.script_address &&
+            interpreter.contract.bytecode_address.unwrap_or_default() == self.script_address
+        {
+            // Log the reason for revert
+            let _ = sh_err!(
+                "Usage of `address(this)` detected in script contract. Script contracts are ephemeral and their addresses should not be relied upon."
+            );
+            // Set the instruction result to Revert to stop execution
+            interpreter.instruction_result = InstructionResult::Revert;
+        }
+        // Note: We don't return anything here as step returns void.
+        // The original check returned InstructionResult::Continue, but that's the default
+        // behavior.
+    }
+}

crates/evm/evm/src/inspectors/stack.rs

@@ -1,6 +1,6 @@
 use super::{
     Cheatcodes, CheatsConfig, ChiselState, CoverageCollector, Fuzzer, LogCollector,
-    TracingInspector,
+    ScriptExecutionInspector, TracingInspector,
 };
 use alloy_primitives::{map::AddressHashMap, Address, Bytes, Log, TxKind, U256};
 use foundry_cheatcodes::{CheatcodesExecutor, Wallets};
@@ -199,6 +199,7 @@ impl InspectorStackBuilder {
         if let Some(chisel_state) = chisel_state {
             stack.set_chisel(chisel_state);
         }
+
         stack.collect_coverage(coverage.unwrap_or(false));
         stack.collect_logs(logs.unwrap_or(true));
         stack.print(print.unwrap_or(false));
@@ -290,6 +291,7 @@ pub struct InspectorStackInner {
     pub log_collector: Option<LogCollector>,
     pub printer: Option<CustomPrintTracer>,
     pub tracer: Option<TracingInspector>,
+    pub script_execution_inspector: Option<ScriptExecutionInspector>,
     pub enable_isolation: bool,
     pub odyssey: bool,
     pub create2_deployer: Address,
@@ -437,6 +439,13 @@ impl InspectorStack {
         }
     }
 
+    /// Set whether to enable script execution inspector.
+    #[inline]
+    pub fn script(&mut self, script_address: Address) {
+        self.script_execution_inspector.get_or_insert_with(Default::default).script_address =
+            script_address;
+    }
+
     /// Collects all the data gathered during inspection into a single struct.
     #[inline]
     pub fn collect(self) -> InspectorData {
@@ -757,7 +766,13 @@ impl Inspector<&mut dyn DatabaseExt> for InspectorStackRefMut<'_> {
         ecx: &mut EvmContext<&mut dyn DatabaseExt>,
     ) {
         call_inspectors!(
-            [&mut self.coverage, &mut self.tracer, &mut self.cheatcodes, &mut self.printer],
+            [
+                &mut self.coverage,
+                &mut self.tracer,
+                &mut self.cheatcodes,
+                &mut self.script_execution_inspector,
+                &mut self.printer
+            ],
             |inspector| inspector.initialize_interp(interpreter, ecx),
         );
     }
@@ -769,7 +784,8 @@ impl Inspector<&mut dyn DatabaseExt> for InspectorStackRefMut<'_> {
                 &mut self.tracer,
                 &mut self.coverage,
                 &mut self.cheatcodes,
-                &mut self.printer,
+                &mut self.script_execution_inspector,
+                &mut self.printer
             ],
             |inspector| inspector.step(interpreter, ecx),
         );

crates/forge/tests/cli/script.rs

@@ -2645,3 +2645,28 @@ Script ran successfully.
 
 "#]]);
 });
+
+// Tests that script reverts if it uses `address(this)`.
+forgetest_init!(should_revert_on_address_opcode, |prj, cmd| {
+    prj.add_script(
+        "ScriptWithAddress.s.sol",
+        r#"
+        import {Script, console} from "forge-std/Script.sol";
+
+    contract ScriptWithAddress is Script {
+        function run() public view {
+            console.log("script address", address(this));
+        }
+    }
+    "#,
+    )
+    .unwrap();
+
+    cmd.arg("script").arg("ScriptWithAddress").assert_failure().stderr_eq(str![[r#"
+...
+Error: Usage of `address(this)` detected in script contract. Script contracts are ephemeral and their addresses should not be relied upon.
+Error: script failed: <empty revert data>
+...
+
+"#]]);
+});

crates/script/src/runner.rs

@@ -157,6 +157,9 @@ impl ScriptRunner {
             self.executor.set_nonce(self.evm_opts.sender, prev_sender_nonce)?;
         }
 
+        // set script address to be used by execution inspector
+        self.executor.set_script(address);
+
         traces.extend(constructor_traces.map(|traces| (TraceKind::Deployment, traces)));
 
         // Optionally call the `setUp` function

testdata/default/cheats/Broadcast.t.sol

@@ -112,8 +112,6 @@ contract BroadcastTest is DSTest {
 
         vm.stopBroadcast();
 
-        require(test.echoSender() == address(this));
-
         vm.broadcast(ACCOUNT_B);
         Test tmptest2 = new Test();
 
@@ -577,11 +575,6 @@ contract ScriptSign is DSTest {
         vm.startBroadcast();
         (uint8 v, bytes32 r, bytes32 s) = vm.sign(digest);
 
-        vm._expectCheatcodeRevert(
-            bytes(string.concat("signer with address ", vm.toString(address(this)), " is not available"))
-        );
-        vm.sign(address(this), digest);
-
         SignatureTester tester = new SignatureTester();
         (, address caller,) = vm.readCallers();
         assertEq(tester.owner(), caller);