diff --git a/pcapkit/protocols/protocol.py b/pcapkit/protocols/protocol.py index e25506d6b7..6d673d42a6 100644 --- a/pcapkit/protocols/protocol.py +++ b/pcapkit/protocols/protocol.py @@ -1767,8 +1767,7 @@ def _read_unpack(self, size: 'int' = 1, *, signed: 'bool' = False, Raises: StructError: If the buffer is empty, whatever ``quiet`` says, or if - unpack (:func:`struct.unpack`) failed, and :exc:`struct.error` - raised, without ``quiet``. + fewer than ``size`` octets were read, without ``quiet``. """ endian = '<' if lilendian else '>' @@ -1788,6 +1787,10 @@ def _read_unpack(self, size: 'int' = 1, *, signed: 'bool' = False, raise StructError('unpack: empty buffer', quiet=True, eof=True) if kind is None: + # A short read fails here as :func:`struct.unpack` makes it fail + # for the sizes that have a format code (#2279). + if len(mem) < size and not quiet: + raise StructError(f'{self.__class__.__name__}: unpack failed') end = 'little' if lilendian else 'big' # type: Literal['little', 'big'] buf = int.from_bytes(mem, end, signed=signed) else: diff --git a/tests/protocols/test_protocol_base_unit.py b/tests/protocols/test_protocol_base_unit.py index 257b9bfb07..ee91e7e2b0 100644 --- a/tests/protocols/test_protocol_base_unit.py +++ b/tests/protocols/test_protocol_base_unit.py @@ -234,6 +234,25 @@ class Number(enum.IntEnum): with self.assertRaises(ProtocolNotImplemented): DummyProtocol._make_index('missing') + def test_read_unpack_short_read_fails_for_every_size(self) -> None: + # #2279: a size with no struct format code read a short buffer as a + # smaller integer instead of failing as sizes 1, 2, 4 and 8 do. + DummyProtocol, _, _ = self._make_protocol_class() + from pcapkit.utilities.exceptions import StructError + + proto = DummyProtocol(packet=b'\x01') + for size in (2, 3, 4, 5, 8): + for signed in (False, True): + with self.subTest(size=size, signed=signed): + proto._file = io.BytesIO(bytes(range(1, size))) + with self.assertRaises(StructError) as context: + proto._read_unpack(size, signed=signed) + self.assertFalse(context.exception.eof) + + proto._file = io.BytesIO(bytes(range(1, size))) + self.assertEqual(proto._read_unpack(size, signed=signed, quiet=True), + int.from_bytes(bytes(range(1, size)), 'big', signed=signed)) + def test_register_analyze_and_next_layer_paths(self) -> None: DummyProtocol, DummyData, _ = self._make_protocol_class() from pcapkit.corekit.module import ModuleDescriptor