AVRO-4330: [python] Bound bytes/string allocation on non-seekable streams

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

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()