From 7c010560f3d31ba5e63078f66903c357622f8480 Mon Sep 17 00:00:00 2001 From: Beinan Date: Mon, 19 Oct 2026 05:01:43 +0000 Subject: [PATCH 1/2] feat: share validated Arrow IPC decoding for WAL materialization --- .github/workflows/rust-test.yml | 4 +- .github/workflows/style.yml | 2 +- crates/lance-context-wal/Cargo.toml | 4 +- crates/lance-context-wal/README.md | 14 ++ crates/lance-context-wal/src/bin/lab.rs | 32 ++--- crates/lance-context-wal/src/ipc.rs | 157 ++++++++++++++++++++++ crates/lance-context-wal/src/ipc/tests.rs | 116 ++++++++++++++++ crates/lance-context-wal/src/lib.rs | 3 + 8 files changed, 306 insertions(+), 26 deletions(-) create mode 100644 crates/lance-context-wal/src/ipc.rs create mode 100644 crates/lance-context-wal/src/ipc/tests.rs diff --git a/.github/workflows/rust-test.yml b/.github/workflows/rust-test.yml index c9e1023..88ae877 100644 --- a/.github/workflows/rust-test.yml +++ b/.github/workflows/rust-test.yml @@ -63,8 +63,8 @@ jobs: run: | cargo run -p lance-context-wal --features lab --bin context-wal-lab -- codec --payload-bytes 1024 --rows 8 --iterations 2 cargo run -p lance-context-wal --features lab --bin context-wal-lab -- disk --directory "$RUNNER_TEMP/wal-disk-smoke" --payload-bytes 1024 --batch 4 --writers 2 --groups-per-writer 2 --rounds 8 - - name: WAL placement and shared admission tests - run: cargo test -p lance-context-wal --features placement,admission --lib + - name: WAL placement, admission and IPC decoding tests + run: cargo test -p lance-context-wal --features placement,admission,ipc --lib - name: Start etcd for HA scheduler test run: | ETCD_VERSION=3.7.0 diff --git a/.github/workflows/style.yml b/.github/workflows/style.yml index 43c97b1..70b94a0 100644 --- a/.github/workflows/style.yml +++ b/.github/workflows/style.yml @@ -56,7 +56,7 @@ jobs: - name: Clippy (lance-context, default features) run: cargo clippy --manifest-path crates/lance-context/Cargo.toml --all-targets -- -D warnings - name: Clippy (WAL experiments) - run: cargo clippy -p lance-context-wal --features lab,placement,admission --all-targets -- -D warnings + run: cargo clippy -p lance-context-wal --features lab,placement,admission,ipc --all-targets -- -D warnings typos: name: Spell Check diff --git a/crates/lance-context-wal/Cargo.toml b/crates/lance-context-wal/Cargo.toml index 58b5958..340fa6c 100644 --- a/crates/lance-context-wal/Cargo.toml +++ b/crates/lance-context-wal/Cargo.toml @@ -9,7 +9,8 @@ description = "Independent Rust WAL experiments owned by lance-context" [features] placement = ["dep:anyhow", "dep:serde", "dep:serde_json", "dep:sha2", "dep:uuid"] admission = ["dep:anyhow", "dep:serde"] -lab = ["dep:anyhow", "dep:arrow-array", "dep:arrow-buffer", "dep:arrow-ipc", "dep:arrow-schema", "dep:bytes", "dep:clap", "dep:hdrhistogram", "dep:raft-engine", "dep:raft-proto", "dep:serde_json"] +ipc = ["dep:anyhow", "dep:arrow-array", "dep:arrow-buffer", "dep:arrow-ipc", "dep:arrow-schema", "dep:bytes"] +lab = ["ipc", "dep:clap", "dep:hdrhistogram", "dep:raft-engine", "dep:raft-proto", "dep:serde_json"] [dependencies] anyhow = { version = "1", optional = true } @@ -28,6 +29,7 @@ sha2 = { version = "0.10", optional = true } uuid = { version = "1", features = ["serde"], optional = true } [dev-dependencies] +arrow-ipc = { version = "58", features = ["zstd"] } tempfile = "3" [[bin]] diff --git a/crates/lance-context-wal/README.md b/crates/lance-context-wal/README.md index f99e6c1..aedd41a 100644 --- a/crates/lance-context-wal/README.md +++ b/crates/lance-context-wal/README.md @@ -3,6 +3,20 @@ Feature-gated Rust measurements for choosing a WAL payload path and durable local log engine. This crate does not change production ingestion or Lance OSS. +The optional `ipc` feature exposes `ipc::decode(Bytes, expected_schema)` for +materializers. It decodes multiple uncompressed record batches and dictionaries +while sharing binary payload ownership with the input. Numeric alignment repair, +schema parsing and Arrow metadata still allocate. The codec validates footer and +block ranges, schema equality, row/count limits and unsupported expansion before +constructing arrays. Compressed batches and dictionary deltas are rejected. +Callers still reserve input and decoding workspace within their own memory budget. +Replication does not need to enable this module or decode Arrow. The `lab` codec +measurement now uses the same validated path, including its validation overhead. + +```bash +cargo test -p lance-context-wal --features ipc --lib +``` + The independent `placement` feature provides `placement::Layout`, session bucket hashing and scoped `Receipt` types. A validated layout assigns each stable virtual bucket to one RF3 group, with explicit voters and StatefulSet DNS addresses. Use a diff --git a/crates/lance-context-wal/src/bin/lab.rs b/crates/lance-context-wal/src/bin/lab.rs index 733282e..33a1820 100644 --- a/crates/lance-context-wal/src/bin/lab.rs +++ b/crates/lance-context-wal/src/bin/lab.rs @@ -9,8 +9,7 @@ use std::time::Instant; use anyhow::{ensure, Context, Result}; use arrow_array::{Array, LargeBinaryArray, RecordBatch, StringArray, UInt64Array}; -use arrow_buffer::Buffer; -use arrow_ipc::reader::{read_footer_length, FileDecoder, FileReader}; +use arrow_ipc::reader::FileReader; use arrow_ipc::writer::FileWriter; use bytes::Bytes; use clap::{Parser, Subcommand}; @@ -131,21 +130,10 @@ fn encode(batch: &RecordBatch) -> Result> { Ok(data) } -fn decode(buffer: &Buffer, schema: arrow_schema::SchemaRef) -> Result { - let end = buffer.len().checked_sub(10).context("short IPC file")?; - let len = read_footer_length(buffer[end..].try_into()?)?; - let start = end.checked_sub(len).context("bad footer")?; - let footer = arrow_ipc::root_as_footer(&buffer[start..end]) - .map_err(|e| anyhow::anyhow!("invalid IPC footer: {e}"))?; - let decoder = FileDecoder::new(schema, footer.version()); - let blocks = footer.recordBatches().context("no batches")?; - ensure!(blocks.len() == 1, "one batch required"); - let block = blocks.get(0); - let length = block.bodyLength() as usize + block.metaDataLength() as usize; - let bytes = buffer.slice_with_length(block.offset() as usize, length); - decoder - .read_record_batch(block, &bytes)? - .context("missing batch") +fn decode(buffer: &Bytes, schema: arrow_schema::SchemaRef) -> Result { + let mut batches = lance_context_wal::ipc::decode(buffer.clone(), schema)?; + ensure!(batches.len() == 1, "one benchmark batch required"); + Ok(batches.remove(0)) } fn measured( @@ -178,15 +166,15 @@ fn measured( fn codec(size: usize, rows: usize, iterations: usize) -> Result<()> { ensure!(size > 0 && rows > 0 && iterations > 0); let batch = make_batch(rows, size)?; - let buffer = Buffer::from_vec(encode(&batch)?); + let buffer = Bytes::from(encode(&batch)?); verify_shared_decode(&buffer, &batch)?; - let reader_batch = FileReader::try_new(Cursor::new(buffer.as_slice()), None)? + let reader_batch = FileReader::try_new(Cursor::new(buffer.as_ref()), None)? .next() .context("missing reader batch")??; ensure!(reader_batch == batch, "reader round trip mismatch"); measured("encode_ipc", iterations, buffer.len(), || encode(&batch))?; measured("file_reader", iterations, buffer.len(), || { - Ok(FileReader::try_new(Cursor::new(buffer.as_slice()), None)? + Ok(FileReader::try_new(Cursor::new(buffer.as_ref()), None)? .next() .context("missing batch")??) })?; @@ -197,7 +185,7 @@ fn codec(size: usize, rows: usize, iterations: usize) -> Result<()> { let mut unaligned = Vec::with_capacity(buffer.len() + 1); unaligned.push(0); unaligned.extend_from_slice(&buffer); - let unaligned = Buffer::from_vec(unaligned).slice(1); + let unaligned = Bytes::from(unaligned).slice(1..); verify_shared_decode(&unaligned, &batch)?; measured( "unaligned_shared_buffer_decoder", @@ -226,7 +214,7 @@ fn codec(size: usize, rows: usize, iterations: usize) -> Result<()> { /// Check equality and actual payload ownership outside the measured interval. /// Numeric offset buffers can still allocate when fixing alignment. -fn verify_shared_decode(buffer: &Buffer, expected: &RecordBatch) -> Result<()> { +fn verify_shared_decode(buffer: &Bytes, expected: &RecordBatch) -> Result<()> { let decoded = decode(buffer, expected.schema())?; ensure!(&decoded == expected, "full Arrow round trip mismatch"); let bodies = decoded diff --git a/crates/lance-context-wal/src/ipc.rs b/crates/lance-context-wal/src/ipc.rs new file mode 100644 index 0000000..fa3c14f --- /dev/null +++ b/crates/lance-context-wal/src/ipc.rs @@ -0,0 +1,157 @@ +//! IPC file decoding for materializers over reference-counted input buffers. +//! Replication does not need this optional module or decode Arrow rows. +//! Compressed IPC is rejected until expansion is separately budgeted. +use anyhow::{ensure, Context, Result}; +use arrow_array::RecordBatch; +use arrow_buffer::Buffer; +use arrow_ipc::{ + reader::{read_footer_length, FileDecoder}, + Block, +}; +use arrow_schema::SchemaRef; +use bytes::Bytes; + +const MAX_METADATA: usize = 1024 * 1024; +const MAX_ROWS: i64 = 1_000_000; + +#[cfg(test)] +mod tests; + +/// Decode an uncompressed IPC file against the expected table schema. +/// +/// Binary values share the input's ownership; alignment repair and Arrow array +/// metadata can still allocate. Callers must bound/reserve input and decoding +/// workspace before calling. The format checks and row/count limits here do not +/// replace process-wide memory admission or a transport checksum. +/// Dictionary deltas and compressed batches are intentionally unsupported. +pub fn decode(body: Bytes, schema: SchemaRef) -> Result> { + let buffer = Buffer::from(body); + ensure!( + buffer.len() >= 16 && &buffer[..6] == b"ARROW1", + "not an IPC file" + ); + let end = buffer.len() - 10; + let footer_len = read_footer_length(buffer[end..].try_into()?)?; + ensure!(footer_len <= MAX_METADATA, "IPC footer too large"); + let start = end + .checked_sub(footer_len) + .context("invalid IPC footer range")?; + let footer = arrow_ipc::root_as_footer(&buffer[start..end]) + .map_err(|e| anyhow::anyhow!("invalid IPC footer: {e}"))?; + let source_schema = + arrow_ipc::convert::fb_to_schema(footer.schema().context("missing IPC schema")?); + ensure!( + source_schema == *schema, + "IPC schema differs from expected schema" + ); + let records = footer.recordBatches().context("missing record blocks")?; + ensure!( + !records.is_empty() && records.len() <= 128, + "invalid batch count" + ); + let dictionaries = footer.dictionaries(); + ensure!( + dictionaries.is_none_or(|d| d.len() <= 128), + "too many dictionary blocks" + ); + let mut blocks = Vec::new(); + if let Some(dictionaries) = dictionaries { + blocks.extend(dictionaries.iter().map(|b| (*b, true))); + } + blocks.extend(records.iter().map(|b| (*b, false))); + blocks.sort_by_key(|(b, _)| b.offset()); + let mut previous_end = 8; + let mut rows = 0; + // Validate the complete block layout and compression before decoding any + // arrays. Disallow overlapping blocks and bound logical rows, including + // null/dictionary values that could otherwise expand from tiny messages. + for (block, dictionary) in &blocks { + let offset = usize::try_from(block.offset())?; + let length = usize::try_from(block.metaDataLength())? + .checked_add(usize::try_from(block.bodyLength())?) + .context("IPC block overflow")?; + let block_end = offset.checked_add(length).context("IPC offset overflow")?; + ensure!( + offset >= previous_end && block_end <= start, + "overlapping/out-of-range IPC blocks" + ); + previous_end = block_end; + rows += validate_message(block, &buffer[offset..block_end], *dictionary)?; + ensure!(rows <= MAX_ROWS, "IPC row limit exceeded"); + } + let mut decoder = FileDecoder::new(schema, footer.version()); + // Match FileReader: load dictionary blocks before records, preserving + // dictionary replacement/delta validation in Arrow's decoder. + for (block, _) in blocks.iter().filter(|(_, dictionary)| *dictionary) { + decoder.read_dictionary(block, &slice(&buffer, block))?; + } + let mut output = Vec::new(); + for (block, _) in blocks.iter().filter(|(_, dictionary)| !dictionary) { + output.push( + decoder + .read_record_batch(block, &slice(&buffer, block))? + .context("missing IPC batch")?, + ); + } + Ok(output) +} + +fn slice(buffer: &Buffer, block: &Block) -> Buffer { + buffer.slice_with_length( + block.offset() as usize, + block.metaDataLength() as usize + block.bodyLength() as usize, + ) +} + +fn validate_message(block: &Block, bytes: &[u8], dictionary: bool) -> Result { + let metadata_len = usize::try_from(block.metaDataLength())?; + ensure!( + (8..=MAX_METADATA).contains(&metadata_len), + "invalid IPC metadata size" + ); + let prefix = if bytes.get(..4) == Some(&[255; 4]) { + 8 + } else { + 4 + }; + let metadata = bytes + .get(prefix..metadata_len) + .context("truncated IPC metadata")?; + let message = arrow_ipc::root_as_message(metadata) + .map_err(|e| anyhow::anyhow!("invalid IPC message: {e}"))?; + ensure!( + message.bodyLength() == block.bodyLength(), + "IPC body length mismatch" + ); + let batch = if dictionary { + let dict = message + .header_as_dictionary_batch() + .context("invalid dictionary block")?; + ensure!( + !dict.isDelta(), + "dictionary deltas require separate expansion accounting" + ); + dict.data().context("missing dictionary data")? + } else { + message + .header_as_record_batch() + .context("invalid record block")? + }; + ensure!( + batch.compression().is_none(), + "compressed IPC requires expansion accounting" + ); + ensure!( + (0..=MAX_ROWS).contains(&batch.length()), + "invalid IPC row count" + ); + if let Some(nodes) = batch.nodes() { + ensure!( + nodes + .iter() + .all(|node| (0..=MAX_ROWS).contains(&node.length())), + "nested IPC row limit exceeded" + ); + } + Ok(batch.length()) +} diff --git a/crates/lance-context-wal/src/ipc/tests.rs b/crates/lance-context-wal/src/ipc/tests.rs new file mode 100644 index 0000000..f9f62fc --- /dev/null +++ b/crates/lance-context-wal/src/ipc/tests.rs @@ -0,0 +1,116 @@ +use super::*; +use arrow_array::{ + types::Int8Type, ArrayRef, DictionaryArray, Int8Array, LargeBinaryArray, NullArray, StringArray, +}; +use arrow_ipc::writer::{FileWriter, IpcWriteOptions}; +use arrow_schema::Schema; +use std::sync::Arc; + +fn fixture() -> RecordBatch { + let payload = vec![73u8; 256 * 1024]; + RecordBatch::try_from_iter([ + ( + "payload", + Arc::new(LargeBinaryArray::from(vec![ + Some(payload.as_slice()), + None, + Some(&[]), + ])) as ArrayRef, + ), + ( + "role", + Arc::new( + DictionaryArray::::try_new( + Int8Array::from(vec![Some(0), None, Some(1)]), + Arc::new(StringArray::from(vec!["request", "response"])), + ) + .unwrap(), + ) as ArrayRef, + ), + ]) + .unwrap() +} + +fn encode(batches: &[RecordBatch], compressed: bool) -> Result { + let mut bytes = Vec::new(); + let options = IpcWriteOptions::default() + .try_with_compression(compressed.then_some(arrow_ipc::CompressionType::ZSTD))?; + let mut writer = FileWriter::try_new_with_options(&mut bytes, &batches[0].schema(), options)?; + for batch in batches { + writer.write(batch)?; + } + writer.finish()?; + drop(writer); + Ok(bytes.into()) +} + +#[test] +fn dictionaries_nulls_and_unaligned_binary_keep_input_ownership() -> Result<()> { + let batch = fixture(); + let encoded = encode(&[batch.clone(), batch.clone()], false)?; + for prefix in [0, 1] { + let mut bytes = vec![0; prefix]; + bytes.extend_from_slice(&encoded); + let bytes = Bytes::from(bytes).slice(prefix..); + let start = bytes.as_ptr() as usize; + let end = start + bytes.len(); + let decoded = decode(bytes, batch.schema())?; + // The caller's Bytes has been consumed. Arrow owns the backing storage. + assert_eq!(decoded, vec![batch.clone(), batch.clone()]); + for batch in &decoded { + let binary = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let pointer = binary.value_data().as_ptr() as usize; + assert!(pointer >= start && pointer + binary.value_data().len() <= end); + } + } + Ok(()) +} + +#[test] +fn rejects_truncation_wrong_schema_and_unsupported_expansion() -> Result<()> { + let batch = fixture(); + let bytes = encode(std::slice::from_ref(&batch), false)?; + for len in [0, 5, 15, bytes.len() / 2, bytes.len() - 8, bytes.len() - 1] { + assert!(decode(bytes.slice(..len), batch.schema()).is_err()); + } + assert!(decode(bytes, Arc::new(Schema::empty())).is_err()); + let compressed = encode(std::slice::from_ref(&batch), true)?; + assert!(decode(compressed, batch.schema()) + .unwrap_err() + .to_string() + .contains("compressed IPC")); + // Null rows encode almost no payload; row limits must apply before decoding. + let huge = RecordBatch::try_from_iter([( + "nulls", + Arc::new(NullArray::new(MAX_ROWS as usize + 1)) as ArrayRef, + )])?; + assert!(decode(encode(std::slice::from_ref(&huge), false)?, huge.schema()).is_err()); + assert!(decode(encode(&vec![batch.clone(); 129], false)?, batch.schema()).is_err()); + Ok(()) +} + +#[test] +fn rejects_invalid_footer_and_block_ranges_before_slicing() -> Result<()> { + let batch = fixture(); + let bytes = encode(std::slice::from_ref(&batch), false)?; + let end = bytes.len() - 10; + let footer_len = read_footer_length(bytes[end..].try_into()?)?; + let block_offset = { + let footer = arrow_ipc::root_as_footer(&bytes[end - footer_len..end]).unwrap(); + let block = footer.recordBatches().unwrap().get(0); + block as *const Block as usize - bytes.as_ptr() as usize + }; + for offset in [1i64, -1, i64::MAX] { + let mut corrupt = bytes.to_vec(); + corrupt[block_offset..block_offset + 8].copy_from_slice(&offset.to_le_bytes()); + assert!(decode(corrupt.into(), batch.schema()).is_err()); + } + let mut corrupt = bytes.to_vec(); + corrupt[end..end + 4].copy_from_slice(&((MAX_METADATA + 1) as i32).to_le_bytes()); + assert!(decode(corrupt.into(), batch.schema()).is_err()); + Ok(()) +} diff --git a/crates/lance-context-wal/src/lib.rs b/crates/lance-context-wal/src/lib.rs index 5ac9820..7ea0eee 100644 --- a/crates/lance-context-wal/src/lib.rs +++ b/crates/lance-context-wal/src/lib.rs @@ -10,3 +10,6 @@ pub mod placement; #[cfg(feature = "admission")] pub mod admission; + +#[cfg(feature = "ipc")] +pub mod ipc; From d8b03014d583ae8d1963bea16c13796478fc6f5d Mon Sep 17 00:00:00 2001 From: Beinan Date: Mon, 19 Oct 2026 05:12:20 +0000 Subject: [PATCH 2/2] fix: reject malformed IPC metadata without unwinding callers --- crates/lance-context-wal/README.md | 7 +- crates/lance-context-wal/src/ipc.rs | 30 +++++++- crates/lance-context-wal/src/ipc/tests.rs | 89 +++++++++++++++++++++++ 3 files changed, 122 insertions(+), 4 deletions(-) diff --git a/crates/lance-context-wal/README.md b/crates/lance-context-wal/README.md index aedd41a..c6621cd 100644 --- a/crates/lance-context-wal/README.md +++ b/crates/lance-context-wal/README.md @@ -7,11 +7,14 @@ The optional `ipc` feature exposes `ipc::decode(Bytes, expected_schema)` for materializers. It decodes multiple uncompressed record batches and dictionaries while sharing binary payload ownership with the input. Numeric alignment repair, schema parsing and Arrow metadata still allocate. The codec validates footer and -block ranges, schema equality, row/count limits and unsupported expansion before -constructing arrays. Compressed batches and dictionary deltas are rejected. +block/array-buffer ranges, native endianness, schema equality, row/count limits +and unsupported expansion before constructing arrays. Compressed batches and +dictionary deltas are rejected. Callers still reserve input and decoding workspace within their own memory budget. Replication does not need to enable this module or decode Arrow. The `lab` codec measurement now uses the same validated path, including its validation overhead. +Arrow dependency panics from malformed metadata become errors in panic-unwind +builds; panic-abort builds cannot recover from dependency panics. ```bash cargo test -p lance-context-wal --features ipc --lib diff --git a/crates/lance-context-wal/src/ipc.rs b/crates/lance-context-wal/src/ipc.rs index fa3c14f..6e89c1c 100644 --- a/crates/lance-context-wal/src/ipc.rs +++ b/crates/lance-context-wal/src/ipc.rs @@ -24,7 +24,17 @@ mod tests; /// workspace before calling. The format checks and row/count limits here do not /// replace process-wide memory admission or a transport checksum. /// Dictionary deltas and compressed batches are intentionally unsupported. +/// Arrow dependency panics are converted to errors in panic-unwind builds; +/// panic-abort builds cannot recover from dependency panics. pub fn decode(body: Bytes, schema: SchemaRef) -> Result> { + // Arrow's schema conversion is infallible in its API, but malformed type + // metadata can panic even after FlatBuffers verification. All decoder state + // is local to this attempt and is discarded on unwind, never reused. + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| decode_inner(body, schema))) + .map_err(|_| anyhow::anyhow!("Arrow IPC decoder panicked on invalid metadata"))? +} + +fn decode_inner(body: Bytes, schema: SchemaRef) -> Result> { let buffer = Buffer::from(body); ensure!( buffer.len() >= 16 && &buffer[..6] == b"ARROW1", @@ -38,8 +48,12 @@ pub fn decode(body: Bytes, schema: SchemaRef) -> Result> { .context("invalid IPC footer range")?; let footer = arrow_ipc::root_as_footer(&buffer[start..end]) .map_err(|e| anyhow::anyhow!("invalid IPC footer: {e}"))?; - let source_schema = - arrow_ipc::convert::fb_to_schema(footer.schema().context("missing IPC schema")?); + let ipc_schema = footer.schema().context("missing IPC schema")?; + ensure!( + ipc_schema.endianness().equals_to_target_endianness(), + "IPC endianness differs from this system" + ); + let source_schema = arrow_ipc::convert::fb_to_schema(ipc_schema); ensure!( source_schema == *schema, "IPC schema differs from expected schema" @@ -153,5 +167,17 @@ fn validate_message(block: &Block, bytes: &[u8], dictionary: bool) -> Result= 0 && buffer.length() >= 0 && end <= block.bodyLength(), + "IPC array buffer outside block body" + ); + } + } Ok(batch.length()) } diff --git a/crates/lance-context-wal/src/ipc/tests.rs b/crates/lance-context-wal/src/ipc/tests.rs index f9f62fc..085d0c3 100644 --- a/crates/lance-context-wal/src/ipc/tests.rs +++ b/crates/lance-context-wal/src/ipc/tests.rs @@ -114,3 +114,92 @@ fn rejects_invalid_footer_and_block_ranges_before_slicing() -> Result<()> { assert!(decode(corrupt.into(), batch.schema()).is_err()); Ok(()) } + +#[test] +fn malformed_dictionary_schema_returns_an_error() -> Result<()> { + let batch = fixture(); + let bytes = encode(std::slice::from_ref(&batch), false)?; + let end = bytes.len() - 10; + let footer_len = read_footer_length(bytes[end..].try_into()?)?; + let footer_start = end - footer_len; + let bit_width = { + let footer = arrow_ipc::root_as_footer(&bytes[footer_start..end]).unwrap(); + let field = footer.schema().unwrap().fields().unwrap().get(1); + let index = field.dictionary().unwrap().indexType().unwrap(); + footer_start + + index._tab.loc() + + usize::from(index._tab.vtable().get(arrow_ipc::Int::VT_BITWIDTH)) + }; + let mut corrupt = bytes.to_vec(); + // Structurally valid FlatBuffers can still contain invalid Arrow types. + corrupt[bit_width..bit_width + 4].copy_from_slice(&7i32.to_le_bytes()); + assert!(decode(corrupt.into(), batch.schema()).is_err()); + Ok(()) +} + +#[test] +fn rejects_array_buffers_outside_the_record_body() -> Result<()> { + let batch = fixture(); + let bytes = encode(std::slice::from_ref(&batch), false)?; + let end = bytes.len() - 10; + let footer_len = read_footer_length(bytes[end..].try_into()?)?; + let position = { + let footer = arrow_ipc::root_as_footer(&bytes[end - footer_len..end]).unwrap(); + let block = footer.recordBatches().unwrap().get(0); + let start = block.offset() as usize; + let message = + arrow_ipc::root_as_message(&bytes[start + 8..start + block.metaDataLength() as usize]) + .unwrap(); + let buffer = message + .header_as_record_batch() + .unwrap() + .buffers() + .unwrap() + .get(0); + buffer as *const arrow_ipc::Buffer as usize - bytes.as_ptr() as usize + }; + for (offset, length) in [(-1i64, 1i64), (0, -1), (i64::MAX, 1), (0, i64::MAX)] { + let mut corrupt = bytes.to_vec(); + corrupt[position..position + 8].copy_from_slice(&offset.to_le_bytes()); + corrupt[position + 8..position + 16].copy_from_slice(&length.to_le_bytes()); + let error = decode(corrupt.into(), batch.schema()).unwrap_err(); + assert!(error.to_string().starts_with("IPC array buffer")); + } + Ok(()) +} + +#[test] +fn rejects_foreign_endianness_before_schema_conversion() { + let schema = Arc::new(Schema::empty()); + let mut builder = arrow_ipc::convert::IpcSchemaEncoder::new().schema_to_fb(&schema); + builder.reset(); + let endianness = if cfg!(target_endian = "little") { + arrow_ipc::Endianness::Big + } else { + arrow_ipc::Endianness::Little + }; + let ipc_schema = arrow_ipc::Schema::create( + &mut builder, + &arrow_ipc::SchemaArgs { + endianness, + ..Default::default() + }, + ); + let footer = arrow_ipc::Footer::create( + &mut builder, + &arrow_ipc::FooterArgs { + schema: Some(ipc_schema), + ..Default::default() + }, + ); + builder.finish(footer, None); + let footer = builder.finished_data(); + let mut bytes = b"ARROW1\0\0".to_vec(); + bytes.extend_from_slice(footer); + bytes.extend_from_slice(&(footer.len() as i32).to_le_bytes()); + bytes.extend_from_slice(b"ARROW1"); + assert!(decode(bytes.into(), schema) + .unwrap_err() + .to_string() + .contains("endianness")); +}