Merge pull request #503 from norhh/machinestack_refactor

Machinestack refactor

mythril/laser/ethereum/evm_exceptions.py

@@ -2,14 +2,13 @@ class VmException(Exception):
     pass
 
 
-class StackUnderflowException(VmException):
+class StackUnderflowException(IndexError, VmException):
     pass
 
 
-class InvalidJumpDestination(VmException):
+class StackOverflowException(VmException):
     pass
 
 
-class StopSignal(Exception):
+class InvalidJumpDestination(VmException):
     pass
-

mythril/laser/ethereum/instructions.py

@@ -13,7 +13,7 @@ from mythril.laser.ethereum.state import GlobalState, CalldataType
 import mythril.laser.ethereum.natives as natives
 from mythril.laser.ethereum.transaction import MessageCallTransaction, TransactionStartSignal, \
     ContractCreationTransaction
-from mythril.laser.ethereum.exceptions import VmException, StackUnderflowException, InvalidJumpDestination
+from mythril.laser.ethereum.evm_exceptions import VmException, StackUnderflowException, InvalidJumpDestination
 from mythril.laser.ethereum.keccak import KeccakFunctionManager
 
 TT256 = 2 ** 256
@@ -106,60 +106,47 @@ class Instruction:
     @StateTransition()
     def dup_(self, global_state):
         value = int(global_state.get_current_instruction()['opcode'][3:], 10)
-        try:
-            global_state.mstate.stack.append(global_state.mstate.stack[-value])
-        except IndexError:
-            raise VmException('Trying to duplicate out of bounds stack elements')
+        global_state.mstate.stack.append(global_state.mstate.stack[-value])
         return [global_state]
 
     @StateTransition()
     def swap_(self, global_state):
         depth = int(self.op_code[4:])
-        try:
-            stack = global_state.mstate.stack
-            stack[-depth - 1], stack[-1] = stack[-1], stack[-depth - 1]
-        except IndexError:
-            raise StackUnderflowException
+        stack = global_state.mstate.stack
+        stack[-depth - 1], stack[-1] = stack[-1], stack[-depth - 1]
         return [global_state]
 
     @StateTransition()
     def pop_(self, global_state):
-        try:
-            global_state.mstate.stack.pop()
-        except IndexError:
-            raise StackUnderflowException
+        global_state.mstate.stack.pop()
         return [global_state]
 
     @StateTransition()
     def and_(self, global_state):
-        try:
-            stack = global_state.mstate.stack
-            op1, op2 = stack.pop(), stack.pop()
-            if type(op1) == BoolRef:
-                op1 = If(op1, BitVecVal(1, 256), BitVecVal(0, 256))
-            if type(op2) == BoolRef:
-                op2 = If(op2, BitVecVal(1, 256), BitVecVal(0, 256))
-
-            stack.append(op1 & op2)
-        except IndexError:
-            raise StackUnderflowException
+        stack = global_state.mstate.stack
+        op1, op2 = stack.pop(), stack.pop()
+        if type(op1) == BoolRef:
+            op1 = If(op1, BitVecVal(1, 256), BitVecVal(0, 256))
+        if type(op2) == BoolRef:
+            op2 = If(op2, BitVecVal(1, 256), BitVecVal(0, 256))
+
+        stack.append(op1 & op2)
+
         return [global_state]
 
     @StateTransition()
     def or_(self, global_state):
         stack = global_state.mstate.stack
-        try:
-            op1, op2 = stack.pop(), stack.pop()
+        op1, op2 = stack.pop(), stack.pop()
+
+        if type(op1) == BoolRef:
+            op1 = If(op1, BitVecVal(1, 256), BitVecVal(0, 256))
 
-            if type(op1) == BoolRef:
-                op1 = If(op1, BitVecVal(1, 256), BitVecVal(0, 256))
+        if type(op2) == BoolRef:
+            op2 = If(op2, BitVecVal(1, 256), BitVecVal(0, 256))
 
-            if type(op2) == BoolRef:
-                op2 = If(op2, BitVecVal(1, 256), BitVecVal(0, 256))
+        stack.append(op1 | op2)
 
-            stack.append(op1 | op2)
-        except IndexError:  # Stack underflow
-            raise StackUnderflowException
         return [global_state]
 
     @StateTransition()
@@ -703,10 +690,7 @@ class Instruction:
     @StateTransition()
     def mstore_(self, global_state):
         state = global_state.mstate
-        try:
-            op0, value = state.stack.pop(), state.stack.pop()
-        except IndexError:
-            raise VmException('Stack underflow exception')
+        op0, value = state.stack.pop(), state.stack.pop()
 
         try:
             mstart = util.get_concrete_int(op0)

mythril/laser/ethereum/state.py

@@ -4,6 +4,8 @@ from copy import copy, deepcopy
 from enum import Enum
 from random import randint
 
+from mythril.laser.ethereum.evm_exceptions import StackOverflowException, StackUnderflowException
+
 
 class CalldataType(Enum):
     CONCRETE = 1
@@ -127,6 +129,56 @@ class Environment:
                     calldata_type=self.calldata_type)
 
 
+class MachineStack(list):
+    """
+    Defines EVM stack, overrides the default list to handle overflows
+    """
+    STACK_LIMIT = 1024
+
+    def __init__(self, default_list=[]):
+        super(MachineStack, self).__init__(default_list)
+
+    def append(self, element):
+        """
+        :param element: element to be appended to the list
+        :function: appends the element to list if the size is less than STACK_LIMIT, else throws an error
+        """
+        if super(MachineStack, self).__len__() >= self.STACK_LIMIT:
+            raise StackOverflowException("Reached the EVM stack limit of {}, you can't append more "
+                                         "elements".format(self.STACK_LIMIT))
+        super(MachineStack, self).append(element)
+
+    def pop(self, index=-1):
+        """
+        :param index:index to be popped, same as the list() class.
+        :returns popped value
+        :function: same as list() class but throws StackUnderflowException for popping from an empty list
+        """
+
+        try:
+            return super(MachineStack, self).pop(index)
+        except IndexError:
+            raise StackUnderflowException("Trying to pop from an empty stack")
+
+    def __getitem__(self, item):
+        try:
+            return super(MachineStack, self).__getitem__(item)
+        except IndexError:
+            raise StackUnderflowException("Trying to access a stack element which doesn't exist")
+
+    def __add__(self, other):
+        """
+        Implement list concatenation if needed
+        """
+        raise NotImplementedError('Implement this if needed')
+
+    def __iadd__(self, other):
+        """
+        Implement list concatenation if needed
+        """
+        raise NotImplementedError('Implement this if needed')
+
+
 class MachineState:
     """
     MachineState represents current machine state also referenced to as \mu
@@ -134,7 +186,7 @@ class MachineState:
     def __init__(self, gas):
         """ Constructor for machineState """
         self.pc = 0
-        self.stack = []
+        self.stack = MachineStack()
         self.memory = []
         self.gas = gas
         self.constraints = []
@@ -159,7 +211,7 @@ class MachineState:
     def pop(self, amount=1):
         """ Pops amount elements from the stack"""
         if amount >= len(self.stack):
-            raise IndexError()
+            raise StackUnderflowException
         values = self.stack[-amount:][::-1]
         del self.stack[-amount:]
 

mythril/laser/ethereum/svm.py

@@ -2,6 +2,7 @@ import logging
 from mythril.laser.ethereum.state import WorldState
 from mythril.laser.ethereum.transaction import TransactionStartSignal, TransactionEndSignal, \
     ContractCreationTransaction
+from mythril.laser.ethereum.evm_exceptions import StackUnderflowException
 from mythril.laser.ethereum.instructions import Instruction
 from mythril.laser.ethereum.cfg import NodeFlags, Node, Edge, JumpType
 from mythril.laser.ethereum.strategy.basic import DepthFirstSearchStrategy
@@ -9,7 +10,7 @@ from datetime import datetime, timedelta
 from copy import copy
 from mythril.laser.ethereum.transaction import execute_contract_creation, execute_message_call
 from functools import reduce
-from mythril.laser.ethereum.exceptions import VmException
+from mythril.laser.ethereum.evm_exceptions import VmException
 
 
 class SVMError(Exception):
@@ -221,7 +222,7 @@ class LaserEVM:
                     new_node.flags |= NodeFlags.CALL_RETURN
                 else:
                     new_node.flags |= NodeFlags.FUNC_ENTRY
-            except IndexError:
+            except StackUnderflowException:
                 new_node.flags |= NodeFlags.FUNC_ENTRY
         address = state.environment.code.instruction_list[state.mstate.pc]['address']
 

tests/laser/evm_testsuite/evm_test.py

@@ -1,4 +1,4 @@
-from mythril.laser.ethereum.exceptions import VmException
+from mythril.laser.ethereum.evm_exceptions import VmException
 from mythril.laser.ethereum.svm import LaserEVM
 from mythril.laser.ethereum.state import Account
 from mythril.disassembler.disassembly import Disassembly

tests/laser/state/mstack_test.py

@@ -0,0 +1,56 @@
+import pytest
+
+from mythril.laser.ethereum.state import MachineStack
+from mythril.laser.ethereum.evm_exceptions import *
+from tests import BaseTestCase
+
+
+class MachineStackTest(BaseTestCase):
+
+    @staticmethod
+    def test_mstack_constructor():
+        mstack = MachineStack([1, 2])
+        assert(mstack == [1, 2])
+
+    @staticmethod
+    def test_mstack_append_single_element():
+        mstack = MachineStack()
+
+        mstack.append(0)
+
+        assert(mstack == [0])
+
+    @staticmethod
+    def test_mstack_append_multiple_elements():
+
+        mstack = MachineStack()
+
+        for i in range(mstack.STACK_LIMIT):
+            mstack.append(1)
+
+        with pytest.raises(StackOverflowException):
+            mstack.append(1000)
+
+    @staticmethod
+    def test_mstack_pop():
+        mstack = MachineStack([2])
+
+        assert mstack.pop() == 2
+
+        with pytest.raises(StackUnderflowException):
+            mstack.pop()
+
+    @staticmethod
+    def test_mstack_no_support_add():
+        mstack = MachineStack([0, 1])
+
+        with pytest.raises(NotImplementedError):
+            mstack = mstack + [2]
+
+    @staticmethod
+    def test_mstack_no_support_iadd():
+        mstack = MachineStack()
+
+        with pytest.raises(NotImplementedError):
+            mstack += mstack
+

tests/laser/state/mstate_test.py

@@ -1,5 +1,6 @@
 import pytest
 from mythril.laser.ethereum.state import MachineState
+from mythril.laser.ethereum.evm_exceptions import StackUnderflowException
 
 memory_extension_test_data = [
     (0, 0, 10),
@@ -37,7 +38,7 @@ def test_stack_pop_too_many(initial_size, overflow):
     machine_state.stack = [42] * initial_size
 
     # Act + Assert
-    with pytest.raises(IndexError):
+    with pytest.raises(StackUnderflowException):
         machine_state.pop(initial_size + overflow)