Save code to world state (#1314)

* Save code

* Refactor code with black

* return the account as the function describes

* use bytecode param

* fix missing edge case in conditional

* use world state helper function over custom code

* back with black

* Refactor test file with black

* Align the type: ignore comment

* Fix black again

* Raise ValueError if dynamic loader is None

Co-authored-by: JoranHonig <JoranHonig@users.noreply.github.com>

mythril/analysis/modules/dependence_on_predictable_vars.py

@@ -193,7 +193,9 @@ class PredictableDependenceModule(DetectionModule):
 
                     # Why the second constraint? Because without it Z3 returns a solution where param overflows.
 
-                    solver.get_model(state.world_state.constraints + constraint)  # type: ignore
+                    solver.get_model(
+                        state.world_state.constraints + constraint  # type: ignore
+                    )
                     state.annotate(OldBlockNumberUsedAnnotation(constraint))
 
                 except UnsatError:

mythril/laser/ethereum/call.py

@@ -53,10 +53,9 @@ def get_call_parameters(
 
     callee_account = None
     call_data = get_call_data(global_state, memory_input_offset, memory_input_size)
-    if (
-        isinstance(callee_address, BitVec)
-        or int(callee_address, 16) > PRECOMPILE_COUNT
-        or int(callee_address, 16) == 0
+    if isinstance(callee_address, BitVec) or (
+        isinstance(callee_address, str)
+        and (int(callee_address, 16) > PRECOMPILE_COUNT or int(callee_address, 16) == 0)
     ):
         callee_account = get_callee_account(
             global_state, callee_address, dynamic_loader
@@ -141,37 +140,9 @@ def get_callee_account(
         else:
             callee_address = hex(callee_address.value)[2:]
 
-    try:
-        return global_state.accounts[int(callee_address, 16)]
-    except KeyError:
-        # We have a valid call address, but contract is not in the modules list
-        log.debug("Module with address %s not loaded.", callee_address)
-
-    if dynamic_loader is None:
-        raise ValueError("dynamic_loader is None")
-
-    log.debug("Attempting to load dependency")
-
-    try:
-        code = dynamic_loader.dynld(callee_address)
-    except ValueError as error:
-        log.debug("Unable to execute dynamic loader because: %s", error)
-        raise error
-    if code is None:
-        log.debug("No code returned, not a contract account?")
-        raise ValueError("No code returned")
-    log.debug("Dependency loaded: " + callee_address)
-
-    callee_account = Account(
-        symbol_factory.BitVecVal(int(callee_address, 16), 256),
-        code,
-        callee_address,
-        dynamic_loader=dynamic_loader,
-        balances=global_state.world_state.balances,
+    return global_state.world_state.accounts_exist_or_load(
+        callee_address, dynamic_loader
     )
-    global_state.accounts[int(callee_address, 16)] = callee_account
-
-    return callee_account
 
 
 def get_call_data(

mythril/laser/ethereum/instructions.py

@@ -1140,7 +1140,7 @@ class Instruction:
         try:
             code = global_state.world_state.accounts_exist_or_load(
                 addr, self.dynamic_loader
-            )
+            ).code.bytecode
         except (ValueError, AttributeError) as e:
             log.debug("error accessing contract storage due to: " + str(e))
             state.stack.append(global_state.new_bitvec("extcodesize_" + str(addr), 256))
@@ -1236,7 +1236,7 @@ class Instruction:
         try:
             code = global_state.world_state.accounts_exist_or_load(
                 addr, self.dynamic_loader
-            )
+            ).code.bytecode
         except (ValueError, AttributeError) as e:
             log.debug("error accessing contract storage due to: " + str(e))
             return [global_state]
@@ -1267,7 +1267,9 @@ class Instruction:
             code_hash = symbol_factory.BitVecVal(0, 256)
         else:
             addr = "0" * (40 - len(hex(address.value)[2:])) + hex(address.value)[2:]
-            code = world_state.accounts_exist_or_load(addr, self.dynamic_loader)
+            code = world_state.accounts_exist_or_load(
+                addr, self.dynamic_loader
+            ).code.bytecode
             code_hash = symbol_factory.BitVecVal(int(get_code_hash(code), 16), 256)
         stack.append(code_hash)
         return [global_state]
@@ -1870,7 +1872,7 @@ class Instruction:
         for i in range(memory_out_size.value):
             global_state.mstate.memory[memory_out_offset + i] = global_state.new_bitvec(
                 "call_output_var({})_{}".format(
-                    simplify(memory_out_offset + i), global_state.mstate.pc,
+                    simplify(memory_out_offset + i), global_state.mstate.pc
                 ),
                 8,
             )

mythril/laser/ethereum/state/account.py

@@ -7,13 +7,7 @@ from copy import copy, deepcopy
 from typing import Any, Dict, Union, Set
 
 
-from mythril.laser.smt import (
-    Array,
-    K,
-    BitVec,
-    simplify,
-    BaseArray,
-)
+from mythril.laser.smt import Array, K, BitVec, simplify, BaseArray
 from mythril.disassembler.disassembly import Disassembly
 from mythril.laser.smt import symbol_factory
 

mythril/laser/ethereum/state/world_state.py

@@ -73,7 +73,7 @@ class WorldState:
         new_world_state.constraints = copy(self.constraints)
         return new_world_state
 
-    def accounts_exist_or_load(self, addr: str, dynamic_loader: DynLoader) -> str:
+    def accounts_exist_or_load(self, addr: str, dynamic_loader: DynLoader) -> Account:
         """
         returns account if it exists, else it loads from the dynamic loader
         :param addr: address
@@ -81,18 +81,17 @@ class WorldState:
         :return: The code
         """
         addr_bitvec = symbol_factory.BitVecVal(int(addr, 16), 256)
+
         if addr_bitvec.value in self.accounts:
-            code = self.accounts[addr_bitvec.value].code
-        else:
-            code = dynamic_loader.dynld(addr)
-            self.create_account(
-                balance=0, address=addr_bitvec.value, dynamic_loader=dynamic_loader
-            )
-        if code is None:
-            code = ""
-        else:
-            code = code.bytecode
-        return code
+            return self.accounts[addr_bitvec.value]
+        if dynamic_loader is None:
+            raise ValueError("dynamic_loader is None")
+        return self.create_account(
+            balance=0,
+            address=addr_bitvec.value,
+            dynamic_loader=dynamic_loader,
+            code=dynamic_loader.dynld(addr),
+        )
 
     def create_account(
         self,
@@ -101,6 +100,7 @@ class WorldState:
         concrete_storage=False,
         dynamic_loader=None,
         creator=None,
+        code=None,
     ) -> Account:
         """Create non-contract account.
 
@@ -108,6 +108,8 @@ class WorldState:
         :param balance: Initial balance for the account
         :param concrete_storage: Interpret account storage as concrete
         :param dynamic_loader: used for dynamically loading storage from the block chain
+        :param creator: The address of the creator of the contract if it's a contract
+        :param code: The code of the contract, if it's a contract
         :return: The new account
         """
         address = (
@@ -122,6 +124,8 @@ class WorldState:
             dynamic_loader=dynamic_loader,
             concrete_storage=concrete_storage,
         )
+        if code:
+            new_account.code = code
         if balance:
             new_account.add_balance(symbol_factory.BitVecVal(balance, 256))
 

tests/laser/state/world_state_account_exist_load_test.py

@@ -46,5 +46,7 @@ def _get_global_state():
 def test_extraction(addr, eth, code_len):
     global_state = _get_global_state()
     dynamic_loader = DynLoader(eth=eth)
-    code = global_state.world_state.accounts_exist_or_load(addr, dynamic_loader)
+    code = global_state.world_state.accounts_exist_or_load(
+        addr, dynamic_loader
+    ).code.bytecode
     assert len(code) == code_len