Handle very short packets

Yay quickcheck.

2 files changed, 19 insertions, 1 deletions

diff --git a/python/subunit/tests/test_test_protocol2.py b/python/subunit/tests/test_test_protocol2.py
index c21392c..9a84a84 100644
--- a/python/subunit/tests/test_test_protocol2.py
+++ b/python/subunit/tests/test_test_protocol2.py
@@ -17,6 +17,11 @@
 from io import BytesIO
 import datetime
 
+from hypothesis import given
+# To debug hypothesis
+# from hypothesis import Settings, Verbosity
+# Settings.default.verbosity = Verbosity.verbose
+import hypothesis.strategies as st
 from testtools import TestCase
 from testtools.matchers import Contains, HasLength
 from testtools.tests.test_testresult import TestStreamResultContract
@@ -434,3 +439,12 @@ class TestByteStreamToStreamResult(TestCase):
             file_bytes=b'foo')
         self.check_event(content.getvalue(), test_id=None, file_name='bar',
             route_code='0', mime_type='text/plain', file_bytes=b'foo')
+
+    @given(st.binary())
+    def test_hypothesis_decoding(self, code_bytes):
+        source = BytesIO(code_bytes)
+        result = StreamResult()
+        stream = subunit.ByteStreamToStreamResult(
+            source, non_subunit_name="stdout")
+        stream.run(result)
+        self.assertEqual(b'', source.read())
diff --git a/python/subunit/v2.py b/python/subunit/v2.py
index 057f65c..f649895 100644
--- a/python/subunit/v2.py
+++ b/python/subunit/v2.py
@@ -386,7 +386,11 @@ class ByteStreamToStreamResult(object):
     def _parse(self, packet, result):
             # 2 bytes flags, at most 3 bytes length.
             packet.append(self.source.read(5))
-            flags = struct.unpack(FMT_16, packet[-1][:2])[0]
+            if len(packet[-1]) != 5:
+                raise ParseError(
+                    'Short read - got %d bytes, wanted 5' % len(packet[-1]))
+            flag_bytes = packet[-1][:2]
+            flags = struct.unpack(FMT_16, flag_bytes)[0]
             length, consumed = self._parse_varint(
                 packet[-1], 2, max_3_bytes=True)
             remainder = self.source.read(length - 6)