// SPDX-License-Identifier: AGPL-3.0-only
//! A total Arrow IPC *stream* decoder for exactly the column types a znippy
//! index uses — and nothing else.
//!
//! WHY NOT arrow-ipc's `StreamReader`: on hostile bytes it panics or aborts.
//! This crate's own property tests and fuzz targets found three: a 1 TiB
//! allocation from an attacker-chosen `bodyLength` (an abort, which no
//! `catch_unwind` can stop), a buffer sliced past its body
//! (`Buffer::slice_with_length`), a boolean buffer shorter than its length
//! (`BooleanBuffer::new`), and a schema conversion panic (`fb_to_schema`). In
//! the wasm guest every one of them is a trap, not a named refusal.
//!
//! So the only thing taken from arrow-ipc is its generated flatbuffers code:
//! `root_as_message` runs the flatbuffers verifier, after which the accessors
//! are bounds-checked by construction. Every buffer, offset and length is then
//! checked here against the bytes actually present, before it is read. Types:
//! Utf8, Int8, UInt32, UInt64, Boolean, FixedSizeBinary; no nulls, no
//! children, no dictionaries, no compression, little-endian only. Anything
//! else is a refusal, never a guess.
use std::collections::BTreeMap;
use arrow_ipc::{Endianness, MessageHeader, Type};
use crate::limits::MAX_INDEX_ROWS;
use crate::refusal::{Refusal, Result};
/// The largest IPC message metadata block read.
const MAX_META_BYTES: usize = 16 * 1024 * 1024;
/// The most fields a schema may declare.
const MAX_FIELDS: usize = 256;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ColType {
Utf8,
Int8,
UInt32,
UInt64,
Boolean,
FixedSizeBinary(usize),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Column {
Utf8(Vec<String>),
Int8(Vec<i8>),
UInt32(Vec<u32>),
UInt64(Vec<u64>),
Boolean(Vec<bool>),
FixedSizeBinary(usize, Vec<Vec<u8>>),
}
impl Column {
pub fn len(&self) -> usize {
match self {
Column::Utf8(v) => v.len(),
Column::Int8(v) => v.len(),
Column::UInt32(v) => v.len(),
Column::UInt64(v) => v.len(),
Column::Boolean(v) => v.len(),
Column::FixedSizeBinary(_, v) => v.len(),
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
/// One decoded record batch: columns by field name.
#[derive(Debug, Clone, Default)]
pub struct Batch {
pub rows: usize,
pub columns: BTreeMap<String, Column>,
}
/// One decoded stream.
#[derive(Debug, Clone, Default)]
pub struct Stream {
pub metadata: BTreeMap<String, String>,
pub fields: Vec<(String, ColType)>,
pub batches: Vec<Batch>,
}
fn bad(what: &str, why: impl core::fmt::Display) -> Refusal {
Refusal::Archive(format!("{what}: {why}"))
}
fn u32_at(b: &[u8], at: usize) -> Option<u32> {
Some(u32::from_le_bytes(
b.get(at..at.checked_add(4)?)?.try_into().ok()?,
))
}
/// Decode an IPC stream: a schema message, then record batches, then an
/// optional end marker (or the end of the bytes).
pub fn decode(bytes: &[u8], what: &str) -> Result<Stream> {
let mut pos = 0usize;
let mut out = Stream::default();
let mut have_schema = false;
while pos < bytes.len() {
let mut len = u32_at(bytes, pos).ok_or_else(|| bad(what, "truncated message length"))?;
pos += 4;
if len == 0xFFFF_FFFF {
len = u32_at(bytes, pos).ok_or_else(|| bad(what, "truncated message length"))?;
pos += 4;
}
if len == 0 {
break; // end-of-stream marker
}
let len = len as usize;
if len > MAX_META_BYTES || len > bytes.len() - pos {
return Err(bad(what, "message metadata runs past the end"));
}
let msg = arrow_ipc::root_as_message(&bytes[pos..pos + len])
.map_err(|e| bad(what, format!("message metadata: {e}")))?;
pos += len;
let body_len =
usize::try_from(msg.bodyLength()).map_err(|_| bad(what, "negative body length"))?;
if body_len > bytes.len() - pos {
return Err(bad(
what,
format!("a message claims a {body_len}-byte body, more than the stream holds"),
));
}
let body = &bytes[pos..pos + body_len];
pos += body_len;
match msg.header_type() {
MessageHeader::Schema => {
if have_schema {
return Err(bad(what, "a second schema message"));
}
let s = msg
.header_as_schema()
.ok_or_else(|| bad(what, "schema header missing"))?;
if s.endianness() != Endianness::Little {
return Err(bad(what, "a big-endian stream"));
}
read_schema(&s, what, &mut out)?;
have_schema = true;
}
MessageHeader::RecordBatch => {
if !have_schema {
return Err(bad(what, "a record batch before the schema"));
}
let rb = msg
.header_as_record_batch()
.ok_or_else(|| bad(what, "record batch header missing"))?;
let batch = read_batch(&rb, body, &out.fields, what)?;
let total: usize = out.batches.iter().map(|b| b.rows).sum::<usize>() + batch.rows;
if total > MAX_INDEX_ROWS {
return Err(Refusal::LimitExceeded(format!(
"{what}: more than {MAX_INDEX_ROWS} rows"
)));
}
out.batches.push(batch);
}
_ => return Err(bad(what, "an IPC message of a kind znippy never writes")),
}
}
if !have_schema {
return Err(bad(what, "no schema message"));
}
Ok(out)
}
fn read_schema(s: &arrow_ipc::Schema<'_>, what: &str, out: &mut Stream) -> Result<()> {
if let Some(kv) = s.custom_metadata() {
for e in kv.iter() {
if let (Some(k), Some(v)) = (e.key(), e.value()) {
out.metadata.insert(k.to_string(), v.to_string());
}
}
}
let fields = s
.fields()
.ok_or_else(|| bad(what, "schema has no fields"))?;
if fields.len() > MAX_FIELDS {
return Err(bad(what, "too many fields"));
}
for f in fields.iter() {
let name = f.name().unwrap_or("").to_string();
if f.dictionary().is_some() {
return Err(bad(what, format!("field {name:?} is dictionary-encoded")));
}
if f.children().is_some_and(|c| !c.is_empty()) {
return Err(bad(what, format!("field {name:?} is nested")));
}
let t = match f.type_type() {
Type::Utf8 => ColType::Utf8,
Type::Bool => ColType::Boolean,
Type::Int => {
let i = f
.type_as_int()
.ok_or_else(|| bad(what, format!("field {name:?}: int type missing")))?;
match (i.bitWidth(), i.is_signed()) {
(8, true) => ColType::Int8,
(32, false) => ColType::UInt32,
(64, false) => ColType::UInt64,
(w, s) => {
return Err(bad(
what,
format!("field {name:?}: unsupported int{w} signed={s}"),
))
}
}
}
Type::FixedSizeBinary => {
let w = f
.type_as_fixed_size_binary()
.ok_or_else(|| bad(what, format!("field {name:?}: width missing")))?
.byteWidth();
if !(1..=4096).contains(&w) {
return Err(bad(what, format!("field {name:?}: byte width {w}")));
}
ColType::FixedSizeBinary(w as usize)
}
other => {
return Err(bad(
what,
format!("field {name:?} has a type znippy's index never uses ({other:?})"),
))
}
};
out.fields.push((name, t));
}
Ok(())
}
fn read_batch(
rb: &arrow_ipc::RecordBatch<'_>,
body: &[u8],
fields: &[(String, ColType)],
what: &str,
) -> Result<Batch> {
if rb.compression().is_some() {
return Err(Refusal::Codec(format!(
"{what}: a compressed IPC record batch; znippy writes none"
)));
}
let rows = usize::try_from(rb.length())
.ok()
.filter(|&n| n <= MAX_INDEX_ROWS)
.ok_or_else(|| bad(what, format!("a record batch of {} rows", rb.length())))?;
let nodes = rb.nodes().ok_or_else(|| bad(what, "no field nodes"))?;
let buffers = rb.buffers().ok_or_else(|| bad(what, "no buffers"))?;
if nodes.len() != fields.len() {
return Err(bad(what, "field node count differs from the schema"));
}
let mut buf_i = 0usize;
let mut next_buf = || -> Result<&[u8]> {
if buf_i >= buffers.len() {
return Err(bad(what, "fewer buffers than the schema needs"));
}
let b = buffers.get(buf_i);
buf_i += 1;
let (off, len) = (b.offset(), b.length());
if off < 0 || len < 0 {
return Err(bad(what, "a negative buffer offset or length"));
}
let (off, len) = (off as u64, len as u64);
let end = off.saturating_add(len);
if end > body.len() as u64 {
return Err(bad(what, "a buffer lies outside its message body"));
}
Ok(&body[off as usize..end as usize])
};
let mut batch = Batch {
rows,
columns: BTreeMap::new(),
};
for (i, (name, t)) in fields.iter().enumerate() {
let node = nodes.get(i);
if node.length() != rows as i64 {
return Err(bad(
what,
format!("column {name:?} length differs from the batch"),
));
}
if node.null_count() != 0 {
return Err(bad(what, format!("column {name:?} holds nulls")));
}
let _validity = next_buf()?;
let col = match t {
ColType::Utf8 => {
let offs = next_buf()?;
let data = next_buf()?;
let need = rows
.checked_add(1)
.and_then(|n| n.checked_mul(4))
.ok_or_else(|| bad(what, "row count overflows"))?;
if offs.len() < need {
return Err(bad(
what,
format!("column {name:?}: offsets buffer too short"),
));
}
let at =
|k: usize| i32::from_le_bytes(offs[k * 4..k * 4 + 4].try_into().expect("4"));
let mut v = Vec::with_capacity(rows);
for r in 0..rows {
let (a, b) = (at(r), at(r + 1));
if a < 0 || b < a || b as usize > data.len() {
return Err(bad(what, format!("column {name:?}: bad string offsets")));
}
let s = core::str::from_utf8(&data[a as usize..b as usize]).map_err(|_| {
bad(what, format!("column {name:?}: a string is not UTF-8"))
})?;
v.push(s.to_string());
}
Column::Utf8(v)
}
ColType::Int8 => {
let d = next_buf()?;
if d.len() < rows {
return Err(bad(what, format!("column {name:?}: buffer too short")));
}
Column::Int8(d[..rows].iter().map(|&b| b as i8).collect())
}
ColType::UInt32 => {
let d = next_buf()?;
if d.len() / 4 < rows {
return Err(bad(what, format!("column {name:?}: buffer too short")));
}
Column::UInt32(
(0..rows)
.map(|r| u32::from_le_bytes(d[r * 4..r * 4 + 4].try_into().expect("4")))
.collect(),
)
}
ColType::UInt64 => {
let d = next_buf()?;
if d.len() / 8 < rows {
return Err(bad(what, format!("column {name:?}: buffer too short")));
}
Column::UInt64(
(0..rows)
.map(|r| u64::from_le_bytes(d[r * 8..r * 8 + 8].try_into().expect("8")))
.collect(),
)
}
ColType::Boolean => {
let d = next_buf()?;
if d.len() < rows.div_ceil(8) {
return Err(bad(what, format!("column {name:?}: buffer too short")));
}
Column::Boolean((0..rows).map(|r| d[r / 8] & (1 << (r % 8)) != 0).collect())
}
ColType::FixedSizeBinary(w) => {
let d = next_buf()?;
if d.len() / w < rows {
return Err(bad(what, format!("column {name:?}: buffer too short")));
}
Column::FixedSizeBinary(
*w,
(0..rows).map(|r| d[r * w..(r + 1) * w].to_vec()).collect(),
)
}
};
if batch.columns.insert(name.clone(), col).is_some() {
return Err(bad(what, format!("column {name:?} appears twice")));
}
}
Ok(batch)
}