fix(avro): bound Cython decoder input

Author: akashchamp

pyiceberg/avro/decoder_basic.c

@@ -23,43 +23,54 @@
   Decode an an array of zig-zag encoded integers from a buffer.
 
   The buffer is advanced to the end of the integers.
+  `end` is the first byte after the buffer.
   `count` is the number of integers to decode.
   `result` is where the decoded integers are stored.
 
   The result is guaranteed to be 64 bits wide.
 
 */
-static inline void decode_zigzag_ints(const unsigned char **buffer, const uint64_t count, uint64_t *result) {
+static inline int decode_zigzag_ints(
+    const unsigned char **buffer, const unsigned char *end, const uint64_t count, uint64_t *result) {
   uint64_t current_index;
   const unsigned char *current_position = *buffer;
   uint64_t temp;
-  // The largest shift will always be < 64
   unsigned char shift;
+  unsigned char byte;
 
   for (current_index = 0; current_index < count; current_index++) {
-    shift = 7;
-    temp = *current_position & 0x7F;
-    while(*current_position & 0x80) {
-        current_position += 1;
-        temp |= (uint64_t)(*current_position & 0x7F) << shift;
-        shift += 7;
+    temp = 0;
+    shift = 0;
+    while (1) {
+      if (current_position >= end || shift >= 64) {
+        return 0;
+      }
+
+      byte = *current_position;
+      current_position += 1;
+
+      if (shift == 63 && (byte & 0x7E)) {
+        return 0;
+      }
+      temp |= (uint64_t)(byte & 0x7F) << shift;
+
+      if (!(byte & 0x80)) {
+        break;
+      }
+      shift += 7;
     }
     result[current_index] = (temp >> 1) ^ (~(temp & 1) + 1);
-    current_position += 1;
   }
   *buffer = current_position;
+  return 1;
 }
 
-
-
 /*
   Skip a zig-zag encoded integer in a buffer.
 
   The buffer is advanced to the end of the integer.
 */
-static inline void skip_zigzag_int(const unsigned char **buffer) {
-  while(**buffer & 0x80) {
-    *buffer += 1;
-  }
-  *buffer += 1;
+static inline int skip_zigzag_int(const unsigned char **buffer, const unsigned char *end) {
+  uint64_t ignored;
+  return decode_zigzag_ints(buffer, end, 1, &ignored);
 }

pyiceberg/avro/decoder_fast.pyx

@@ -25,8 +25,8 @@ import array
 
 
 cdef extern from "decoder_basic.c":
-  void decode_zigzag_ints(const unsigned char **buffer, const uint64_t count, uint64_t *result);
-  void skip_zigzag_int(const unsigned char **buffer);
+  int decode_zigzag_ints(const unsigned char **buffer, const unsigned char *end, const uint64_t count, uint64_t *result);
+  int skip_zigzag_int(const unsigned char **buffer, const unsigned char *end);
 
 unsigned_long_long_array_template = cython.declare(array.array, array.array('Q', []))
 
@@ -61,6 +61,14 @@ cdef class CythonBinaryDecoder:
     def __dealloc__(self):
         PyMem_Free(self._data)
 
+    cdef inline void _ensure_available(self, uint64_t length):
+        if length > <uint64_t>(self._end - self._current):
+            raise EOFError(f"EOF: read {length} bytes")
+
+    cdef inline void _decode_zigzag_ints(self, uint64_t count, uint64_t *result):
+        if not decode_zigzag_ints(&self._current, self._end, count, result):
+            raise EOFError("EOF: read 1 bytes")
+
     cpdef unsigned int tell(self):
         """Return the current stream position."""
         return self._current - self._data
@@ -69,16 +77,19 @@ cdef class CythonBinaryDecoder:
         """Read n bytes."""
         if n < 0:
             raise ValueError(f"Requested {n} bytes to read, expected positive integer.")
+        cdef uint64_t length = n
+        self._ensure_available(length)
         cdef const unsigned char *r = self._current
-        self._current += n
-        return r[0:n]
+        self._current += length
+        return r[0:length]
 
     def read_boolean(self) -> bool:
         """Reads a value from the stream as a boolean.
 
         A boolean is written as a single byte
         whose value is either 0 (false) or 1 (true).
         """
+        self._ensure_available(1)
         self._current += 1;
         return self._current[-1] != 0
 
@@ -88,46 +99,43 @@ cdef class CythonBinaryDecoder:
         int/long values are written using variable-length, zigzag coding.
         """
         cdef uint64_t result;
-        if self._current >= self._end:
-          raise EOFError(f"EOF: read 1 bytes")
-        decode_zigzag_ints(&self._current, 1, &result)
+        self._decode_zigzag_ints(1, &result)
         return result
 
     def read_ints(self, count: int) -> array.array[int]:
         """Reads a list of integers."""
         newarray = array.clone(unsigned_long_long_array_template, count, zero=False)
-        if self._current >= self._end:
-          raise EOFError(f"EOF: read 1 bytes")
-        decode_zigzag_ints(&self._current, count, <uint64_t *>newarray.data.as_ulonglongs)
+        self._decode_zigzag_ints(count, <uint64_t *>newarray.data.as_ulonglongs)
         return newarray
 
     cpdef void read_int_bytes_dict(self, count: int, dest: Dict[int, bytes]):
         """Reads a dictionary of integers for keys and bytes for values into a destination dict."""
-        cdef uint64_t result[2];
-        if self._current >= self._end:
-          raise EOFError(f"EOF: read 1 bytes")
+        cdef uint64_t raw_result[2];
+        cdef int64_t key
+        cdef int64_t length
 
         for _ in range(count):
-          decode_zigzag_ints(&self._current, 2, <uint64_t *>&result)
-          if result[1] <= 0:
-              dest[result[0]] = b""
+          self._decode_zigzag_ints(2, raw_result)
+          key = <int64_t>raw_result[0]
+          length = <int64_t>raw_result[1]
+          if length <= 0:
+              dest[key] = b""
           else:
-              dest[result[0]] = self._current[0:result[1]]
-              self._current += result[1]
+              self._ensure_available(<uint64_t>length)
+              dest[key] = self._current[0:length]
+              self._current += length
 
     cpdef inline bytes read_bytes(self):
         """Bytes are encoded as a long followed by that many bytes of data."""
-        cdef uint64_t length;
-        if self._current >= self._end:
-          raise EOFError(f"EOF: read 1 bytes")
-
-        decode_zigzag_ints(&self._current, 1, &length)
+        cdef uint64_t raw_length;
+        self._decode_zigzag_ints(1, &raw_length)
 
-        if length <= 0:
+        if <int64_t>raw_length <= 0:
             return b""
+        self._ensure_available(raw_length)
         cdef const unsigned char *r = self._current
-        self._current += length
-        return r[0:length]
+        self._current += raw_length
+        return r[0:raw_length]
 
     cpdef float read_float(self):
         """Reads a value from the stream as a float.
@@ -156,25 +164,32 @@ cdef class CythonBinaryDecoder:
         return self.read_bytes().decode("utf-8")
 
     def skip_int(self) -> None:
-        skip_zigzag_int(&self._current)
-        return
+        if not skip_zigzag_int(&self._current, self._end):
+            raise EOFError("EOF: read 1 bytes")
 
     def skip(self, n: int) -> None:
-        self._current += n
+        if n < 0:
+            raise ValueError(f"Requested {n} bytes to skip, expected positive integer.")
+        cdef uint64_t length = n
+        self._ensure_available(length)
+        self._current += length
 
     def skip_boolean(self) -> None:
-        self._current += 1
+        self.skip(1)
 
     def skip_float(self) -> None:
-        self._current += 4
+        self.skip(4)
 
     def skip_double(self) -> None:
-        self._current += 8
+        self.skip(8)
 
     def skip_bytes(self) -> None:
-        cdef uint64_t result;
-        decode_zigzag_ints(&self._current, 1, &result)
-        self._current += result
+        cdef uint64_t raw_length;
+        self._decode_zigzag_ints(1, &raw_length)
+        if <int64_t>raw_length <= 0:
+            return
+        self._ensure_available(raw_length)
+        self._current += raw_length
 
     def skip_utf8(self) -> None:
         self.skip_bytes()

tests/avro/test_decoder.py

@@ -66,6 +66,42 @@ def test_read_int_longer(decoder_class: Callable[[bytes], BinaryDecoder]) -> Non
     assert decoder.read_int() == 1111111
 
 
+def test_cython_decoder_rejects_truncated_varint() -> None:
+    decoder = CythonBinaryDecoder(b"\x80")
+
+    with pytest.raises(EOFError, match="EOF: read 1 bytes"):
+        decoder.read_int()
+
+    assert decoder.tell() == 0
+
+
+def test_cython_decoder_rejects_overlong_varint() -> None:
+    decoder = CythonBinaryDecoder(b"\x80" * 10 + b"\x00")
+
+    with pytest.raises(EOFError, match="EOF: read 1 bytes"):
+        decoder.read_int()
+
+    assert decoder.tell() == 0
+
+
+def test_cython_decoder_rejects_truncated_skipped_varint() -> None:
+    decoder = CythonBinaryDecoder(b"\x80")
+
+    with pytest.raises(EOFError, match="EOF: read 1 bytes"):
+        decoder.skip_int()
+
+    assert decoder.tell() == 0
+
+
+def test_cython_decoder_rejects_truncated_ints() -> None:
+    decoder = CythonBinaryDecoder(b"\x00")
+
+    with pytest.raises(EOFError, match="EOF: read 1 bytes"):
+        decoder.read_ints(2)
+
+    assert decoder.tell() == 0
+
+
 def zigzag_encode(datum: int) -> bytes:
     result = []
     datum = (datum << 1) ^ (datum >> 63)
@@ -192,6 +228,49 @@ def test_read_bytes(decoder_class: Callable[[bytes], BinaryDecoder]) -> None:
     assert actual == b"\x01\x02\x03\x04"
 
 
+def test_cython_decoder_rejects_truncated_bytes() -> None:
+    decoder = CythonBinaryDecoder(b"\x04\x01")
+
+    with pytest.raises(EOFError, match="EOF: read 2 bytes"):
+        decoder.read_bytes()
+
+    assert decoder.tell() == 1
+
+
+def test_cython_decoder_rejects_truncated_skipped_bytes() -> None:
+    decoder = CythonBinaryDecoder(b"\x04\x01")
+
+    with pytest.raises(EOFError, match="EOF: read 2 bytes"):
+        decoder.skip_bytes()
+
+    assert decoder.tell() == 1
+
+
+@pytest.mark.parametrize("decoder_class", AVAILABLE_DECODERS)
+def test_read_negative_length_bytes(decoder_class: Callable[[bytes], BinaryDecoder]) -> None:
+    decoder = decoder_class(b"\x01")
+    assert decoder.read_bytes() == b""
+
+
+@pytest.mark.parametrize("decoder_class", AVAILABLE_DECODERS)
+def test_read_int_bytes_dict_negative_length(decoder_class: Callable[[bytes], BinaryDecoder]) -> None:
+    decoder = decoder_class(b"\x00\x01")
+    dest: dict[int, bytes] = {}
+
+    decoder.read_int_bytes_dict(1, dest)
+
+    assert dest == {0: b""}
+
+
+@pytest.mark.parametrize("decoder_class", AVAILABLE_DECODERS)
+def test_read_int_bytes_dict_rejects_truncated_bytes(decoder_class: Callable[[bytes], BinaryDecoder]) -> None:
+    decoder = decoder_class(b"\x00\x04\x01")
+    dest: dict[int, bytes] = {}
+
+    with pytest.raises(EOFError):
+        decoder.read_int_bytes_dict(1, dest)
+
+
 @pytest.mark.parametrize("decoder_class", AVAILABLE_DECODERS)
 def test_read_utf8(decoder_class: Callable[[bytes], BinaryDecoder]) -> None:
     decoder = decoder_class(b"\x04\x76\x6f")