Refactor machine stack

mythril/laser/ethereum/exceptions.py

@@ -6,5 +6,9 @@ class StackUnderflowException(VmException):
     pass
 
 
+class StackOverflowException(VmException):
+    pass
+
+
 class StopSignal(Exception):
     pass

mythril/laser/ethereum/instructions.py

@@ -14,7 +14,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
+from mythril.laser.ethereum.exceptions import VmException
 
 TT256 = 2 ** 256
 TT256M1 = 2 ** 256 - 1
@@ -92,51 +92,41 @@ class Instruction:
     @instruction
     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]
 
     @instruction
     def pop_(self, global_state):
-        try:
-            global_state.mstate.stack.pop()
-        except IndexError:
-            raise StackUnderflowException
+        global_state.mstate.stack.pop()
         return [global_state]
 
     @instruction
     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]
 
     @instruction
     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]
 
     @instruction

mythril/laser/ethereum/state.py

@@ -3,6 +3,10 @@ from copy import copy, deepcopy
 from enum import Enum
 from random import randint
 
+from mythril.laser.ethereum.exceptions import StackOverflowException, StackUnderflowException
+
+STACK_LIMIT = 1024
+
 
 class CalldataType(Enum):
     CONCRETE = 1
@@ -122,6 +126,29 @@ class Environment:
                     calldata_type=self.calldata_type)
 
 
+class MachineStack(list):
+
+    def __init__(self):
+        super(MachineStack, self).__init__([])
+
+    def append(self, element):
+        if super(MachineStack, self).__len__() >= STACK_LIMIT:
+            raise StackOverflowException
+        super(MachineStack, self).append(element)
+
+    def pop(self):
+        try:
+            return super(MachineStack, self).pop()
+        except IndexError:
+            raise StackUnderflowException
+
+    def __getitem__(self, item):
+        try:
+            return super(MachineStack, self).__getitem__(item)
+        except IndexError:
+            raise StackUnderflowException
+
+
 class MachineState:
     """
     MachineState represents current machine state also referenced to as \mu
@@ -129,7 +156,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 = []