diff --git a/lang/py/avro/io.py b/lang/py/avro/io.py index f5063a6686f..6a6e9a26f93 100644 --- a/lang/py/avro/io.py +++ b/lang/py/avro/io.py @@ -224,13 +224,41 @@ def read(self, n: int) -> bytes: raise avro.errors.InvalidAvroBinaryEncoding(f"Requested {n} bytes to read, expected positive integer.") if n > self._MAX_UNCHECKED_READ: remaining = self.bytes_remaining() - if remaining is not None and n > remaining: - raise avro.errors.InvalidAvroBinaryEncoding(f"Requested {n} bytes to read, but only {remaining} remain.") + if remaining is not None: + if n > remaining: + raise avro.errors.InvalidAvroBinaryEncoding(f"Requested {n} bytes to read, but only {remaining} remain.") + else: + # The number of bytes remaining is unknown (a non-seekable stream: + # socket, pipe, decompression stream). A single reader.read(n) for + # a huge declared n allocates n bytes up front before a single + # payload byte is validated, so a tiny truncated/hostile input can + # force a large allocation. Read into a buffer that grows in + # bounded chunks instead, so the cost of a hostile length is + # proportional to the bytes actually delivered and a truncated + # stream fails after a bounded allocation. + return self._read_bounded(n) read_bytes = self.reader.read(n) if len(read_bytes) != n: raise avro.errors.InvalidAvroBinaryEncoding(f"Read {len(read_bytes)} bytes, expected {n} bytes") return read_bytes + def _read_bounded(self, n: int) -> bytes: + """Read exactly ``n`` bytes in bounded chunks from a non-seekable stream. + + Reads at most ``_MAX_UNCHECKED_READ`` bytes per step into a growing buffer + so a truncated or hostile declared length fails after a bounded allocation + rather than allocating the full ``n`` bytes up front. + """ + buf = bytearray() + while len(buf) < n: + chunk = self.reader.read(min(self._MAX_UNCHECKED_READ, n - len(buf))) + if not chunk: + break + buf.extend(chunk) + if len(buf) != n: + raise avro.errors.InvalidAvroBinaryEncoding(f"Read {len(buf)} bytes, expected {n} bytes") + return bytes(buf) + def bytes_remaining(self) -> Optional[int]: """ Return the number of bytes still available to read, or ``None`` when diff --git a/lang/py/avro/test/test_bounded_stream_read.py b/lang/py/avro/test/test_bounded_stream_read.py new file mode 100644 index 00000000000..bc0ff092509 --- /dev/null +++ b/lang/py/avro/test/test_bounded_stream_read.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 + +## +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""AVRO-4303: bound bytes/string allocation from a length prefix on a stream.""" + +import io +import unittest + +import avro.errors +import avro.io + + +def _encode_long(value: int) -> bytes: + """Zig-zag + varint encode a long, matching BinaryEncoder.write_long.""" + datum = (value << 1) ^ (value >> 63) + out = bytearray() + while (datum & ~0x7F) != 0: + out.append((datum & 0x7F) | 0x80) + datum >>= 7 + out.append(datum) + return bytes(out) + + +class NonSeekable: + """A minimal non-seekable, tell-less stream wrapper (socket/pipe-like). + + Records the largest single ``read(n)`` request so tests can assert the + decoder never asks for one huge allocation up front (i.e. that it goes + through the bounded, chunked read path). + """ + + def __init__(self, data: bytes) -> None: + self._bio = io.BytesIO(data) + self.max_read_request = 0 + + def read(self, n: int = -1) -> bytes: + if n is not None and n >= 0: + self.max_read_request = max(self.max_read_request, n) + return self._bio.read(n) + + def seekable(self) -> bool: + return False + + +class TestBoundedStreamRead(unittest.TestCase): + # A near-2GB declared length a single up-front allocation could not satisfy, + # so reaching a bounded decode error proves no full allocation was attempted. + HUGE_LENGTH = (1 << 31) - 1 - 8 + + @staticmethod + def _length_prefixed(declared: int, payload: bytes) -> bytes: + return _encode_long(declared) + payload + + def test_huge_bytes_length_on_stream_rejected_without_huge_allocation(self) -> None: + data = self._length_prefixed(self.HUGE_LENGTH, b"\x01\x02\x03\x04\x05") + stream = NonSeekable(data) + decoder = avro.io.BinaryDecoder(stream) # type: ignore[arg-type] + self.assertRaises(avro.errors.InvalidAvroBinaryEncoding, decoder.read_bytes) + # The decoder must never have requested the full declared length in a + # single read; it reads in bounded chunks instead. + self.assertLessEqual(stream.max_read_request, avro.io.BinaryDecoder._MAX_UNCHECKED_READ) + + def test_huge_string_length_on_stream_rejected_without_huge_allocation(self) -> None: + data = self._length_prefixed(self.HUGE_LENGTH, b"abc") + stream = NonSeekable(data) + decoder = avro.io.BinaryDecoder(stream) # type: ignore[arg-type] + self.assertRaises(avro.errors.InvalidAvroBinaryEncoding, decoder.read_utf8) + self.assertLessEqual(stream.max_read_request, avro.io.BinaryDecoder._MAX_UNCHECKED_READ) + + def test_legitimate_large_bytes_round_trips_on_stream(self) -> None: + # Larger than the per-chunk bound so it exercises the chunked-read path, + # but a genuinely present payload must still decode intact. + payload = bytes((i & 0xFF) for i in range(2 * 1024 * 1024)) + data = self._length_prefixed(len(payload), payload) + stream = NonSeekable(data) + decoder = avro.io.BinaryDecoder(stream) # type: ignore[arg-type] + self.assertEqual(decoder.read_bytes(), payload) + # Even for a legitimately large value the read is chunked, so no single + # read request exceeds the per-chunk bound. + self.assertLessEqual(stream.max_read_request, avro.io.BinaryDecoder._MAX_UNCHECKED_READ) + + +if __name__ == "__main__": + unittest.main()