diff --git a/pcapkit/protocols/protocol.py b/pcapkit/protocols/protocol.py index e25506d6b7..40d5885130 100644 --- a/pcapkit/protocols/protocol.py +++ b/pcapkit/protocols/protocol.py @@ -1811,12 +1811,19 @@ def _read_binary(self, size: 'int' = 1) -> 'str': Returns: Binary bits (``0``/``1``). + Raises: + StructError: If the buffer is empty, with ``eof`` set as by + :meth:`_read_unpack`, or if fewer than ``size`` octets were read. + """ - bin_ = [] # type: list[str] - for _ in range(size): - byte = self._file.read(1) - bin_.append(bin(ord(byte))[2:].zfill(8)) - return ''.join(bin_) + # ``max`` makes a non-positive size read nothing rather than the + # rest of the stream. + mem = self._file.read(max(size, 0)) + if len(mem) < size: + if not mem: + raise StructError('unpack: empty buffer', quiet=True, eof=True) + raise StructError(f'{self.__class__.__name__}: unpack failed') + return ''.join(bin(byte)[2:].zfill(8) for byte in mem) @overload # pragma: no cover def _read_packet(self, length: 'Optional[int]' = ..., *, header: 'None' = ...) -> 'bytes': ... diff --git a/tests/protocols/test_protocol_base_unit.py b/tests/protocols/test_protocol_base_unit.py index 257b9bfb07..e8d93e5dc9 100644 --- a/tests/protocols/test_protocol_base_unit.py +++ b/tests/protocols/test_protocol_base_unit.py @@ -357,6 +357,28 @@ def __index__(cls) -> int: no_stop._exproto = 'raw' self.assertFalse(no_stop._check_term_threshold()) + def test_read_binary_raises_struct_error_at_end_of_stream(self) -> None: + # #2292: _read_binary called ord() on an empty read and raised + # TypeError, where _read_unpack raises StructError. + DummyProtocol, _, _ = self._make_protocol_class() + from pcapkit.utilities.exceptions import StructError + + proto = DummyProtocol(packet=b'\x01') + proto._file = io.BytesIO(b'') + with self.assertRaises(StructError) as context: + proto._read_binary() + self.assertTrue(context.exception.eof) + + proto._file = io.BytesIO(b'\x01') + with self.assertRaises(StructError) as context: + proto._read_binary(2) + self.assertFalse(context.exception.eof) + + proto._file = io.BytesIO(b'\x01\x80\xff') + self.assertEqual(proto._read_binary(0), '') + self.assertEqual(proto._read_binary(2), '0000000110000000') + self.assertEqual(proto._file.read(), b'\xff') + def test_register_stays_quiet_on_same_object_reregistration(self) -> None: """GitHub issue #718: replaying a call with the identical class is a no-op.