Update is_valid_opcode security

evm/vm/code_stream.py

@@ -10,13 +10,15 @@
 
 class CodeStream(object):
     stream = None
+    deepest = None
 
     logger = logging.getLogger('evm.vm.CodeStream')
 
     def __init__(self, code_bytes):
         validate_is_bytes(code_bytes)
         self.stream = io.BytesIO(code_bytes)
-        self._validity_cache = {}
+        self._validity_cache = set(range(0))
+        self.deepest = 0
 
     def read(self, size):
         return self.stream.read(size)
@@ -66,28 +68,35 @@ def seek(self, pc):
     _validity_cache = None
 
     def is_valid_opcode(self, position):
+        # if position longer than bytecode return false
         if position >= len(self):
             return False
 
-        if position not in self._validity_cache:
-            with self.seek(max(0, position - 32)):
-                prefix = self.read(min(position, 32))
-
-            for offset, opcode in enumerate(reversed(prefix)):
-                if opcode < opcode_values.PUSH1 or opcode > opcode_values.PUSH32:
-                    continue
-
-                push_size = 1 + opcode - opcode_values.PUSH1
-                if push_size <= offset:
-                    continue
-
-                opcode_position = position - 1 - offset
-                if not self.is_valid_opcode(opcode_position):
-                    continue
-
-                self._validity_cache[position] = False
-                break
+        # return false if position already in val_cache
+        if position in self._validity_cache:
+            return False
+        # return true if position not in val_cache but has been parsed
+        if position <= self.deepest:
+            return True
+        else:
+            # set counter to deepest parsed position
+            i = self.deepest
+            while i <= position:
+                # get opcode
+                with self.seek(i):
+                    opcode = self.next()
+                # if opcode = pushxx
+                if opcode >= opcode_values.PUSH1 and opcode <= opcode_values.PUSH32:
+                    # add range(xx) to val_cache
+                    self._validity_cache.update(range((i + 1), ((i + 1) + (opcode - 95))))
+                    # increment counter to end of invalid bytes
+                    i += (1 + (opcode - 95))
+                else:
+                    # if opcode != pushxx : update deepest processed and increment counter
+                    self.deepest = i
+                    i += 1
+
+            if position in self._validity_cache:
+                return False
             else:
-                self._validity_cache[position] = True
-
-        return self._validity_cache[position]
+                return True

tests/core/code-stream/test_code_stream.py

@@ -56,7 +56,65 @@ def test_seek_reverts_to_original_stream_position_when_context_exits():
 
 
 def test_is_valid_opcode_invalidates_bytes_after_PUSHXX_opcodes():
-    code_stream = CodeStream(b'\x01\x60\x02')
+    code_stream = CodeStream(b'\x02\x60\x02\x04')
+    assert code_stream.is_valid_opcode(0) is True  # x02
+    assert code_stream.is_valid_opcode(1) is True  # x60
+    assert code_stream.is_valid_opcode(2) is False  # x02
+    assert code_stream.is_valid_opcode(3) is True  # x04
+    assert code_stream.is_valid_opcode(4) is False  # too long
+
+
+def test_harder_is_valid_opcode():
+    code_stream = CodeStream(b'\x02\x03\x72' + (b'\x04' * 32) + b'\x05')
+    # valid: 0-2
+    # invalid: 3-21 #PUSH19
+    # valid: 22-35
+    # invalid: 36+ # too long
+    assert code_stream.is_valid_opcode(0) is True
+    assert code_stream.is_valid_opcode(1) is True
+    assert code_stream.is_valid_opcode(2) is True
+    assert code_stream.is_valid_opcode(3) is False
+    assert code_stream.is_valid_opcode(21) is False
+    assert code_stream.is_valid_opcode(22) is True
+    assert code_stream.is_valid_opcode(35) is True
+    assert code_stream.is_valid_opcode(36) is False
+
+
+def test_even_harder_is_valid_opcode():
+    test = b'\x02\x03\x7d' + (b'\x04' * 32) + b'\x05\x7e' + (b'\x04' * 35) + b'\x01\x61\x01\x01\x01'
+    code_stream = CodeStream(test)
+    # valid: 0-2
+    # invalid: 3 - 32 # PUSH30
+    # valid: 33 - 36
+    # invalid: 37 - 67 # PUSH31
+    # valid: 68 - 73
+    # invalid: 74, 75 # PUSH2
+    # valid: 76
+    # invalid: 77+ # too long
+    assert code_stream.is_valid_opcode(0) is True
+    assert code_stream.is_valid_opcode(1) is True
+    assert code_stream.is_valid_opcode(2) is True
+    assert code_stream.is_valid_opcode(3) is False
+    assert code_stream.is_valid_opcode(32) is False
+    assert code_stream.is_valid_opcode(33) is True
+    assert code_stream.is_valid_opcode(36) is True
+    assert code_stream.is_valid_opcode(37) is False
+    assert code_stream.is_valid_opcode(67) is False
+    assert code_stream.is_valid_opcode(68) is True
+    assert code_stream.is_valid_opcode(71) is True
+    assert code_stream.is_valid_opcode(72) is True
+    assert code_stream.is_valid_opcode(73) is True
+    assert code_stream.is_valid_opcode(74) is False
+    assert code_stream.is_valid_opcode(75) is False
+    assert code_stream.is_valid_opcode(76) is True
+    assert code_stream.is_valid_opcode(77) is False
+
+
+def test_right_number_of_bytes_invalidated_after_pushxx():
+    code_stream = CodeStream(b'\x02\x03\x60\x02\x02')
     assert code_stream.is_valid_opcode(0) is True
     assert code_stream.is_valid_opcode(1) is True
-    assert code_stream.is_valid_opcode(2) is False
+    assert code_stream.is_valid_opcode(2) is True
+    assert code_stream.is_valid_opcode(3) is False
+    assert code_stream.is_valid_opcode(4) is True
+    assert code_stream.is_valid_opcode(5) is False