Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions arrow-avro/src/reader/block.rs
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ impl BlockDecoder {
while !buf.is_empty() {
match self.state {
BlockDecoderState::Count => {
if let Some(c) = self.vlq_decoder.long(&mut buf) {
if let Some(c) = self.vlq_decoder.long(&mut buf)? {
self.in_progress.count = c.try_into().map_err(|_| {
AvroError::ParseError(format!(
"Block count cannot be negative, got {c}"
Expand All @@ -91,7 +91,7 @@ impl BlockDecoder {
}
}
BlockDecoderState::Size => {
if let Some(c) = self.vlq_decoder.long(&mut buf) {
if let Some(c) = self.vlq_decoder.long(&mut buf)? {
self.bytes_remaining = c.try_into().map_err(|_| {
AvroError::ParseError(format!("Block size cannot be negative, got {c}"))
})?;
Expand Down
8 changes: 4 additions & 4 deletions arrow-avro/src/reader/header.rs
Original file line number Diff line number Diff line change
Expand Up @@ -253,7 +253,7 @@ impl HeaderDecoder {
}
}
HeaderDecoderState::BlockCount => {
if let Some(block_count) = self.vlq_decoder.long(&mut buf) {
if let Some(block_count) = self.vlq_decoder.long(&mut buf)? {
match block_count.try_into() {
Ok(0) => {
self.state = HeaderDecoderState::Sync;
Expand All @@ -271,7 +271,7 @@ impl HeaderDecoder {
}
}
HeaderDecoderState::BlockLen => {
if self.vlq_decoder.long(&mut buf).is_some() {
if self.vlq_decoder.long(&mut buf)?.is_some() {
self.state = HeaderDecoderState::KeyLen
}
}
Expand Down Expand Up @@ -301,13 +301,13 @@ impl HeaderDecoder {
}
}
HeaderDecoderState::KeyLen => {
if let Some(len) = self.vlq_decoder.long(&mut buf) {
if let Some(len) = self.vlq_decoder.long(&mut buf)? {
self.bytes_remaining = len as _;
self.state = HeaderDecoderState::Key;
}
}
HeaderDecoderState::ValueLen => {
if let Some(len) = self.vlq_decoder.long(&mut buf) {
if let Some(len) = self.vlq_decoder.long(&mut buf)? {
self.bytes_remaining = len as _;
self.state = HeaderDecoderState::Value;
}
Expand Down
62 changes: 59 additions & 3 deletions arrow-avro/src/reader/vlq.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
// specific language governing permissions and limitations
// under the License.

use crate::errors::AvroError;

/// Decoder for zig-zag encoded variable length (VLW) integers
///
/// See also:
Expand All @@ -29,19 +31,29 @@ pub struct VLQDecoder {

impl VLQDecoder {
/// Decode a signed long from `buf`
pub fn long(&mut self, buf: &mut &[u8]) -> Option<i64> {
///
/// Returns `Err` if more than 10 continuation bytes accumulate across calls without a
/// terminator, which would otherwise overflow the accumulated `i64`.
pub fn long(&mut self, buf: &mut &[u8]) -> Result<Option<i64>, AvroError> {
while let Some(byte) = buf.first().copied() {
if self.shift == 63 && byte >= 0x02 {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Adding more branching to the loop makes me nervous...I'd like to see some benchmarks for this.

Perhaps something like

            self.in_progress |=
                ((byte & 0x7F) as u64)
                    .checked_shl(self.shift)
                    .ok_or(AvroError::ParseError(
                        "Malformed Avro varint: too many continuation bytes".to_string(),
                    ))?;

would work?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good news on the benchmark front first: Jefffrey's bot run above shows every avro_reader case at ~1.00x ratio vs. base (both wall-clock and CPU), so the extra branch isn't measurable here.

On the checked_shl suggestion specifically — I don't think it's equivalent, and I'd like to explain why before swapping it in. checked_shl(shift) only returns None when the shift amount itself is >= 64 (i.e. an 11th+ continuation byte). It does not detect the case that actually matters at the boundary: on the 10th byte (shift == 63), a byte like 0x7F (0b111_1111) shifted left by 63 doesn't overflow the shift-amount check at all — it just silently drops the top 6 bits of that byte's payload via ordinary << semantics, and checked_shl returns Some(...) with a truncated, wrong value rather than None. So that version would still decode a malformed 10-byte varint successfully, just to the wrong i64 — which is the same "silently produce something instead of erroring" outcome @alamb's review on the stale #9887 was trying to move away from (discussion).

The byte >= 0x02 check at shift == 63 is what actually catches that: it mirrors the exact bound read_varint_array already enforces for its own 10th byte a few lines below ((b < 0x02).then_some(...)), so this isn't a new convention, just applying the file's existing one to the stateful decoder too.

Given the benchmarks come back flat, I'd lean toward keeping the explicit check for correctness — but happy to revisit if you'd rather accept the narrower checked_shl guarantee (catches unbounded/DoS-style inputs, just not this one specific truncation case) for less branching.

self.in_progress = 0;
self.shift = 0;
return Err(AvroError::ParseError(
"Malformed Avro varint: too many continuation bytes".to_string(),
));
}
*buf = &buf[1..];
self.in_progress |= ((byte & 0x7F) as u64) << self.shift;
self.shift += 7;
if byte & 0x80 == 0 {
let val = self.in_progress;
self.in_progress = 0;
self.shift = 0;
return Some((val >> 1) as i64 ^ -((val & 1) as i64));
return Ok(Some((val >> 1) as i64 ^ -((val & 1) as i64)));
}
}
None
Ok(None)
}
}

Expand Down Expand Up @@ -163,4 +175,48 @@ mod tests {
varint_test(rand::random());
}
}

fn zigzag_encode(n: i64) -> u64 {
((n << 1) ^ (n >> 63)) as u64
}

fn long_test(n: i64) {
let mut buf = [0_u8; 10];
let len = encode_var(zigzag_encode(n), &mut buf);
let mut decoder = VLQDecoder::default();
let mut slice = &buf[..len];
assert_eq!(decoder.long(&mut slice).unwrap(), Some(n));
assert!(slice.is_empty());
}

#[test]
fn test_long_roundtrip() {
long_test(0);
long_test(1);
long_test(-1);
long_test(4395932);
long_test(-4395932);
long_test(i64::MAX);
long_test(i64::MIN);

for _ in 0..1000 {
long_test(rand::random());
}
}

#[test]
fn test_long_overlong_varint_returns_error() {
// The Avro file magic followed by 13 continuation bytes and a terminator, as
// reported against #10290: previously drove `self.shift` past 63 and panicked
// with "attempt to shift left with overflow" inside `<< self.shift`.
let overlong = [0xFFu8; 13];
let mut decoder = VLQDecoder::default();
let mut buf = &overlong[..];
assert!(decoder.long(&mut buf).is_err());

// The decoder's state must be reset after a malformed varint, so a subsequent
// call decodes a fresh varint cleanly rather than staying stuck mid-decode.
let mut fresh = &[0x02u8][..]; // zigzag-encoded 1
assert_eq!(decoder.long(&mut fresh).unwrap(), Some(1));
}
}
Loading