iemejia commented on code in PR #3861:
URL: https://github.com/apache/avro/pull/3861#discussion_r3565728596


##########
lang/py/avro/io.py:
##########
@@ -214,11 +222,43 @@ def read(self, n: int) -> bytes:
         """
         if n < 0:
             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.")
         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 bytes_remaining(self) -> Optional[int]:
+        """
+        Return the number of bytes still available to read, or ``None`` when
+        that count is not known (a non-seekable reader, or one whose position
+        cannot be obtained). Used to reject a declared length or collection
+        block count that exceeds the data actually available before allocating
+        for it.
+        """
+        reader = self.reader
+        try:
+            pos = reader.tell()
+        except (OSError, ValueError, AttributeError):
+            # Not seekable, or the position could not be determined.
+            return None
+        try:
+            reader.seek(0, os.SEEK_END)
+            end = reader.tell()
+            return end - pos
+        except (OSError, ValueError, AttributeError):
+            return None
+        finally:
+            # Always restore the original position, even if seeking to the end
+            # or reading it failed, so the reader is never left at EOF.
+            try:
+                reader.seek(pos)
+            except (OSError, ValueError):

Review Comment:
   Fixed in 102061267c: the finally block now also catches AttributeError when 
restoring the position, so a reader that implements tell() but not seek() 
reliably falls back to None instead of raising.



##########
lang/py/avro/test/test_io.py:
##########
@@ -550,6 +551,136 @@ def test_deserialization_fails(self) -> None:
             )
 
 
+class TestBinaryDecoderAvailableBytes(unittest.TestCase):
+    """A bytes/string value declares a length prefix; a malicious or truncated
+    input can declare far more bytes than actually exist. On a seekable reader
+    that is rejected before allocating for it."""
+
+    @staticmethod
+    def _encode_length_prefix(length: int) -> bytes:
+        buf = io.BytesIO()
+        avro.io.BinaryEncoder(buf).write_long(length)
+        return buf.getvalue()
+
+    def test_read_rejects_length_beyond_stream(self) -> None:
+        # Declares 100 MiB but provides no data.
+        prefix = self._encode_length_prefix(100 * 1024 * 1024)
+        with io.BytesIO(prefix) as bio:
+            decoder = avro.io.BinaryDecoder(bio)
+            self.assertRaises(avro.errors.InvalidAvroBinaryEncoding, 
decoder.read_bytes)
+
+    def test_read_within_stream_still_reads(self) -> None:
+        # A well-formed large value whose data is actually present still reads.
+        payload = b"x" * (2 * 1024 * 1024)
+        buf = io.BytesIO()
+        avro.io.BinaryEncoder(buf).write_bytes(payload)
+        with io.BytesIO(buf.getvalue()) as bio:
+            decoder = avro.io.BinaryDecoder(bio)
+            self.assertEqual(decoder.read_bytes(), payload)
+
+    def test_bytes_remaining_restores_position(self) -> None:
+        # bytes_remaining() must leave the reader position unchanged.
+        with io.BytesIO(b"abcdefghij") as bio:
+            bio.seek(3)
+            decoder = avro.io.BinaryDecoder(bio)
+            self.assertEqual(decoder.bytes_remaining(), 7)
+            self.assertEqual(bio.tell(), 3)
+
+    def test_bytes_remaining_restores_position_on_error(self) -> None:
+        # If reading the end offset fails after seeking, the original position
+        # must still be restored (via the finally block).
+        class FailingEndStream(io.BytesIO):
+            def __init__(self, data: bytes) -> None:
+                super().__init__(data)
+                self._calls = 0
+

Review Comment:
   Removed the unused self._calls assignment from the FailingEndStream test 
helper in 102061267c.



-- 
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]

Reply via email to