iemejia commented on code in PR #3931: URL: https://github.com/apache/avro/pull/3931#discussion_r3737407510
########## lang/py/avro/test/test_bounded_stream_read.py: ########## @@ -0,0 +1,82 @@ +#!/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).""" + + def __init__(self, data: bytes) -> None: + self._bio = io.BytesIO(data) + + def read(self, n: int = -1) -> bytes: + return self._bio.read(n) + + def seekable(self) -> bool: + return False Review Comment: Good point. Strengthened the tests in the latest commit: the `NonSeekable` wrapper now records the largest single `read(n)` request, and each test asserts `max_read_request <= BinaryDecoder._MAX_UNCHECKED_READ`. This fails if the decoder ever requests one big allocation, so it genuinely exercises the bounded-chunk path (both the truncated huge-length case and the 2 MB legitimate round-trip stay within the per-chunk bound) rather than only asserting that decoding raises. ########## lang/py/avro/io.py: ########## @@ -224,13 +224,43 @@ 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. + """ + chunks: List[bytes] = [] + got = 0 + while got < n: + chunk = self.reader.read(min(self._MAX_UNCHECKED_READ, n - got)) + if not chunk: + break + chunks.append(chunk) + got += len(chunk) + if got != n: + raise avro.errors.InvalidAvroBinaryEncoding(f"Read {got} bytes, expected {n} bytes") + return b"".join(chunks) Review Comment: Fixed: `_read_bounded` now accumulates into a growing `bytearray` (via `extend`) and returns `bytes(buf)`, matching the docstring and dropping the intermediate list of chunks. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
