feat: add getCode cheatcode

.gitignore

@@ -14,6 +14,7 @@ __pycache__/
 
 # Environments
 .env
+env/
 .venv
 venv/
 

src/halmos/__main__.py

@@ -438,6 +438,7 @@ def main() -> int:
         'modByConst': args.smt_mod_by_const,
         'expByConst': args.smt_exp_by_const,
         'timeout': args.solver_timeout_branching,
+        'target': args.target,
     }
 
     if args.width is not None:

src/halmos/sevm.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: AGPL-3.0
 
+import json
 import math
 
 from copy import deepcopy
@@ -707,6 +708,11 @@ def call_unknown() -> None:
             ex.solver.add(exit_code_var == exit_code)
             ex.st.push(exit_code_var)
 
+            ret = None
+            if ret_size > 0:
+                f_ret = Function('ret_'+str(ret_size*8), BitVecSort(256), BitVecSort(ret_size*8))
+                ret = f_ret(exit_code_var)
+
             # TODO: cover other precompiled
             if to == con(1): # ecrecover exit code is always 1
                 ex.solver.add(exit_code_var != con(0))
@@ -722,22 +728,50 @@ def call_unknown() -> None:
                     assume_cond = simplify(is_non_zero(Extract(255, 0, arg)))
                     ex.solver.add(assume_cond)
                     ex.path.append(str(assume_cond))
+                elif simplify(Extract(arg_size*8-1, arg_size*8-32, arg)) == hevm_cheat_code.get_code_sig:
+                    calldata = bytes.fromhex(hex(arg.as_long())[2:])
+                    path_len = int.from_bytes(calldata[36:68], 'big')
+                    path = calldata[68:68+path_len].decode('utf-8')
+
+                    if ':' in path:
+                        [filename, contract_name] = path.split(':')
+                        path = 'out/' + filename + '/' + contract_name + '.json'
+
+                    target = self.options['target'].rstrip('/')
+                    path = target + '/' + path
+
+                    with open(path) as f:
+                        artifact = json.loads(f.read())
+
+
+                    if artifact['bytecode']['object']:
+                        bytecode = artifact['bytecode']['object'].replace('0x', '')
+                    else:
+                        bytecode = artifact['bytecode'].replace('0x', '')
+
+                    bytecode_len = (len(bytecode) + 1) // 2
+                    bytecode_len_enc = hex(bytecode_len).replace('0x', '').rjust(64, '0')
+
+                    bytecode_len_ceil = (bytecode_len + 31) // 32 * 32
+
+                    ret_bytes = '00' * 31 + '20' + bytecode_len_enc + bytecode.ljust(bytecode_len_ceil*2, '0')
+                    ret_len = len(ret_bytes) // 2
+                    ret_bytes = bytes.fromhex(ret_bytes)
+
+                    ret = BitVecVal(int.from_bytes(ret_bytes, 'big'), ret_len * 8)
                 else:
                     # TODO: support other cheat codes
                     ex.error = str('Unsupported cheat code: calldata: ' + str(arg))
                     out.append(ex)
                     return
 
-            # TODO: The actual return data size may be different from the given ret_size.
-            #       In that case, ex.output should be set to the actual return data.
-            #       And, if the actual size is smaller than the given size, then the memory is updated only up to the actual size.
-
+            # TODO: handle inconsistent return sizes for unknown functions
             # store return value
-            if ret_size > 0:
-                f_ret = Function('ret_'+str(ret_size*8), BitVecSort(256), BitVecSort(ret_size*8))
-                ret = f_ret(exit_code_var)
+            if ret_size > 0 and ret != None:
                 wstore(ex.st.memory, ret_loc, ret_size, ret)
                 ex.output = ret
+            elif ret != None:
+                ex.output = ret
             else:
                 ex.output = None
 

src/halmos/utils.py

@@ -252,6 +252,9 @@ class hevm_cheat_code:
     # bytes4(keccak256("assume(bool)"))
     assume_sig: int = 0x4C63E562
 
+    # bytes4(keccak256("getCode(string)))
+    get_code_sig: int = 0x8d1cc925
+
 sha3_inv: Dict[int, int] = { # sha3(x) -> x
     0x290decd9548b62a8d60345a988386fc84ba6bc95484008f6362f93160ef3e563 : 0,
     0xb10e2d527612073b26eecdfd717e6a320cf44b4afac2b0732d9fcbe2b7fa0cf6 : 1,

tests/test/Foundry.t.sol

@@ -2,6 +2,9 @@
 pragma solidity >=0.8.0 <0.9.0;
 
 import "forge-std/Test.sol";
+import "forge-std/StdCheats.sol";
+
+import "../src/Counter.sol";
 
 contract FoundryTest is Test {
     /* TODO: support testFail prefix
@@ -15,6 +18,16 @@ contract FoundryTest is Test {
         assertLt(x, 100);
     }
 
+    function testGetCode(uint x) public {
+        Counter counter = Counter(deployCode("./out/Counter.sol/Counter.json"));
+        counter.set(x);
+        assertEq(counter.cnt(), x);
+
+        Counter counter2 = Counter(deployCode("Counter.sol:Counter"));
+        counter2.set(x);
+        assertEq(counter2.cnt(), x);
+    }
+
     /* TODO: support vm.prank cheatcode
     function testPrank(address x) public {
         vm.prank(x); // not supported