znippy-common/src/ipc_guard.rs
//! The one door every Arrow IPC stream read from an archive goes through.
//!
//! **Why it exists.** arrow-ipc's `StreamReader` (arrow 59) trusts the bytes it
//! is handed. On a hostile `.znippy` it does not return `Err`, it panics or
//! aborts the process — found by the Heimdall verify-bundle guest lane and
//! reproduced on `ZnippyArchive::open` (`tests/tests/hostile_ipc.rs`):
//!
//! * a message claiming a 1 TiB `bodyLength` is allocated before a byte of it is
//! read (`MutableBuffer::from_len_zeroed`): an **abort** no `catch_unwind` stops;
//! a negative one is a capacity-overflow panic in the same place;
//! * a buffer whose offset/length lies past its message body panics in
//! `Buffer::slice_with_length`;
//! * a validity bitmap shorter than its array panics in `BooleanBuffer::new`;
//! * a schema with no field vector (or an unknown/half-described type) panics in
//! `fb_to_schema` / `get_data_type`.
//!
//! **What it does.** Before arrow sees the stream, [`check_stream`] walks it with
//! the flatbuffers verifier (arrow's own generated `root_as_message`, which is
//! total) and checks every length arrow is about to trust against the bytes that
//! are actually there: message framing, body length, every field's type (only
//! the flat types znippy writes: bool, int 8/16/32/64, float 16/32/64, utf8,
//! binary and their large forms, fixed-size binary, list and large list of
//! those), every field node and every buffer of every record batch, including
//! the minimum size each buffer needs for its node's length. Anything else is an
//! `Err` naming what was wrong. It allocates nothing proportional to a declared
//! length; it reads the bytes once and copies none.
//!
//! What is left to arrow is what arrow already turns into `Err`: offsets that go
//! backwards or past their values, invalid UTF-8, column lengths that disagree
//! with the batch (full `ArrayData` validation stays on).
//!
//! Every reader in this crate opens its stream with [`stream_reader`] (or, for
//! the hot journal, which is a log with a possibly torn tail, [`valid_prefix`]).
use std::io::Cursor;
use anyhow::{Result, anyhow, bail};
use arrow::ipc::reader::StreamReader;
use arrow::ipc::{self as fb, MessageHeader, MetadataVersion, Type};
/// The IPC continuation marker that precedes a message length since Arrow 0.15.
const CONTINUATION: u32 = 0xFFFF_FFFF;
/// Fields per schema, counting nested children. znippy writes a few dozen.
const MAX_FIELDS: usize = 4096;
/// List nesting depth. znippy writes one level.
const MAX_DEPTH: usize = 8;
/// Widest `FixedSizeBinary` accepted (znippy writes 20 and 32).
const MAX_FIXED_WIDTH: i32 = 1 << 20;
/// The physical layout of one field, as the reader will consume it.
#[derive(Debug)]
enum Shape {
Bool,
/// A primitive of this many bytes per value (1, 2, 4 or 8).
Fixed(usize),
/// A fixed-size binary of this many bytes per value.
Fsb(usize),
/// Utf8/Binary (4-byte offsets) or their Large forms (8-byte offsets).
Var(usize),
/// List (4-byte offsets) or LargeList (8-byte offsets) of one child.
List(usize, Box<Shape>),
}
/// Open an Arrow IPC stream held in memory, after [`check_stream`] passed it.
pub fn stream_reader<B: AsRef<[u8]>>(bytes: B) -> Result<StreamReader<Cursor<B>>> {
check_stream(bytes.as_ref())?;
Ok(StreamReader::try_new(Cursor::new(bytes), None)?)
}
/// Refuse a stream arrow-ipc could panic, abort or over-allocate on.
pub fn check_stream(bytes: &[u8]) -> Result<()> {
let end = walk(bytes, false)?;
debug_assert!(end <= bytes.len());
Ok(())
}
/// Refuse a record-batch message (or a run of whole messages) that arrow-ipc
/// could panic, abort or over-allocate on, read under the schema of a SEPARATE
/// schema message — the shape of a table that keeps its schema message once and
/// feeds a `StreamDecoder` one batch at a time from wherever the batch lies
/// (`schema` = one whole schema message, `batches` = the bytes handed to the
/// decoder after it). Nothing is copied or concatenated.
pub fn check_batches_after_schema(schema: &[u8], batches: &[u8]) -> Result<()> {
let mut shapes = None;
let mut pos = 0usize;
if !message(schema, &mut pos, &mut shapes)? || shapes.is_none() {
bail!("IPC: the schema side holds no schema message");
}
if pos != schema.len() {
bail!("IPC: the schema side holds more than one message");
}
let mut pos = 0usize;
while pos < batches.len() {
if !message(batches, &mut pos, &mut shapes)? {
break;
}
}
Ok(())
}
/// The longest prefix of `bytes` made of whole, well-formed messages — for a log
/// whose last message may be torn. Never errors: a stream that is bad from its
/// first byte has an empty valid prefix.
pub fn valid_prefix(bytes: &[u8]) -> &[u8] {
let end = walk(bytes, true).unwrap_or(0);
&bytes[..end]
}
/// Walk every message. Strict: any defect is an `Err`. Lenient (`prefix`): stop
/// at the first defect and return how far the stream was good.
fn walk(bytes: &[u8], prefix: bool) -> Result<usize> {
let mut pos = 0usize;
let mut schema: Option<Vec<Shape>> = None;
while pos < bytes.len() {
let start = pos;
match message(bytes, &mut pos, &mut schema) {
Ok(true) => continue,
Ok(false) => break, // end-of-stream marker
Err(_) if prefix => return Ok(start),
Err(e) => return Err(e),
}
}
if schema.is_none() && !prefix {
bail!("Arrow IPC stream has no schema message");
}
Ok(pos)
}
fn u32_at(b: &[u8], at: usize) -> Option<u32> {
Some(u32::from_le_bytes(
b.get(at..at.checked_add(4)?)?.try_into().ok()?,
))
}
/// Check one message at `*pos` and advance past it. `Ok(false)` at the
/// end-of-stream marker.
fn message(bytes: &[u8], pos: &mut usize, schema: &mut Option<Vec<Shape>>) -> Result<bool> {
let mut at = *pos;
let mut len = u32_at(bytes, at).ok_or_else(|| anyhow!("IPC: truncated message length"))?;
at += 4;
if len == CONTINUATION {
len = u32_at(bytes, at).ok_or_else(|| anyhow!("IPC: truncated message length"))?;
at += 4;
}
// arrow reads the length as an i32 and refuses a negative one.
if len as i32 <= 0 {
if len == 0 {
*pos = at;
return Ok(false);
}
bail!("IPC: negative message length");
}
let len = len as usize;
let rest = bytes.len() - at;
if len > rest {
bail!("IPC: message metadata of {len} bytes runs past the {rest} left");
}
let msg = fb::root_as_message(&bytes[at..at + len])
.map_err(|e| anyhow!("IPC: message metadata is not a valid flatbuffer: {e}"))?;
at += len;
if msg.version() < MetadataVersion::V4 || msg.version() > MetadataVersion::V5 {
bail!(
"IPC: metadata version {:?} is not one znippy writes",
msg.version()
);
}
let body_len = usize::try_from(msg.bodyLength())
.map_err(|_| anyhow!("IPC: negative body length {}", msg.bodyLength()))?;
let rest = bytes.len() - at;
if body_len > rest {
bail!("IPC: a message claims a {body_len}-byte body, only {rest} bytes follow");
}
match msg.header_type() {
MessageHeader::Schema => {
if schema.is_some() {
bail!("IPC: a second schema message");
}
let s = msg
.header_as_schema()
.ok_or_else(|| anyhow!("IPC: schema header missing"))?;
*schema = Some(check_schema(&s)?);
}
MessageHeader::RecordBatch => {
let shapes = schema
.as_ref()
.ok_or_else(|| anyhow!("IPC: a record batch before the schema"))?;
let rb = msg
.header_as_record_batch()
.ok_or_else(|| anyhow!("IPC: record batch header missing"))?;
check_batch(&rb, body_len, shapes)?;
}
other => bail!("IPC: a {other:?} message, which znippy never writes"),
}
*pos = at + body_len;
Ok(true)
}
fn check_schema(s: &fb::Schema<'_>) -> Result<Vec<Shape>> {
if s.endianness() != fb::Endianness::Little {
bail!("IPC: a big-endian schema");
}
let fields = s
.fields()
.ok_or_else(|| anyhow!("IPC: schema has no field list"))?;
let mut count = 0usize;
fields
.iter()
.map(|f| field_shape(&f, 0, &mut count))
.collect()
}
fn field_shape(f: &fb::Field<'_>, depth: usize, count: &mut usize) -> Result<Shape> {
*count += 1;
if *count > MAX_FIELDS {
bail!("IPC: more than {MAX_FIELDS} fields");
}
let name = f.name().unwrap_or("");
if f.dictionary().is_some() {
bail!("IPC: field {name:?} is dictionary-encoded");
}
let children = f.children().map(|c| c.len()).unwrap_or(0);
let leaf = |shape: Shape| -> Result<Shape> {
if children != 0 {
bail!("IPC: field {name:?} of a flat type has children");
}
Ok(shape)
};
let missing = || anyhow!("IPC: field {name:?}: type table missing");
match f.type_type() {
Type::Bool => leaf(Shape::Bool),
Type::Int => {
let i = f.type_as_int().ok_or_else(missing)?;
match i.bitWidth() {
8 | 16 | 32 | 64 => leaf(Shape::Fixed(i.bitWidth() as usize / 8)),
w => bail!("IPC: field {name:?}: int of {w} bits"),
}
}
Type::FloatingPoint => {
let p = f.type_as_floating_point().ok_or_else(missing)?;
match p.precision() {
fb::Precision::HALF => leaf(Shape::Fixed(2)),
fb::Precision::SINGLE => leaf(Shape::Fixed(4)),
fb::Precision::DOUBLE => leaf(Shape::Fixed(8)),
p => bail!("IPC: field {name:?}: float precision {p:?}"),
}
}
Type::Utf8 | Type::Binary => leaf(Shape::Var(4)),
Type::LargeUtf8 | Type::LargeBinary => leaf(Shape::Var(8)),
Type::FixedSizeBinary => {
let w = f
.type_as_fixed_size_binary()
.ok_or_else(missing)?
.byteWidth();
if !(1..=MAX_FIXED_WIDTH).contains(&w) {
bail!("IPC: field {name:?}: fixed-size binary of width {w}");
}
leaf(Shape::Fsb(w as usize))
}
t @ (Type::List | Type::LargeList) => {
if depth >= MAX_DEPTH {
bail!("IPC: field {name:?} nests deeper than {MAX_DEPTH}");
}
let kids = f
.children()
.ok_or_else(|| anyhow!("IPC: list {name:?} has no child"))?;
if kids.len() != 1 {
bail!("IPC: list {name:?} has {} children, not one", kids.len());
}
let child = field_shape(&kids.get(0), depth + 1, count)?;
Ok(Shape::List(
if t == Type::List { 4 } else { 8 },
Box::new(child),
))
}
t => bail!("IPC: field {name:?} has type {t:?}, which znippy never writes"),
}
}
/// The field nodes and buffers of one record batch, consumed in arrow's order.
/// `node(i)` / `buf(i)` read entry `i` straight out of the verified flatbuffer
/// (there are as many as the metadata really holds, never a declared count):
/// collecting them into vectors first was a fifth of the guard's time on a log
/// of small batches.
struct Batch<N, B> {
node_at: N,
n_nodes: usize,
buf_at: B,
n_buffers: usize,
node_i: usize,
buf_i: usize,
body_len: u64,
}
impl<N: Fn(usize) -> (i64, i64), B: Fn(usize) -> (i64, i64)> Batch<N, B> {
/// The next field node: its length and null count, both checked.
fn node(&mut self) -> Result<usize> {
if self.node_i >= self.n_nodes {
bail!("IPC: fewer field nodes than the schema needs");
}
let (len, nulls) = (self.node_at)(self.node_i);
self.node_i += 1;
if len < 0 || nulls < 0 || nulls > len {
bail!("IPC: a field node of length {len} with {nulls} nulls");
}
usize::try_from(len).map_err(|_| anyhow!("IPC: field node length {len}"))
}
/// The next buffer: inside the body, at least `need` bytes long and, when it
/// is read as typed values (`unit` > 1), a whole number of them — arrow's
/// `Buffer::typed_data` asserts that on the WHOLE buffer during validation.
fn buffer(&mut self, need: Option<u64>, unit: u64, what: &str) -> Result<()> {
if self.buf_i >= self.n_buffers {
bail!("IPC: fewer buffers than the schema needs");
}
let (off, len) = (self.buf_at)(self.buf_i);
self.buf_i += 1;
if off < 0 || len < 0 {
bail!("IPC: {what} buffer at {off} of length {len}");
}
let (off, len) = (off as u64, len as u64);
if off.checked_add(len).is_none_or(|end| end > self.body_len) {
bail!(
"IPC: {what} buffer [{off}, +{len}) lies outside its {}-byte body",
self.body_len
);
}
if let Some(need) = need
&& len < need
{
bail!("IPC: {what} buffer of {len} bytes, {need} needed");
}
if len % unit != 0 {
bail!("IPC: {what} buffer of {len} bytes is not a whole number of {unit}-byte values");
}
Ok(())
}
/// One array of `shape`, recursively, exactly as arrow's `create_array` walks it.
fn array(&mut self, shape: &Shape, expect_len: Option<usize>) -> Result<()> {
// A node is checked by its own length and null count first.
let node_i = self.node_i;
let len = self.node()?;
if let Some(e) = expect_len
&& len != e
{
bail!("IPC: a column of {len} rows in a batch of {e}");
}
let nulls = (self.node_at)(node_i).1;
let len64 = len as u64;
let bits = len64.div_ceil(8);
// The validity bitmap is read (BooleanBuffer::new) only when there are nulls.
self.buffer((nulls > 0).then_some(bits), 1, "validity")?;
let mul = |w: usize| {
len64
.checked_mul(w as u64)
.ok_or_else(|| anyhow!("IPC: buffer size overflows"))
};
// Offsets: len + 1 of them; an empty array may carry an empty buffer.
let offsets = |w: usize| -> Result<Option<u64>> {
if len == 0 {
return Ok(None);
}
Ok(Some(
(len64 + 1)
.checked_mul(w as u64)
.ok_or_else(|| anyhow!("IPC: buffer size overflows"))?,
))
};
match shape {
Shape::Bool => self.buffer(Some(bits), 1, "boolean values"),
Shape::Fixed(w) => self.buffer(Some(mul(*w)?), *w as u64, "values"),
Shape::Fsb(w) => self.buffer(Some(mul(*w)?), 1, "fixed-size values"),
Shape::Var(w) => {
self.buffer(offsets(*w)?, *w as u64, "offsets")?;
self.buffer(None, 1, "data")
}
Shape::List(w, child) => {
self.buffer(offsets(*w)?, *w as u64, "list offsets")?;
self.array(child, None)
}
}
}
}
fn check_batch(rb: &fb::RecordBatch<'_>, body_len: usize, shapes: &[Shape]) -> Result<()> {
if rb.compression().is_some() {
bail!("IPC: a compressed record batch; znippy writes none");
}
if rb.variadicBufferCounts().is_some_and(|v| !v.is_empty()) {
bail!("IPC: variadic buffer counts on a stream with no view types");
}
let rows = usize::try_from(rb.length())
.map_err(|_| anyhow!("IPC: a record batch of {} rows", rb.length()))?;
let (nodes, buffers) = match (rb.nodes(), rb.buffers()) {
(Some(n), Some(b)) => (n, b),
_ if shapes.is_empty() => return Ok(()),
_ => bail!("IPC: a record batch without field nodes or buffers"),
};
let mut b = Batch {
node_at: |i: usize| {
let n = nodes.get(i);
(n.length(), n.null_count())
},
n_nodes: nodes.len(),
buf_at: |i: usize| {
let x = buffers.get(i);
(x.offset(), x.length())
},
n_buffers: buffers.len(),
node_i: 0,
buf_i: 0,
body_len: body_len as u64,
};
for shape in shapes {
b.array(shape, Some(rows))?;
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use arrow::array::{
ArrayRef, BooleanArray, FixedSizeBinaryArray, Int64Array, ListArray, StringArray,
UInt64Array,
};
use arrow::datatypes::Int32Type;
use arrow::ipc::writer::StreamWriter;
use arrow::record_batch::RecordBatch;
use std::sync::Arc;
fn stream(cols: Vec<(&str, ArrayRef)>) -> Vec<u8> {
let batch = RecordBatch::try_from_iter(cols).unwrap();
let mut out = Vec::new();
let mut w = StreamWriter::try_new(&mut out, &batch.schema()).unwrap();
w.write(&batch).unwrap();
w.write(&batch).unwrap();
w.finish().unwrap();
out
}
fn every_type() -> Vec<u8> {
let list = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
Some(vec![Some(1), None]),
None,
Some(vec![]),
]);
stream(vec![
(
"s",
Arc::new(StringArray::from(vec![Some("a"), None, Some("ccc")])) as ArrayRef,
),
("u", Arc::new(UInt64Array::from(vec![1, 2, 3]))),
(
"i",
Arc::new(Int64Array::from(vec![Some(-1), None, Some(3)])),
),
(
"b",
Arc::new(BooleanArray::from(vec![Some(true), None, Some(false)])),
),
(
"k",
Arc::new(
FixedSizeBinaryArray::try_from_iter(vec![[1u8; 4], [2; 4], [3; 4]].into_iter())
.unwrap(),
),
),
("l", Arc::new(list)),
])
}
#[test]
fn every_stream_arrow_writes_for_znippy_types_passes_and_reads_back() {
let bytes = every_type();
check_stream(&bytes).unwrap();
let batches: Vec<_> = stream_reader(&bytes[..])
.unwrap()
.collect::<Result<_, _>>()
.unwrap();
assert_eq!(batches.len(), 2);
assert_eq!(valid_prefix(&bytes).len(), bytes.len());
}
#[test]
fn every_truncation_is_an_error_and_never_a_panic() {
let bytes = every_type();
for cut in 0..bytes.len() {
// A cut on a message boundary is a shorter valid stream; anything
// else must be refused. Either way arrow then never panics on it.
if check_stream(&bytes[..cut]).is_ok() {
for b in StreamReader::try_new(Cursor::new(&bytes[..cut]), None).unwrap() {
let _ = b;
}
}
let p = valid_prefix(&bytes[..cut]);
assert!(p.len() <= cut);
if !p.is_empty() {
check_stream(p).unwrap();
}
}
}
#[test]
fn every_single_byte_flip_is_checked_or_survived_by_arrow() {
// The guard's contract in one loop: whatever it lets through, arrow
// decodes without panicking (std::panic::catch_unwind here is only the
// test's detector, not part of the fix).
let bytes = every_type();
for i in 0..bytes.len() {
for v in [0x00u8, 0xFF, 0x7F, 0x80, bytes[i] ^ 0x01] {
let mut m = bytes.clone();
m[i] = v;
if check_stream(&m).is_ok() {
let r = std::panic::catch_unwind(|| {
for b in StreamReader::try_new(Cursor::new(&m), None).unwrap() {
let _ = b;
}
});
assert!(
r.is_ok(),
"byte {i} = {v:#x} passed the guard and panicked arrow"
);
}
}
}
}
#[test]
fn a_huge_body_length_is_refused_before_anything_is_allocated() {
let bytes = every_type();
// The first record batch's bodyLength, patched to 1 TiB.
let schema_meta = u32::from_le_bytes(bytes[4..8].try_into().unwrap()) as usize;
let rb_at = 8 + schema_meta;
let meta_len = u32::from_le_bytes(bytes[rb_at + 4..rb_at + 8].try_into().unwrap()) as usize;
let meta = &bytes[rb_at + 8..rb_at + 8 + meta_len];
let body = fb::root_as_message(meta).unwrap().bodyLength();
let at = meta
.windows(8)
.position(|w| w == body.to_le_bytes())
.unwrap()
+ rb_at
+ 8;
for v in [1i64 << 40, -8, i64::MAX] {
let mut m = bytes.clone();
m[at..at + 8].copy_from_slice(&v.to_le_bytes());
let e = check_stream(&m).unwrap_err().to_string();
assert!(e.contains("body"), "{v}: {e}");
}
}
/// The first message of `stream` (the schema) and the rest, split at the
/// message boundary — what a table holding its schema message apart hands over.
fn split_schema(stream: &[u8]) -> (&[u8], &[u8]) {
let meta = u32::from_le_bytes(stream[4..8].try_into().unwrap()) as usize;
stream.split_at(8 + meta)
}
fn decode_after_schema(schema: &[u8], batches: &[u8]) {
use arrow::buffer::Buffer;
use arrow::ipc::reader::StreamDecoder;
let mut dec = StreamDecoder::new();
let mut head = Buffer::from(schema.to_vec());
if dec.decode(&mut head).is_err() {
return;
}
let mut body = Buffer::from(batches.to_vec());
while !body.is_empty() {
match dec.decode(&mut body) {
Ok(_) => {}
Err(_) => return,
}
}
}
#[test]
fn a_batch_after_a_separate_schema_message_passes_and_decodes() {
let bytes = every_type();
let (schema, rest) = split_schema(&bytes);
check_batches_after_schema(schema, rest).unwrap();
decode_after_schema(schema, rest);
}
#[test]
fn a_batch_after_a_separate_schema_is_checked_like_a_stream() {
let bytes = every_type();
let (schema, rest) = split_schema(&bytes);
// A huge body, a schema that is not a schema, garbage, and a batch with
// no schema at all are each refused.
let rb_meta = u32::from_le_bytes(rest[4..8].try_into().unwrap()) as usize;
let meta = &rest[8..8 + rb_meta];
let body = fb::root_as_message(meta).unwrap().bodyLength();
let at = meta
.windows(8)
.position(|w| w == body.to_le_bytes())
.unwrap()
+ 8;
for v in [1i64 << 40, -8, i64::MAX] {
let mut m = rest.to_vec();
m[at..at + 8].copy_from_slice(&v.to_le_bytes());
let e = check_batches_after_schema(schema, &m)
.unwrap_err()
.to_string();
assert!(e.contains("body"), "{v}: {e}");
}
assert!(check_batches_after_schema(rest, rest).is_err());
assert!(check_batches_after_schema(schema, b"garbage!").is_err());
assert!(check_batches_after_schema(b"", rest).is_err());
assert!(
check_batches_after_schema(&bytes, rest).is_err(),
"two messages are not one schema"
);
// Every truncation of the batch side is refused unless it ends on a
// message boundary, and every single-byte flip on either side that the
// guard lets through decodes without a panic.
for cut in 0..rest.len() {
if check_batches_after_schema(schema, &rest[..cut]).is_ok() {
decode_after_schema(schema, &rest[..cut]);
}
}
for side in 0..2 {
let src = if side == 0 { schema } else { rest };
for i in 0..src.len() {
for v in [0x00u8, 0xFF, 0x7F, 0x80, src[i] ^ 0x01] {
let mut m = src.to_vec();
m[i] = v;
let (s, b) = if side == 0 {
(&m[..], rest)
} else {
(schema, &m[..])
};
if check_batches_after_schema(s, b).is_ok() {
let r = std::panic::catch_unwind(|| decode_after_schema(s, b));
assert!(
r.is_ok(),
"side {side} byte {i} = {v:#x} passed the guard and panicked arrow"
);
}
}
}
}
}
#[test]
fn empty_and_garbage_are_errors() {
assert!(check_stream(b"").is_err());
assert!(check_stream(&[0xFF; 3]).is_err());
assert!(check_stream(&[0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x7F]).is_err());
assert!(check_stream(&[0xFF, 0xFF, 0xFF, 0xFF, 0x10, 0, 0, 0, 1, 2, 3]).is_err());
assert!(valid_prefix(b"garbage").is_empty());
}
}