Merge pull request #499 from dmuhs/fix/487

Add class-based decorator for instructions to skip PC incrementation

mythril/laser/ethereum/instructions.py

@@ -22,16 +22,37 @@ TT256M1 = 2 ** 256 - 1
 keccak_function_manager = KeccakFunctionManager()
 
 
-def instruction(func):
-    """ Wrapper that handles copy and original return """
+class StateTransition(object):
+    """Decorator that handles global state copy and original return.
 
-    def wrapper(self, global_state):
-        global_state_copy = copy(global_state)
-        new_global_states = func(self, global_state_copy)
-        for state in new_global_states:
-            state.mstate.pc += 1
-        return new_global_states
-    return wrapper
+    This decorator calls the decorated instruction mutator function on a copy of the state that
+    is passed to it. After the call, the resulting new states' program counter is automatically
+    incremented if `increment_pc=True`.
+    """
+
+    def __init__(self, increment_pc=True):
+        self.increment_pc = increment_pc
+
+    @staticmethod
+    def call_on_state_copy(func, func_obj, state):
+        global_state_copy = copy(state)
+        return func(func_obj, global_state_copy)
+
+    def increment_states_pc(self, states):
+        if self.increment_pc:
+            for state in states:
+                state.mstate.pc += 1
+        return states
+
+    def __call__(self, func):
+        def wrapper(func_obj, global_state):
+            new_global_states = self.call_on_state_copy(
+                func,
+                func_obj,
+                global_state
+            )
+            return self.increment_states_pc(new_global_states)
+        return wrapper
 
 
 class Instruction:
@@ -64,11 +85,11 @@ class Instruction:
 
         return instruction_mutator(global_state)
 
-    @instruction
+    @StateTransition()
     def jumpdest_(self, global_state):
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def push_(self, global_state):
         push_instruction = global_state.get_current_instruction()
         push_value = push_instruction['argument'][2:]
@@ -82,7 +103,7 @@ class Instruction:
         global_state.mstate.stack.append(BitVecVal(int(push_value, 16), 256))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def dup_(self, global_state):
         value = int(global_state.get_current_instruction()['opcode'][3:], 10)
         try:
@@ -91,7 +112,7 @@ class Instruction:
             raise VmException('Trying to duplicate out of bounds stack elements')
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def swap_(self, global_state):
         depth = int(self.op_code[4:])
         try:
@@ -101,7 +122,7 @@ class Instruction:
             raise StackUnderflowException
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def pop_(self, global_state):
         try:
             global_state.mstate.stack.pop()
@@ -109,7 +130,7 @@ class Instruction:
             raise StackUnderflowException
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def and_(self, global_state):
         try:
             stack = global_state.mstate.stack
@@ -124,7 +145,7 @@ class Instruction:
             raise StackUnderflowException
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def or_(self, global_state):
         stack = global_state.mstate.stack
         try:
@@ -141,19 +162,19 @@ class Instruction:
             raise StackUnderflowException
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def xor_(self, global_state):
         mstate = global_state.mstate
         mstate.stack.append(mstate.stack.pop() ^ mstate.stack.pop())
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def not_(self, global_state: GlobalState):
         mstate = global_state.mstate
         mstate.stack.append(TT256M1 - mstate.stack.pop())
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def byte_(self, global_state):
         mstate = global_state.mstate
         op0, op1 = mstate.stack.pop(), mstate.stack.pop()
@@ -174,25 +195,25 @@ class Instruction:
         return [global_state]
 
     # Arithmetic
-    @instruction
+    @StateTransition()
     def add_(self, global_state):
         global_state.mstate.stack.append(
             (helper.pop_bitvec(global_state.mstate) + helper.pop_bitvec(global_state.mstate)))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def sub_(self, global_state):
         global_state.mstate.stack.append(
             (helper.pop_bitvec(global_state.mstate) - helper.pop_bitvec(global_state.mstate)))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def mul_(self, global_state):
         global_state.mstate.stack.append(
             (helper.pop_bitvec(global_state.mstate) * helper.pop_bitvec(global_state.mstate)))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def div_(self, global_state):
         op0, op1 = util.pop_bitvec(global_state.mstate), util.pop_bitvec(global_state.mstate)
         if op1 == 0:
@@ -201,7 +222,7 @@ class Instruction:
             global_state.mstate.stack.append(UDiv(op0, op1))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def sdiv_(self, global_state):
         s0, s1 = util.pop_bitvec(global_state.mstate), util.pop_bitvec(global_state.mstate)
         if s1 == 0:
@@ -210,33 +231,33 @@ class Instruction:
             global_state.mstate.stack.append(s0 / s1)
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def mod_(self, global_state):
         s0, s1 = util.pop_bitvec(global_state.mstate), util.pop_bitvec(global_state.mstate)
         global_state.mstate.stack.append(0 if s1 == 0 else URem(s0, s1))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def smod_(self, global_state):
         s0, s1 = util.pop_bitvec(global_state.mstate), util.pop_bitvec(global_state.mstate)
         global_state.mstate.stack.append(0 if s1 == 0 else SRem(s0, s1))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def addmod_(self, global_state):
         s0, s1, s2 = util.pop_bitvec(global_state.mstate), util.pop_bitvec(global_state.mstate), util.pop_bitvec(
             global_state.mstate)
         global_state.mstate.stack.append(URem(URem(s0, s2) + URem(s1, s2), s2))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def mulmod_(self, global_state):
         s0, s1, s2 = util.pop_bitvec(global_state.mstate), util.pop_bitvec(global_state.mstate), util.pop_bitvec(
             global_state.mstate)
         global_state.mstate.stack.append(URem(URem(s0, s2) * URem(s1, s2), s2))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def exp_(self, global_state):
         state = global_state.mstate
 
@@ -248,7 +269,7 @@ class Instruction:
 
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def signextend_(self, global_state):
         state = global_state.mstate
         s0, s1 = state.stack.pop(), state.stack.pop()
@@ -272,28 +293,28 @@ class Instruction:
         return [global_state]
 
     # Comparisons
-    @instruction
+    @StateTransition()
     def lt_(self, global_state):
         state = global_state.mstate
         exp = ULT(util.pop_bitvec(state), util.pop_bitvec(state))
         state.stack.append(exp)
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def gt_(self, global_state):
         state = global_state.mstate
         exp = UGT(util.pop_bitvec(state), util.pop_bitvec(state))
         state.stack.append(exp)
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def slt_(self, global_state):
         state = global_state.mstate
         exp = util.pop_bitvec(state) < util.pop_bitvec(state)
         state.stack.append(exp)
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def sgt_(self, global_state):
         state = global_state.mstate
 
@@ -301,7 +322,7 @@ class Instruction:
         state.stack.append(exp)
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def eq_(self, global_state):
         state = global_state.mstate
 
@@ -319,7 +340,7 @@ class Instruction:
         state.stack.append(exp)
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def iszero_(self, global_state):
         state = global_state.mstate
 
@@ -330,7 +351,7 @@ class Instruction:
         return [global_state]
 
     # Call data
-    @instruction
+    @StateTransition()
     def callvalue_(self, global_state):
         state = global_state.mstate
         environment = global_state.environment
@@ -338,7 +359,7 @@ class Instruction:
 
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def calldataload_(self, global_state):
         state = global_state.mstate
         environment = global_state.environment
@@ -378,7 +399,7 @@ class Instruction:
 
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def calldatasize_(self, global_state):
         state = global_state.mstate
         environment = global_state.environment
@@ -388,7 +409,7 @@ class Instruction:
             state.stack.append(BitVecVal(len(environment.calldata), 256))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def calldatacopy_(self, global_state):
         state = global_state.mstate
         environment = global_state.environment
@@ -453,35 +474,35 @@ class Instruction:
         return [global_state]
 
     # Environment
-    @instruction
+    @StateTransition()
     def address_(self, global_state):
         state = global_state.mstate
         environment = global_state.environment
         state.stack.append(environment.address)
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def balance_(self, global_state):
         state = global_state.mstate
         address = state.stack.pop()
         state.stack.append(global_state.new_bitvec("balance_at_" + str(address), 256))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def origin_(self, global_state):
         state = global_state.mstate
         environment = global_state.environment
         state.stack.append(environment.origin)
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def caller_(self, global_state):
         state = global_state.mstate
         environment = global_state.environment
         state.stack.append(environment.sender)
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def codesize_(self, global_state):
         state = global_state.mstate
         environment = global_state.environment
@@ -489,7 +510,7 @@ class Instruction:
         state.stack.append(len(disassembly.bytecode) // 2)
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def sha3_(self, global_state):
         global keccak_function_manager
 
@@ -526,12 +547,12 @@ class Instruction:
         state.stack.append(BitVecVal(util.concrete_int_from_bytes(keccak, 0), 256))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def gasprice_(self, global_state):
         global_state.mstate.stack.append(global_state.new_bitvec("gasprice", 256))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def codecopy_(self, global_state):
         memory_offset, code_offset, size = global_state.mstate.stack.pop(), global_state.mstate.stack.pop(), global_state.mstate.stack.pop()
 
@@ -580,7 +601,7 @@ class Instruction:
 
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def extcodesize_(self, global_state):
         state = global_state.mstate
         addr = state.stack.pop()
@@ -606,7 +627,7 @@ class Instruction:
 
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def extcodecopy_(self, global_state):
         # FIXME: not implemented
         state = global_state.mstate
@@ -614,45 +635,45 @@ class Instruction:
         start, s2, size = state.stack.pop(), state.stack.pop(), state.stack.pop()
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def returndatasize_(self, global_state):
         global_state.mstate.stack.append(global_state.new_bitvec("returndatasize", 256))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def blockhash_(self, global_state):
         state = global_state.mstate
         blocknumber = state.stack.pop()
         state.stack.append(global_state.new_bitvec("blockhash_block_" + str(blocknumber), 256))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def coinbase_(self, global_state):
         global_state.mstate.stack.append(global_state.new_bitvec("coinbase", 256))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def timestamp_(self, global_state):
         global_state.mstate.stack.append(global_state.new_bitvec("timestamp", 256))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def number_(self, global_state):
         global_state.mstate.stack.append(global_state.new_bitvec("block_number", 256))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def difficulty_(self, global_state):
         global_state.mstate.stack.append(global_state.new_bitvec("block_difficulty", 256))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def gaslimit_(self, global_state):
         global_state.mstate.stack.append(global_state.new_bitvec("block_gaslimit", 256))
         return [global_state]
 
     # Memory operations
-    @instruction
+    @StateTransition()
     def mload_(self, global_state):
         state = global_state.mstate
         op0 = state.stack.pop()
@@ -679,7 +700,7 @@ class Instruction:
         state.stack.append(data)
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def mstore_(self, global_state):
         state = global_state.mstate
         try:
@@ -717,7 +738,7 @@ class Instruction:
 
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def mstore8_(self, global_state):
         state = global_state.mstate
         op0, value = state.stack.pop(), state.stack.pop()
@@ -733,7 +754,7 @@ class Instruction:
         state.memory[offset] = value % 256
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def sload_(self, global_state):
         global keccak_function_manager
         
@@ -786,6 +807,7 @@ class Instruction:
         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:
@@ -794,7 +816,7 @@ class Instruction:
             keccak_argument = keccak_function_manager.get_argument(keccak_key)
             yield keccak_argument != argument
 
-    @instruction
+    @StateTransition()
     def sstore_(self, global_state):
         global keccak_function_manager
         state = global_state.mstate
@@ -852,7 +874,7 @@ class Instruction:
 
         return [global_state]
 
-    @instruction
+    @StateTransition(increment_pc=False)
     def jump_(self, global_state):
         state = global_state.mstate
         disassembly = global_state.environment.code
@@ -881,7 +903,7 @@ class Instruction:
 
         return [new_state]
 
-    @instruction
+    @StateTransition(increment_pc=False)
     def jumpi_(self, global_state):
         state = global_state.mstate
         disassembly = global_state.environment.code
@@ -894,6 +916,7 @@ class Instruction:
             # FIXME: to broad exception handler
         except:
             logging.debug("Skipping JUMPI to invalid destination.")
+            global_state.mstate.pc += 1
             return [global_state]
 
         # False case
@@ -902,6 +925,7 @@ class Instruction:
         if (type(negated) == bool and negated) or (type(negated) == BoolRef and not is_false(negated)):
             new_state = copy(global_state)
             new_state.mstate.depth += 1
+            new_state.mstate.pc += 1
             new_state.mstate.constraints.append(negated)
             states.append(new_state)
         else:
@@ -924,29 +948,27 @@ class Instruction:
                 new_state.mstate.pc = index
                 new_state.mstate.depth += 1
                 new_state.mstate.constraints.append(condi)
-
                 states.append(new_state)
             else:
                 logging.debug("Pruned unreachable states.")
-
         return states
 
-    @instruction
+    @StateTransition()
     def pc_(self, global_state):
         global_state.mstate.stack.append(global_state.mstate.pc - 1)
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def msize_(self, global_state):
         global_state.mstate.stack.append(global_state.new_bitvec("msize", 256))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def gas_(self, global_state):
         global_state.mstate.stack.append(global_state.new_bitvec("gas", 256))
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def log_(self, global_state):
         # TODO: implement me
         state = global_state.mstate
@@ -956,7 +978,7 @@ class Instruction:
         # Not supported
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def create_(self, global_state):
         # TODO: implement me
         state = global_state.mstate
@@ -965,7 +987,7 @@ class Instruction:
         state.stack.append(0)
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def return_(self, global_state):
         state = global_state.mstate
         offset, length = state.stack.pop(), state.stack.pop()
@@ -976,7 +998,7 @@ class Instruction:
             logging.debug("Return with symbolic length or offset. Not supported")
         global_state.current_transaction.end(global_state, return_data)
 
-    @instruction
+    @StateTransition()
     def suicide_(self, global_state):
         target = global_state.mstate.stack.pop()
 
@@ -995,23 +1017,23 @@ class Instruction:
 
         global_state.current_transaction.end(global_state)
 
-    @instruction
+    @StateTransition()
     def revert_(self, global_state):
         return []
 
-    @instruction
+    @StateTransition()
     def assert_fail_(self, global_state):
         return []
 
-    @instruction
+    @StateTransition()
     def invalid_(self, global_state):
         return []
 
-    @instruction
+    @StateTransition()
     def stop_(self, global_state):
         global_state.current_transaction.end(global_state)
 
-    @instruction
+    @StateTransition()
     def call_(self, global_state):
 
         instr = global_state.get_current_instruction()
@@ -1070,7 +1092,7 @@ class Instruction:
                                              call_data_type=call_data_type)
         raise TransactionStartSignal(transaction, self.op_code)
 
-    @instruction
+    @StateTransition()
     def call_post(self, global_state):
         instr = global_state.get_current_instruction()
 
@@ -1111,7 +1133,7 @@ class Instruction:
 
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def callcode_(self, global_state):
         instr = global_state.get_current_instruction()
         environment = global_state.environment
@@ -1138,7 +1160,7 @@ class Instruction:
                                              )
         raise TransactionStartSignal(transaction, self.op_code)
 
-    @instruction
+    @StateTransition()
     def callcode_post(self, global_state):
         instr = global_state.get_current_instruction()
 
@@ -1180,7 +1202,7 @@ class Instruction:
         return [global_state]
 
 
-    @instruction
+    @StateTransition()
     def delegatecall_(self, global_state):
         instr = global_state.get_current_instruction()
         environment = global_state.environment
@@ -1208,7 +1230,7 @@ class Instruction:
         raise TransactionStartSignal(transaction, self.op_code)
 
 
-    @instruction
+    @StateTransition()
     def delegatecall_post(self, global_state):
         instr = global_state.get_current_instruction()
 
@@ -1252,7 +1274,7 @@ class Instruction:
 
         return [global_state]
 
-    @instruction
+    @StateTransition()
     def staticcall_(self, global_state):
         # TODO: implement me
         instr = global_state.get_current_instruction()

mythril/laser/ethereum/svm.py

@@ -218,7 +218,7 @@ class LaserEVM:
                     new_node.flags |= NodeFlags.FUNC_ENTRY
             except IndexError:
                 new_node.flags |= NodeFlags.FUNC_ENTRY
-        address = state.environment.code.instruction_list[state.mstate.pc - 1]['address']
+        address = state.environment.code.instruction_list[state.mstate.pc]['address']
 
         environment = state.environment
         disassembly = environment.code