From 2dea33c82ff8ccda50bf978490c1ea96fbf81980 Mon Sep 17 00:00:00 2001 From: Akash Kumar <116457960+akashchamp@users.noreply.github.com> Date: Sat, 12 Sep 2026 07:31:09 +0530 Subject: [PATCH] fix(avro): bound Cython decoder input --- pyiceberg/avro/decoder_basic.c | 43 ++++++++++------- pyiceberg/avro/decoder_fast.pyx | 85 +++++++++++++++++++-------------- tests/avro/test_decoder.py | 79 ++++++++++++++++++++++++++++++ 3 files changed, 156 insertions(+), 51 deletions(-) diff --git a/pyiceberg/avro/decoder_basic.c b/pyiceberg/avro/decoder_basic.c index c3954f330d..6b31ed64ca 100644 --- a/pyiceberg/avro/decoder_basic.c +++ b/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); } diff --git a/pyiceberg/avro/decoder_fast.pyx b/pyiceberg/avro/decoder_fast.pyx index ffd23dd977..a24697a036 100644 --- a/pyiceberg/avro/decoder_fast.pyx +++ b/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 > (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,9 +77,11 @@ 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. @@ -79,6 +89,7 @@ cdef class CythonBinaryDecoder: 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, newarray.data.as_ulonglongs) + self._decode_zigzag_ints(count, 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, &result) - if result[1] <= 0: - dest[result[0]] = b"" + self._decode_zigzag_ints(2, raw_result) + key = raw_result[0] + length = 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(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 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 raw_length <= 0: + return + self._ensure_available(raw_length) + self._current += raw_length def skip_utf8(self) -> None: self.skip_bytes() diff --git a/tests/avro/test_decoder.py b/tests/avro/test_decoder.py index 1cf8346347..9a0961e38d 100644 --- a/tests/avro/test_decoder.py +++ b/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")