This is an automated email from the ASF dual-hosted git repository.
Jefffrey pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new d8c1bf4d22 fix: bound IPC message allocations by the bytes actually
read (#10522)
d8c1bf4d22 is described below
commit d8c1bf4d226b02e9dea06a9890e775d50ad5b627
Author: ranflarion <[email protected]>
AuthorDate: Fri Aug 7 09:59:24 2026 -0400
fix: bound IPC message allocations by the bytes actually read (#10522)
# Which issue does this PR close?
- Closes #10521.
# Rationale for this change
`MessageReader::maybe_next` reserves both the metadata length and the
message body length before reading any of the bytes they describe. Both
come out of the stream itself, so a corrupted or truncated stream is
handed straight to the allocator:
`MutableBuffer::from_len_zeroed(message.bodyLength() as usize)` on an
implausible length either aborts the process (`memory allocation of N
bytes failed`, which is not catchable and takes the host process with
it) or panics on `LayoutError`. A negative `bodyLength` is accepted too,
since `as usize` wraps it to a large positive length.
I hit this fuzzing real IPC blocks rather than crafted ones: single-bit
flips over the framing region of genuine streams produced `memory
allocation of 1125899907497992 bytes failed` and `SIGABRT`. Where those
blocks cross disk or a network, one flipped bit ends the process instead
of failing a read the caller could retry.
# What changes are included in this PR?
`bodyLength` now goes through `usize::try_from`, so a negative length is
a parse error rather than a huge positive one.
Neither length reserves more than `MAX_PREALLOC_BYTES` (64 MiB) before
the bytes behind it have arrived. Bodies up to that size are allocated
in one go exactly as before; larger ones grow as the data arrives, which
costs the reallocations that `MutableBuffer::reserve` doubling implies.
That constant is the one judgement call here, trading the size of the
bounded allocation a malformed stream can still ask for against how
large a body keeps the single-allocation path, so it is worth a second
opinion.
The metadata read switches from `resize(meta_len, 0)` plus `read_exact`
to `take(meta_len).read_to_end(&mut self.buf)`. That reuses the retained
capacity across messages and drops the zeroing entirely, so it should be
slightly cheaper than what it replaces rather than a cost, and `Take`
returns `Ok(0)` at its limit so there is no extra read.
Only the streaming path is touched. `read_block` on the file side has
the same shape at `arrow-ipc/src/reader.rs:875` and two `unwrap()`s on
block metadata besides; I left it alone to keep this reviewable, and
noted it in the issue.
This overlaps #9777, which is after the same zeroing for performance
reasons. The two want the same thing here, and I am happy to rebase onto
whatever lands first.
# Are these changes tested?
Yes, two tests in `arrow-ipc/src/reader.rs`.
`test_stream_reader_rejects_implausible_body_length` covers `i64::MAX`,
`1 << 50` and `-1`;
`test_stream_reader_rejects_unbacked_metadata_length` covers a metadata
length of `i32::MAX` with nineteen bytes behind it. Both fail without
the change: the first panics inside `MutableBuffer::from_len_zeroed`,
and the second spends 7.7s zeroing 2 GiB before reporting the wrong
error.
The existing `arrow-ipc` suite passes (139 tests), along with `cargo fmt
--all --check` and `cargo clippy -p arrow-ipc --all-targets
--all-features -- -D warnings`.
# Are there any user-facing changes?
No API changes. A stream that previously aborted or panicked now returns
an `ArrowError`. No breaking changes.
---------
Co-authored-by: Jeffrey Vo <[email protected]>
---
arrow-ipc/src/reader.rs | 130 ++++++++++++++++++++++++++++++++++++++++++++++--
1 file changed, 126 insertions(+), 4 deletions(-)
diff --git a/arrow-ipc/src/reader.rs b/arrow-ipc/src/reader.rs
index 39cdb731bb..d6aa2d6653 100644
--- a/arrow-ipc/src/reader.rs
+++ b/arrow-ipc/src/reader.rs
@@ -1796,6 +1796,38 @@ pub(crate) enum IpcMessage {
},
}
+/// Upper bound on how much memory a single message length is allowed to
reserve before any
+/// of the bytes it promises have been read.
+///
+/// Message lengths come from the stream itself, so they cannot be trusted: a
truncated or
+/// corrupted stream can declare a body of arbitrary size. Reserving that up
front turns a
+/// malformed input into an allocation failure, which aborts the process
rather than
+/// returning an error the caller can handle. Beyond this size the buffer
grows as the data
+/// arrives instead, so an implausible length costs one bounded allocation and
then fails as
+/// a short read.
+///
+/// The value trades the size of that bounded allocation against how large a
body still gets
+/// read in a single allocation: bodies up to this size behave exactly as
before, larger ones
+/// start here and the buffer doubles as the data arrives, so a body of `n`
bytes costs
+/// `log2(n / MAX_PREALLOC_BYTES)` reallocations and the buffer never runs
more than a factor
+/// of two ahead of the bytes actually received.
+const MAX_PREALLOC_BYTES: usize = 64 * 1024 * 1024;
+
+/// Reads exactly `len` bytes of message body, without reserving `len` before
reading it.
+fn read_body_bounded<R: Read>(reader: &mut R, len: usize) ->
Result<MutableBuffer, ArrowError> {
+ let mut buf = MutableBuffer::from_len_zeroed(len.min(MAX_PREALLOC_BYTES));
+ let mut filled = 0;
+ while filled < len {
+ let target = buf.len();
+ reader.read_exact(&mut buf.as_slice_mut()[filled..target])?;
+ filled = target;
+ if filled < len {
+ buf.resize(len.min(target.saturating_mul(2)), 0);
+ }
+ }
+ Ok(buf)
+}
+
/// A low-level construct that reads [`Message::Message`]s from a reader while
/// re-using a buffer for metadata. This is composed into [`StreamReader`].
struct MessageReader<R> {
@@ -1827,15 +1859,29 @@ impl<R: Read> MessageReader<R> {
return Ok(None);
};
- self.buf.resize(meta_len, 0);
- self.reader.read_exact(&mut self.buf)?;
+ // `read_to_end` on a `Take` grows with the bytes that arrive, so an
implausible
+ // `meta_len` costs a short read rather than the allocation it asks
for.
+ self.buf.clear();
+ let read = (&mut self.reader)
+ .take(meta_len as u64)
+ .read_to_end(&mut self.buf)?;
+ if read != meta_len {
+ return Err(ArrowError::ParseError(format!(
+ "Unexpected end of stream: expected {meta_len} metadata bytes,
got {read}"
+ )));
+ }
let message =
crate::root_as_message(self.buf.as_slice()).map_err(|err| {
ArrowError::ParseError(format!("Unable to get root as message:
{err:?}"))
})?;
- let mut buf = MutableBuffer::from_len_zeroed(message.bodyLength() as
usize);
- self.reader.read_exact(&mut buf)?;
+ let body_len = usize::try_from(message.bodyLength()).map_err(|_| {
+ ArrowError::ParseError(format!(
+ "Invalid IPC message body length: {}",
+ message.bodyLength()
+ ))
+ })?;
+ let buf = read_body_bounded(&mut self.reader, body_len)?;
Ok(Some((message, buf)))
}
@@ -3859,4 +3905,80 @@ mod tests {
assert_eq!(read_batch.column(0).as_ref(), &values);
}
}
+
+ /// Builds a stream whose single message declares `body_length` but is
followed by only
+ /// `body_bytes` bytes of body. The body length is consumed before the
header is
+ /// interpreted, so the header itself is immaterial.
+ fn stream_with_declared_body(body_length: i64, body_bytes: usize) ->
Vec<u8> {
+ let mut fbb = flatbuffers::FlatBufferBuilder::new();
+ let mut message = crate::MessageBuilder::new(&mut fbb);
+ message.add_version(crate::MetadataVersion::V5);
+ message.add_header_type(crate::MessageHeader::NONE);
+ message.add_bodyLength(body_length);
+ let root = message.finish();
+ fbb.finish(root, None);
+ let metadata = fbb.finished_data();
+
+ let mut stream = Vec::new();
+ stream.extend_from_slice(&CONTINUATION_MARKER);
+ stream.extend_from_slice(&(metadata.len() as i32).to_le_bytes());
+ stream.extend_from_slice(metadata);
+ stream.resize(stream.len() + body_bytes, 0xAB);
+ stream
+ }
+
+ /// A message body length is read from the stream before any of the bytes
it describes,
+ /// so it cannot be trusted. Declaring an implausible one used to reserve
it outright,
+ /// and the resulting allocation failure aborts the process instead of
surfacing an error
+ /// the caller can handle.
+ #[test]
+ fn test_stream_reader_rejects_implausible_body_length() {
+ for body_length in [i64::MAX, 1 << 50, -1] {
+ let stream = stream_with_declared_body(body_length, 0);
+ let err = StreamReader::try_new(std::io::Cursor::new(stream),
None).expect_err(
+ &format!("a message declaring {body_length} body bytes must
not be accepted"),
+ );
+ let err = err.to_string();
+ assert!(
+ err.contains("Invalid IPC message body length")
+ || err.contains("Unexpected end of stream")
+ || err.contains("failed to fill whole buffer"),
+ "unexpected error for body_length {body_length}: {err}"
+ );
+ }
+ }
+
+ /// A plausible body length whose bytes end early must fail as a short
read: before the
+ /// first read for a small body, and after at least one growth step for a
body larger
+ /// than `MAX_PREALLOC_BYTES`.
+ #[test]
+ #[cfg_attr(miri, ignore)] // Takes too long
+ fn test_stream_reader_rejects_truncated_body() {
+ let over_prealloc = MAX_PREALLOC_BYTES as i64 + 1;
+ for (body_length, body_bytes) in [(1024, 10), (over_prealloc,
MAX_PREALLOC_BYTES)] {
+ let stream = stream_with_declared_body(body_length, body_bytes);
+ let err = StreamReader::try_new(std::io::Cursor::new(stream), None)
+ .expect_err("a body backed by fewer bytes than declared must
not be accepted");
+ assert!(
+ err.to_string().contains("failed to fill whole buffer"),
+ "unexpected error for body_length {body_length}: {err}"
+ );
+ }
+ }
+
+ /// The metadata length is untrusted for the same reason as the body
length.
+ #[test]
+ fn test_stream_reader_rejects_unbacked_metadata_length() {
+ let mut stream = Vec::new();
+ stream.extend_from_slice(&CONTINUATION_MARKER);
+ stream.extend_from_slice(&i32::MAX.to_le_bytes());
+ stream.extend_from_slice(b"not this many bytes");
+
+ let err = StreamReader::try_new(std::io::Cursor::new(stream), None)
+ .expect_err("metadata length must be backed by the stream");
+ assert!(
+ err.to_string().contains("Unexpected end of stream"),
+ "unexpected error: {err}"
+ );
+ }
}
