Merge remote-tracking branch 'origin/storage/refactor' into berndtz_experimental

mythril/laser/ethereum/instructions.py

@@ -56,8 +56,6 @@ log = logging.getLogger(__name__)
 TT256 = 2 ** 256
 TT256M1 = 2 ** 256 - 1
 
-keccak_function_manager = KeccakFunctionManager()
-
 
 class StateTransition(object):
     """Decorator that handles global state copy and original return.
@@ -903,7 +901,6 @@ class Instruction:
         :param global_state:
         :return:
         """
-        global keccak_function_manager
 
         state = global_state.mstate
         op0, op1 = state.stack.pop(), state.stack.pop()
@@ -958,7 +955,6 @@ class Instruction:
             )
             log.debug("Created BitVecFunc hash.")
 
-            keccak_function_manager.add_keccak(result, state.memory[index])
         else:
             keccak = utils.sha3(data.value.to_bytes(length, byteorder="big"))
             result = symbol_factory.BitVecFuncVal(
@@ -1374,52 +1370,44 @@ class Instruction:
         :param global_state:
         :return:
         """
-        global keccak_function_manager
 
         state = global_state.mstate
         index = state.stack.pop()
         log.debug("Storage access at index " + str(index))
+        index = simplify(index)
 
-        try:
-            index = util.get_concrete_int(index)
+        if not KeccakFunctionManager.is_keccak(index):
             return self._sload_helper(global_state, index)
 
-        except TypeError:
-            if not keccak_function_manager.is_keccak(index):
-                return self._sload_helper(global_state, str(index))
+        storage_keys = global_state.environment.active_account.storage.keys()
+        keccak_keys = list(filter(KeccakFunctionManager.is_keccak, storage_keys))
 
-            storage_keys = global_state.environment.active_account.storage.keys()
-            keccak_keys = list(filter(keccak_function_manager.is_keccak, storage_keys))
+        results = []  # type: List[GlobalState]
+        constraints = []
 
-            results = []  # type: List[GlobalState]
-            constraints = []
+        for keccak_key in keccak_keys:
+            key_argument = KeccakFunctionManager.get_argument(keccak_key)
+            index_argument = KeccakFunctionManager.get_argument(index)
+            if key_argument.size() != index_argument.size():
+                continue
+            constraints.append((keccak_key, key_argument == index_argument))
 
-            for keccak_key in keccak_keys:
-                key_argument = keccak_function_manager.get_argument(keccak_key)
-                index_argument = keccak_function_manager.get_argument(index)
-                constraints.append((keccak_key, key_argument == index_argument))
+        for (keccak_key, constraint) in constraints:
+            if constraint in state.constraints:
+                results += self._sload_helper(global_state, keccak_key, [constraint])
+        if len(results) > 0:
+            return results
 
-            for (keccak_key, constraint) in constraints:
-                if constraint in state.constraints:
-                    results += self._sload_helper(
-                        global_state, keccak_key, [constraint]
-                    )
-            if len(results) > 0:
-                return results
+        for (keccak_key, constraint) in constraints:
+            results += self._sload_helper(copy(global_state), keccak_key, [constraint])
 
-            for (keccak_key, constraint) in constraints:
-                results += self._sload_helper(
-                    copy(global_state), keccak_key, [constraint]
-                )
-            if len(results) > 0:
-                return results
+        if len(results) > 0:
+            return results
 
-            return self._sload_helper(global_state, str(index))
+        return self._sload_helper(global_state, index)
 
     @staticmethod
-    def _sload_helper(
-        global_state: GlobalState, index: Union[str, int], constraints=None
-    ):
+    def _sload_helper(global_state: GlobalState, index: Expression, constraints=None):
         """
 
         :param global_state:
@@ -1435,6 +1423,8 @@ class Instruction:
 
         if constraints is not None:
             global_state.mstate.constraints += constraints
+            if global_state.mstate.constraints.is_possible is False:
+                return []
 
         global_state.mstate.stack.append(data)
         return [global_state]
@@ -1447,11 +1437,10 @@ class Instruction:
         :param this_key:
         :param argument:
         """
-        global keccak_function_manager
         for keccak_key in keccak_keys:
             if keccak_key == this_key:
                 continue
-            keccak_argument = keccak_function_manager.get_argument(keccak_key)
+            keccak_argument = KeccakFunctionManager.get_argument(keccak_key)
             yield keccak_argument != argument
 
     @StateTransition()
@@ -1461,62 +1450,43 @@ class Instruction:
         :param global_state:
         :return:
         """
-        global keccak_function_manager
         state = global_state.mstate
         index, value = state.stack.pop(), state.stack.pop()
         log.debug("Write to storage[" + str(index) + "]")
-
-        try:
-            index = util.get_concrete_int(index)
+        index = simplify(index)
+        is_keccak = KeccakFunctionManager.is_keccak(index)
+        if not is_keccak:
             return self._sstore_helper(global_state, index, value)
-        except TypeError:
-            is_keccak = keccak_function_manager.is_keccak(index)
-            if not is_keccak:
-                return self._sstore_helper(global_state, str(index), value)
-
-            storage_keys = global_state.environment.active_account.storage.keys()
-            keccak_keys = filter(keccak_function_manager.is_keccak, storage_keys)
-
-            results = []  # type: List[GlobalState]
-            new = symbol_factory.Bool(False)
-
-            for keccak_key in keccak_keys:
-                key_argument = keccak_function_manager.get_argument(
-                    keccak_key
-                )  # type: Expression
-                index_argument = keccak_function_manager.get_argument(
-                    index
-                )  # type: Expression
-                condition = key_argument == index_argument
-                condition = (
-                    condition
-                    if type(condition) == bool
-                    else is_true(simplify(cast(Bool, condition)))
-                )
-                if condition:
-                    return self._sstore_helper(
-                        copy(global_state),
-                        keccak_key,
-                        value,
-                        key_argument == index_argument,
-                    )
 
-                results += self._sstore_helper(
-                    copy(global_state),
-                    keccak_key,
-                    value,
-                    key_argument == index_argument,
-                )
+        storage_keys = global_state.environment.active_account.storage.keys()
 
-                new = Or(new, cast(Bool, key_argument != index_argument))
+        keccak_keys = filter(KeccakFunctionManager.is_keccak, storage_keys)
 
-            if len(results) > 0:
-                results += self._sstore_helper(
-                    copy(global_state), str(index), value, new
-                )
-                return results
+        results = []  # type: List[GlobalState]
+        new = symbol_factory.Bool(False)
+
+        for keccak_key in keccak_keys:
+            key_argument = KeccakFunctionManager.get_argument(
+                keccak_key
+            )  # type: Expression
+            index_argument = KeccakFunctionManager.get_argument(
+                index
+            )  # type: Expression
+            if key_argument.size() != index_argument.size():
+                continue
+
+            result = self._sstore_helper(
+                copy(global_state), keccak_key, value, key_argument == index_argument
+            )
+            if len(result) > 0:
+                return result
 
-            return self._sstore_helper(global_state, str(index), value)
+            new = Or(new, cast(Bool, key_argument != index_argument))
+        results += self._sstore_helper(copy(global_state), index, value, new)
+        if len(results) > 0:
+            return results
+        # results += self._sstore_helper(copy(global_state), index, value)
+        return self._sstore_helper(copy(global_state), index, value)
 
     @staticmethod
     def _sstore_helper(global_state, index, value, constraint=None):
@@ -1544,6 +1514,8 @@ class Instruction:
 
         if constraint is not None:
             global_state.mstate.constraints.append(constraint)
+            if global_state.mstate.constraints.is_possible is False:
+                return []
 
         return [global_state]
 

mythril/laser/ethereum/keccak.py

@@ -1,38 +1,30 @@
 """This module contains a function manager to deal with symbolic Keccak
 values."""
-from mythril.laser.smt import Expression
+from mythril.laser.smt import Expression, BitVecFunc, BitVec
+from typing import cast, Union
 
 
 class KeccakFunctionManager:
     """A keccak function manager for symbolic expressions."""
 
-    def __init__(self):
-        """"""
-        self.keccak_expression_mapping = {}
-
-    def is_keccak(self, expression: Expression) -> bool:
+    @staticmethod
+    def is_keccak(expression: Expression) -> bool:
         """
 
         :param expression:
         :return:
         """
-        return str(expression) in self.keccak_expression_mapping.keys()
+        if not isinstance(expression, BitVecFunc):
+            return False
+        return expression.func_name == "keccak256"
 
-    def get_argument(self, expression: Expression) -> Expression:
+    @staticmethod
+    def get_argument(expression: Expression) -> Expression:
         """
 
         :param expression:
         :return:
         """
-        if not self.is_keccak(expression):
-            raise ValueError("Expression is not a recognized keccac result")
-        return self.keccak_expression_mapping[str(expression)][1]
-
-    def add_keccak(self, expression: Expression, argument: Expression) -> None:
-        """
-
-        :param expression:
-        :param argument:
-        """
-        index = str(expression)
-        self.keccak_expression_mapping[index] = (expression, argument)
+        if not KeccakFunctionManager.is_keccak(expression):
+            raise ValueError("Expression is not a recognized keccak result")
+        return cast(BitVecFunc, expression).input_

mythril/laser/ethereum/state/account.py

@@ -7,7 +7,7 @@ from typing import Any, Dict, KeysView, Union
 
 from z3 import ExprRef
 
-from mythril.laser.smt import Array, symbol_factory, BitVec
+from mythril.laser.smt import Array, symbol_factory, BitVec, Expression
 from mythril.disassembler.disassembly import Disassembly
 from mythril.laser.smt import symbol_factory
 
@@ -20,12 +20,12 @@ class Storage:
 
         :param concrete: bool indicating whether to interpret uninitialized storage as concrete versus symbolic
         """
-        self._storage = {}  # type: Dict[Union[int, str], Any]
+        self._storage = {}  # type: Dict[Expression, Any]
         self.concrete = concrete
         self.dynld = dynamic_loader
         self.address = address
 
-    def __getitem__(self, item: Union[str, int]) -> Any:
+    def __getitem__(self, item: Expression) -> Any:
         try:
             return self._storage[item]
         except KeyError:
@@ -38,7 +38,7 @@ class Storage:
                     self._storage[item] = symbol_factory.BitVecVal(
                         int(
                             self.dynld.read_storage(
-                                contract_address=self.address, index=int(item)
+                                contract_address=self.address, index=int(item.raw)
                             ),
                             16,
                         ),
@@ -56,7 +56,7 @@ class Storage:
         )
         return self._storage[item]
 
-    def __setitem__(self, key: Union[int, str], value: Any) -> None:
+    def __setitem__(self, key: Expression, value: Any) -> None:
         self._storage[key] = value
 
     def keys(self) -> KeysView:

mythril/laser/smt/__init__.py

@@ -72,7 +72,7 @@ class SymbolFactory(Generic[T, U]):
         func_name: str,
         size: int,
         annotations: Annotations = None,
-        input_: Union[int, "BitVec"] = None,
+        input_: Expression = None,
     ) -> BitVecFunc:
         """Creates a new bit vector function with a symbolic value.
 
@@ -91,7 +91,7 @@ class SymbolFactory(Generic[T, U]):
         func_name: str,
         size: int,
         annotations: Annotations = None,
-        input_: Union[int, "BitVec"] = None,
+        input_: Expression = None,
     ) -> BitVecFunc:
         """Creates a new bit vector function with a symbolic value.
 
@@ -140,7 +140,7 @@ class _SmtSymbolFactory(SymbolFactory[bool.Bool, BitVec]):
         func_name: str,
         size: int,
         annotations: Annotations = None,
-        input_: Union[int, "BitVec"] = None,
+        input_: Expression = None,
     ) -> BitVecFunc:
         """Creates a new bit vector function with a concrete value."""
         raw = z3.BitVecVal(value, size)
@@ -152,7 +152,7 @@ class _SmtSymbolFactory(SymbolFactory[bool.Bool, BitVec]):
         func_name: str,
         size: int,
         annotations: Annotations = None,
-        input_: Union[int, "BitVec"] = None,
+        input_: Expression = None,
     ) -> BitVecFunc:
         """Creates a new bit vector function with a symbolic value."""
         raw = z3.BitVec(name, size)

mythril/laser/smt/bitvec.py

@@ -243,6 +243,9 @@ class BitVec(Expression[z3.BitVecRef]):
         """
         return self._handle_shift(other, rshift)
 
+    def __hash__(self):
+        return self.raw.__hash__()
+
 
 def _comparison_helper(
     a: BitVec, b: BitVec, operation: Callable, default_value: bool, inputs_equal: bool

mythril/laser/smt/bitvecfunc.py

@@ -2,6 +2,7 @@ from typing import Optional, Union, cast, Callable
 
 import z3
 
+from mythril.laser.smt.expression import Expression
 from mythril.laser.smt.bitvec import BitVec, Bool, And, Annotations
 from mythril.laser.smt.bool import Or
 
@@ -26,6 +27,7 @@ def _arithmetic_helper(
     union = a.annotations + b.annotations
 
     if isinstance(b, BitVecFunc):
+
         # TODO: Find better value to set input and name to in this case?
         return BitVecFunc(raw=raw, func_name=None, input_=None, annotations=union)
 
@@ -79,7 +81,7 @@ class BitVecFunc(BitVec):
         self,
         raw: z3.BitVecRef,
         func_name: Optional[str],
-        input_: Union[int, "BitVec"] = None,
+        input_: Expression = None,
         annotations: Optional[Annotations] = None,
     ):
         """
@@ -223,3 +225,6 @@ class BitVecFunc(BitVec):
         :return The resulting right shifted output:
         """
         return _arithmetic_helper(self, other, operator.rshift)
+
+    def __hash__(self) -> int:
+        return self.raw.__hash__()

mythril/laser/smt/bool.py

@@ -80,6 +80,9 @@ class Bool(Expression[z3.BoolRef]):
         else:
             return False
 
+    def __hash__(self) -> int:
+        return self.raw.__hash__()
+
 
 def And(*args: Union[Bool, bool]) -> Bool:
     """Create an And expression."""

mythril/laser/smt/expression.py

@@ -45,6 +45,12 @@ class Expression(Generic[T]):
     def __repr__(self) -> str:
         return repr(self.raw)
 
+    def size(self):
+        return self.raw.size()
+
+    def __hash__(self) -> int:
+        return self.raw.__hash__()
+
 
 G = TypeVar("G", bound=Expression)
 

tests/laser/evm_testsuite/evm_test.py

@@ -125,7 +125,8 @@ def test_vmtest(
         account.code = Disassembly(details["code"][2:])
         account.nonce = int(details["nonce"], 16)
         for key, value in details["storage"].items():
-            account.storage[int(key, 16)] = int(value, 16)
+            key_bitvec = symbol_factory.BitVecVal(int(key, 16), 256)
+            account.storage[key_bitvec] = int(value, 16)
 
         world_state.put_account(account)
         account.set_balance(int(details["balance"], 16))
@@ -175,8 +176,7 @@ def test_vmtest(
 
             for index, value in details["storage"].items():
                 expected = int(value, 16)
-                actual = account.storage[int(index, 16)]
-
+                actual = account.storage[symbol_factory.BitVecVal(int(index, 16), 256)]
                 if isinstance(actual, Expression):
                     actual = actual.value
                     actual = 1 if actual is True else 0 if actual is False else actual