Revise laser to work using keccak store

mythril/laser/ethereum/instructions.py

@@ -2,22 +2,24 @@ import binascii
 import logging
 from copy import copy, deepcopy
 
-import ethereum.opcodes as opcodes
 from ethereum import utils
 from z3 import Extract, UDiv, simplify, Concat, ULT, UGT, BitVecNumRef, Not, \
-    is_false, is_expr, ExprRef, URem, SRem
+    is_false, is_expr, ExprRef, URem, SRem, BitVec, Solver, is_true, BitVecVal, If, BoolRef, Or
-from z3 import BitVecVal, If, BoolRef
 
 import mythril.laser.ethereum.util as helper
 from mythril.laser.ethereum import util
 from mythril.laser.ethereum.call import get_call_parameters
-from mythril.laser.ethereum.state import GlobalState, MachineState, Environment, CalldataType
+from mythril.laser.ethereum.state import GlobalState, CalldataType
 import mythril.laser.ethereum.natives as natives
-from mythril.laser.ethereum.transaction import MessageCallTransaction, TransactionEndSignal, TransactionStartSignal, ContractCreationTransaction
+from mythril.laser.ethereum.transaction import MessageCallTransaction, TransactionStartSignal, ContractCreationTransaction
+from mythril.laser.ethereum.keccak import KeccakFunctionManager
+
 
 TT256 = 2 ** 256
 TT256M1 = 2 ** 256 - 1
 
+keccak_function_manager: KeccakFunctionManager = KeccakFunctionManager()
+
 
 class StackUnderflowException(Exception):
     pass
@@ -341,12 +343,12 @@ class Instruction:
         except AttributeError:
             logging.debug("CALLDATALOAD: Unsupported symbolic index")
             state.stack.append(global_state.new_bitvec(
-                "calldata_" + str(environment.active_account.contract_name) + str(len(global_state.current_transaction.caller)) +"[" + str(simplify(op0)) + "]", 256))
+                "calldata_" + str(environment.active_account.contract_name) + "[" + str(simplify(op0)) + "]", 256))
             return [global_state]
         except IndexError:
             logging.debug("Calldata not set, using symbolic variable instead")
             state.stack.append(global_state.new_bitvec(
-                "calldata_" + str(environment.active_account.contract_name) +str(global_state.current_transaction.caller)+ "[" + str(simplify(op0)) + "]", 256))
+                "calldata_" + str(environment.active_account.contract_name) + "[" + str(simplify(op0)) + "]", 256))
             return [global_state]
 
         if type(b) == int:
@@ -482,6 +484,8 @@ class Instruction:
 
     @instruction
     def sha3_(self, global_state):
+        global keccak_function_manager
+
         state = global_state.mstate
         environment = global_state.environment
         op0, op1 = state.stack.pop(), state.stack.pop()
@@ -493,7 +497,7 @@ class Instruction:
             # Can't access symbolic memory offsets
             if is_expr(op0):
                 op0 = simplify(op0)
-            state.stack.append(global_state.new_bitvec("KECCAC_mem[" + str(op0) + "]", 256))
+            state.stack.append(BitVec("KECCAC_mem[" + str(op0) + "]", 256))
             return [global_state]
 
         try:
@@ -504,18 +508,17 @@ class Instruction:
                 i += 1
             # FIXME: broad exception catch
         except:
+            argument = str(state.memory[index]).replace(" ", "_")
 
-            svar = str(state.memory[index])
+            result = BitVec("KECCAC[{}]".format(argument), 256)
-
+            keccak_function_manager.add_keccak(result, state.memory[index])
-            svar = svar.replace(" ", "_")
+            state.stack.append(result)
-
-            state.stack.append(global_state.new_bitvec("keccac_" + svar, 256))
             return [global_state]
 
-        keccac = utils.sha3(utils.bytearray_to_bytestr(data))
+        keccak = utils.sha3(utils.bytearray_to_bytestr(data))
-        logging.debug("Computed SHA3 Hash: " + str(binascii.hexlify(keccac)))
+        logging.debug("Computed SHA3 Hash: " + str(binascii.hexlify(keccak)))
 
-        state.stack.append(BitVecVal(util.concrete_int_from_bytes(keccac, 0), 256))
+        state.stack.append(BitVecVal(util.concrete_int_from_bytes(keccak, 0), 256))
         return [global_state]
 
     @instruction
@@ -727,26 +730,68 @@ class Instruction:
 
     @instruction
     def sload_(self, global_state):
+        global keccak_function_manager
+        
         state = global_state.mstate
         index = state.stack.pop()
         logging.debug("Storage access at index " + str(index))
 
         try:
             index = util.get_concrete_int(index)
+            return self._sload_helper(global_state, index)
+
         except AttributeError:
-            index = str(index)
+            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(keccak_function_manager.is_keccak, storage_keys))
+
+            results = []
+            constraints = []
+
+            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:
+                results += self._sload_helper(copy(global_state), keccak_key, [constraint])
+            if len(results) > 0:
+                return results
+            
+            return self._sload_helper(global_state, str(index))
 
+    def _sload_helper(self, global_state, index, constraints=None):
         try:
             data = global_state.environment.active_account.storage[index]
         except KeyError:
             data = global_state.new_bitvec("storage_" + str(index), 256)
             global_state.environment.active_account.storage[index] = data
 
-        state.stack.append(data)
+        if constraints is not None:
+            global_state.mstate.constraints += constraints
+
+        global_state.mstate.stack.append(data)
         return [global_state]
 
+    def _get_constraints(self, keccak_keys, this_key, 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)
+            yield keccak_argument != argument
+
     @instruction
     def sstore_(self, global_state):
+        global keccak_function_manager
         state = global_state.mstate
         index, value = state.stack.pop(), state.stack.pop()
 
@@ -754,17 +799,52 @@ class Instruction:
 
         try:
             index = util.get_concrete_int(index)
+            return self._sstore_helper(global_state, index, value)
         except AttributeError:
-            index = str(index)
+            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)
+
+            solver = Solver()
+            solver.set(timeout=1000)
+
+            results = []
+            new = False
+
+            for keccak_key in keccak_keys:
+                key_argument = keccak_function_manager.get_argument(keccak_key)
+                index_argument = keccak_function_manager.get_argument(index)
 
+                if is_true(key_argument == index_argument):
+                    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)
+
+                new = Or(new, key_argument != index_argument)
+
+            if len(results) > 0:
+                results += self._sstore_helper(copy(global_state), str(index), value, new)
+                return results
+
+            return self._sstore_helper(global_state, str(index), value)
+
+    def _sstore_helper(self, global_state, index, value, constraint=None):
         try:
             global_state.environment.active_account = deepcopy(global_state.environment.active_account)
             global_state.accounts[
                 global_state.environment.active_account.address] = global_state.environment.active_account
 
-            global_state.environment.active_account.storage[index] = value
+            global_state.environment.active_account.storage[index] =\
+                value if not isinstance(value, ExprRef) else simplify(value)
         except KeyError:
             logging.debug("Error writing to storage: Invalid index")
+
+        if constraint is not None:
+            global_state.mstate.constraints.append(constraint)
+
         return [global_state]
 
     @instruction

mythril/laser/ethereum/svm.py

@@ -185,6 +185,10 @@ class LaserEVM:
         elif opcode == "JUMPI":
             for state in new_states:
                 self._new_node_state(state, JumpType.CONDITIONAL, state.mstate.constraints[-1])
+        elif opcode in ("SLOAD", "SSTORE") and len(new_states) > 1:
+            for state in new_states:
+                self._new_node_state(state, JumpType.CONDITIONAL, state.mstate.constraints[-1])
+
         elif opcode in ("CALL", 'CALLCODE', 'DELEGATECALL', 'STATICCALL'):
             assert len(new_states) <= 1
             for state in new_states: