diff --git a/msgpack/_unpacker.pyx b/msgpack/_unpacker.pyx index f4b89d43..0801a4f9 100644 --- a/msgpack/_unpacker.pyx +++ b/msgpack/_unpacker.pyx @@ -535,6 +535,8 @@ cdef class Unpacker: raise RuntimeError("Unpacker.read_bytes() cannot be called while unpacking is in progress") if self.ctx.execute != NULL: raise ValueError("Cannot switch unpacking methods while an object is incomplete") + if nbytes < 0: + raise ValueError("nbytes must be non-negative") nread = min(self.buf_tail - self.buf_head, nbytes) ret = PyBytes_FromStringAndSize(self.buf + self.buf_head, nread) self.buf_head += nread diff --git a/msgpack/fallback.py b/msgpack/fallback.py index 9151a6bd..b26258b6 100644 --- a/msgpack/fallback.py +++ b/msgpack/fallback.py @@ -354,6 +354,8 @@ def _get_extradata(self): return self._buffer[self._buff_i :] def read_bytes(self, n): + if n < 0: + raise ValueError("nbytes must be non-negative") ret = self._read(n, raise_outofdata=False) self._consume() # The C extension returns bytes. A bytearray slice cannot be a dict key. diff --git a/test/test_unpack.py b/test/test_unpack.py index 5811c27f..03e49ca3 100644 --- a/test/test_unpack.py +++ b/test/test_unpack.py @@ -99,6 +99,24 @@ def test_unpacker_tell_read_bytes(): assert unpacker.read_bytes(n) == raw +@mark.parametrize("from_file", [False, True]) +@mark.parametrize("nbytes", [-1, -(2**31)]) +def test_unpacker_read_bytes_negative(from_file, nbytes): + packed = packb([1, 2]) + unpacker = Unpacker(BytesIO(packed)) if from_file else Unpacker() + if not from_file: + unpacker.feed(packed) + assert unpacker.read_array_header() == 2 + + with raises(ValueError, match="non-negative"): + unpacker.read_bytes(nbytes) + + assert unpacker.tell() == 1 + assert unpacker.read_bytes(0) == b"" + assert unpacker.read_bytes(1) == b"\x01" + assert unpacker.unpack() == 2 + + @mark.skipif( Unpacker.__module__ == "msgpack.fallback", reason="specific to C extension reinit leak",