add support for prank(sender, origin) and startPrank(sender, origin) cheatcodes

src/halmos/__main__.py

@@ -681,6 +681,7 @@ def run(
             jumpis={},
             symbolic=args.symbolic_storage,
             prank=Prank(),  # prank is reset after setUp()
+            origin=setup_ex.origin,
             #
             path=path,
             alias=setup_ex.alias.copy(),

src/halmos/cheatcodes.py

@@ -75,58 +75,93 @@ def stringified_bytes_to_bytes(hexstring: str) -> ByteVec:
     return ByteVec(ret_bytes)
 
 
+@dataclass(frozen=True)
+class PrankResult:
+    sender: Address | None = None
+    origin: Address | None = None
+
+    def __bool__(self) -> bool:
+        """
+        True iff either sender or origin is set.
+        """
+        return self.sender is not None or self.origin is not None
+
+    def __str__(self) -> str:
+        return f"{hexify(self.sender)}, {hexify(self.origin)}"
+
+
+NO_PRANK = PrankResult()
+
+
+@dataclass
 class Prank:
-    addr: Any  # prank address
-    keep: bool  # start / stop prank
+    """
+    A mutable object to store current prank context, one per execution context.
+
+    Because it's mutable, it must be copied across contexts.
 
-    def __init__(self, addr: Any = None, keep: bool = False) -> None:
-        if addr is not None:
-            assert_address(addr)
-        self.addr = addr
-        self.keep = keep
+    Can test for the existence of an active prank with `if prank: ...`
+
+    A prank is active if either sender or origin is set.
+    Technically supports pranking origin but not sender, which is not
+    possible with the current cheatcodes:
+    - prank(address) sets sender
+    - prank(address, address) sets both sender and origin
+    """
+
+    active: PrankResult = NO_PRANK  # active prank context
+    keep: bool = False  # start / stop prank
+
+    def __bool__(self) -> bool:
+        """
+        True iff either sender or origin is set.
+        """
+        return bool(self.active)
 
     def __str__(self) -> str:
-        if self.addr:
-            if self.keep:
-                return f"startPrank({str(self.addr)})"
-            else:
-                return f"prank({str(self.addr)})"
-        else:
-            return "None"
+        if not self:
+            return "no active prank"
+
+        fn_name = "startPrank" if self.keep else "prank"
+        return f"{fn_name}({str(self.active)})"
+
+    def lookup(self, to: Address) -> PrankResult:
+        """
+        If `to` is an eligible prank destination, return the active prank context.
+
+        If `keep` is False, this resets the prank context.
+        """
 
-    def lookup(self, this: Any, to: Any) -> Any:
-        assert_address(this)
         assert_address(to)
-        caller = this
         if (
-            self.addr is not None
+            self
             and not eq(to, hevm_cheat_code.address)
             and not eq(to, halmos_cheat_code.address)
         ):
-            caller = self.addr
+            result = self.active
             if not self.keep:
-                self.addr = None
-        return caller
+                self.stopPrank()
+            return result
 
-    def prank(self, addr: Any) -> bool:
-        assert_address(addr)
-        if self.addr is not None:
+        return NO_PRANK
+
+    def prank(self, sender: Address, origin: Address | None = None) -> bool:
+        assert_address(sender)
+        if self.active:
             return False
-        self.addr = addr
+
+        self.active = PrankResult(sender=sender, origin=origin)
         self.keep = False
         return True
 
-    def startPrank(self, addr: Any) -> bool:
-        assert_address(addr)
-        if self.addr is not None:
-            return False
-        self.addr = addr
-        self.keep = True
-        return True
+    def startPrank(self, sender: Address, origin: Address | None = None) -> bool:
+        result = self.prank(sender, origin)
+        self.keep = result if result else self.keep
+        return result
 
     def stopPrank(self) -> bool:
-        # stopPrank is allowed to call even when no active prank exists
-        self.addr = None
+        # stopPrank calls are allowed even when no active prank exists
+        self.active = NO_PRANK
         self.keep = False
         return True
 
@@ -282,9 +317,15 @@ class hevm_cheat_code:
     # bytes4(keccak256("prank(address)"))
     prank_sig: int = 0xCA669FA7
 
+    # bytes4(keccak256("prank(address,address)"))
+    prank_addr_addr_sig: int = 0x47E50CCE
+
     # bytes4(keccak256("startPrank(address)"))
     start_prank_sig: int = 0x06447D56
 
+    # bytes4(keccak256("startPrank(address,address)"))
+    start_prank_addr_addr_sig: int = 0x45B56078
+
     # bytes4(keccak256("stopPrank()"))
     stop_prank_sig: int = 0x90C5013B
 
@@ -381,8 +422,17 @@ def handle(sevm, ex, arg: ByteVec, stack, step_id) -> Optional[ByteVec]:
 
         # vm.prank(address)
         elif funsig == hevm_cheat_code.prank_sig:
-            address = uint160(arg.get_word(4))
-            result = ex.prank.prank(address)
+            sender = uint160(arg.get_word(4))
+            result = ex.prank.prank(sender)
+            if not result:
+                raise HalmosException("You have an active prank already.")
+            return ret
+
+        # vm.prank(address sender, address origin)
+        elif funsig == hevm_cheat_code.prank_addr_addr_sig:
+            sender = uint160(arg.get_word(4))
+            origin = uint160(arg.get_word(36))
+            result = ex.prank.prank(sender, origin)
             if not result:
                 raise HalmosException("You have an active prank already.")
             return ret
@@ -395,6 +445,15 @@ def handle(sevm, ex, arg: ByteVec, stack, step_id) -> Optional[ByteVec]:
                 raise HalmosException("You have an active prank already.")
             return ret
 
+        # vm.startPrank(address sender, address origin)
+        elif funsig == hevm_cheat_code.start_prank_addr_addr_sig:
+            sender = uint160(arg.get_word(4))
+            origin = uint160(arg.get_word(36))
+            result = ex.prank.startPrank(sender, origin)
+            if not result:
+                raise HalmosException("You have an active prank already.")
+            return ret
+
         # vm.stopPrank()
         elif funsig == hevm_cheat_code.stop_prank_sig:
             ex.prank.stopPrank()

src/halmos/sevm.py

@@ -630,6 +630,7 @@ class Exec:  # an execution path
     jumpis: Dict[str, Dict[bool, int]]  # for loop detection
     symbolic: bool  # symbolic or concrete storage
     prank: Prank
+    origin: Address
     addresses_to_delete: Set[Address]
 
     # path
@@ -662,6 +663,7 @@ def __init__(self, **kwargs) -> None:
         self.jumpis = kwargs["jumpis"]
         self.symbolic = kwargs["symbolic"]
         self.prank = kwargs["prank"]
+        self.origin = kwargs["origin"]
         self.addresses_to_delete = kwargs.get("addresses_to_delete") or set()
         #
         self.path = kwargs["path"]
@@ -724,6 +726,13 @@ def current_opcode(self) -> UnionType[int, BitVecRef]:
     def current_instruction(self) -> Instruction:
         return self.pgm.decode_instruction(self.pc)
 
+    def resolve_prank(self, to: Address) -> Tuple[Address, Address]:
+        # this potentially "consumes" the active prank
+        prank_result = self.prank.lookup(to)
+        caller = self.this if prank_result.sender is None else prank_result.sender
+        origin = f_origin() if prank_result.origin is None else prank_result.origin
+        return caller, origin
+
     def set_code(self, who: Address, code: UnionType[ByteVec, Contract]) -> None:
         """
         Sets the code at a given address.
@@ -1558,14 +1567,14 @@ def call(
         if not ret_size >= 0:
             raise ValueError(ret_size)
 
-        caller = ex.prank.lookup(ex.this, to)
+        pranked_caller, pranked_origin = ex.resolve_prank(to)
         arg = ex.st.memory.slice(arg_loc, arg_loc + arg_size)
 
         def send_callvalue(condition=None) -> None:
             # no balance update for CALLCODE which transfers to itself
             if op == EVM.CALL:
                 # TODO: revert if context is static
-                self.transfer_value(ex, caller, to, fund, condition)
+                self.transfer_value(ex, pranked_caller, to, fund, condition)
 
         def call_known(to: Address) -> None:
             # backup current state
@@ -1578,7 +1587,7 @@ def call_known(to: Address) -> None:
 
             message = Message(
                 target=to if op in [EVM.CALL, EVM.STATICCALL] else ex.this,
-                caller=caller if op != EVM.DELEGATECALL else ex.caller(),
+                caller=pranked_caller if op != EVM.DELEGATECALL else ex.caller(),
                 value=fund if op != EVM.DELEGATECALL else ex.callvalue(),
                 data=arg,
                 is_static=(ex.context.message.is_static or op == EVM.STATICCALL),
@@ -1610,6 +1619,7 @@ def callback(new_ex: Exec, stack, step_id):
                 new_ex.jumpis = deepcopy(ex.jumpis)
                 new_ex.symbolic = ex.symbolic
                 new_ex.prank = deepcopy(ex.prank)
+                new_ex.origin = ex.origin
 
                 # set return data (in memory)
                 effective_ret_size = min(ret_size, new_ex.returndatasize())
@@ -1652,6 +1662,7 @@ def callback(new_ex: Exec, stack, step_id):
                 jumpis={},
                 symbolic=ex.symbolic,
                 prank=Prank(),
+                origin=pranked_origin,
                 #
                 path=ex.path,
                 alias=ex.alias,
@@ -1775,7 +1786,7 @@ def call_unknown() -> None:
                 CallContext(
                     message=Message(
                         target=to,
-                        caller=caller,
+                        caller=pranked_caller,
                         value=fund,
                         data=ex.st.memory.slice(arg_loc, arg_loc + arg_size),
                         call_scheme=op,
@@ -1847,8 +1858,8 @@ def create(
         if op == EVM.CREATE2:
             salt = ex.st.pop()
 
-        # lookup prank
-        caller = ex.prank.lookup(ex.this, con_addr(0))
+        # check if there is an active prank
+        pranked_caller, pranked_origin = ex.resolve_prank(address(ex.this))
 
         # contract creation code
         create_hexcode = ex.st.memory.slice(loc, loc + size)
@@ -1867,14 +1878,16 @@ def create(
                 create_hexcode = bytes_to_bv_value(create_hexcode)
 
             code_hash = ex.sha3_data(create_hexcode)
-            hash_data = simplify(Concat(con(0xFF, 8), uint160(caller), salt, code_hash))
+            hash_data = simplify(
+                Concat(con(0xFF, 8), uint160(pranked_caller), salt, code_hash)
+            )
             new_addr = uint160(ex.sha3_data(hash_data))
         else:
             raise HalmosException(f"Unknown CREATE opcode: {op}")
 
         message = Message(
             target=new_addr,
-            caller=caller,
+            caller=pranked_caller,
             value=value,
             data=create_hexcode,
             is_static=False,
@@ -1908,7 +1921,7 @@ def create(
         ex.storage[new_addr] = {}  # existing storage may not be empty and reset here
 
         # transfer value
-        self.transfer_value(ex, caller, new_addr, value)
+        self.transfer_value(ex, pranked_caller, new_addr, value)
 
         def callback(new_ex, stack, step_id):
             subcall = new_ex.context
@@ -1972,6 +1985,7 @@ def callback(new_ex, stack, step_id):
             jumpis={},
             symbolic=False,
             prank=Prank(),
+            origin=pranked_origin,
             #
             path=ex.path,
             alias=ex.alias,
@@ -2092,6 +2106,7 @@ def create_branch(self, ex: Exec, cond: BitVecRef, target: int) -> Exec:
             jumpis=deepcopy(ex.jumpis),
             symbolic=ex.symbolic,
             prank=deepcopy(ex.prank),
+            origin=ex.origin,
             #
             path=new_path,
             alias=ex.alias.copy(),
@@ -2291,7 +2306,7 @@ def finalize(ex: Exec):
                     ex.st.push(uint256(ex.caller()))
 
                 elif opcode == EVM.ORIGIN:
-                    ex.st.push(uint256(f_origin()))
+                    ex.st.push(uint256(ex.origin))
 
                 elif opcode == EVM.ADDRESS:
                     ex.st.push(uint256(ex.this))
@@ -2626,6 +2641,7 @@ def mk_exec(
             jumpis={},
             symbolic=symbolic,
             prank=Prank(),
+            origin=f_origin(),
             #
             path=path,
             alias={},

src/halmos/utils.py

@@ -139,6 +139,10 @@ def int256(x: BitVecRef) -> BitVecRef:
     return simplify(SignExt(256 - bitsize, x))
 
 
+def address(x: Any) -> Address:
+    return uint(x, 160)
+
+
 def con(n: int, size_bits=256) -> Word:
     return BitVecVal(n, BitVecSorts[size_bits])
 

tests/expected/all.json

@@ -1337,7 +1337,7 @@
         ],
         "test/Prank.t.sol:PrankTest": [
             {
-                "name": "check_prank(address)",
+                "name": "check_prank(address,address)",
                 "exitcode": 0,
                 "num_models": 0,
                 "models": null,
@@ -1354,6 +1354,15 @@
                 "time": null,
                 "num_bounded_loops": null
             },
+            {
+                "name": "check_prank_ConstructorCreate2(address,bytes32)",
+                "exitcode": 0,
+                "num_models": 0,
+                "models": null,
+                "num_paths": null,
+                "time": null,
+                "num_bounded_loops": null
+            },
             {
                 "name": "check_prank_External(address)",
                 "exitcode": 0,
@@ -1409,7 +1418,7 @@
                 "num_bounded_loops": null
             },
             {
-                "name": "check_startPrank(address)",
+                "name": "check_startPrank(address,address)",
                 "exitcode": 0,
                 "num_models": 0,
                 "models": null,