znippy-common 0.10.4

//! 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(())
}

/// 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.
struct Batch {
    /// `(length, null_count)` per node, copied out of the verified flatbuffer
    /// (as many as the metadata really holds, never a declared count).
    nodes: Vec<(i64, i64)>,
    /// `(offset, length)` per buffer, likewise.
    buffers: Vec<(i64, i64)>,
    node_i: usize,
    buf_i: usize,
    body_len: u64,
}

impl Batch {
    /// The next field node: its length and null count, both checked.
    fn node(&mut self) -> Result<usize> {
        if self.node_i >= self.nodes.len() {
            bail!("IPC: fewer field nodes than the schema needs");
        }
        let (len, nulls) = self.nodes[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.buffers.len() {
            bail!("IPC: fewer buffers than the schema needs");
        }
        let (off, len) = self.buffers[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.nodes[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 {
        nodes: nodes.iter().map(|n| (n.length(), n.null_count())).collect(),
        buffers: buffers.iter().map(|x| (x.offset(), x.length())).collect(),
        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}");
        }
    }

    #[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());
    }
}