Skip to main content

vgi_rpc/
wire.rs

1//! Low-level IPC stream helpers that preserve per-batch custom metadata.
2//!
3//! The standard `arrow-ipc` `StreamWriter` / `StreamReader` types do not
4//! expose per-message `custom_metadata`, but the vgi_rpc wire protocol
5//! relies on that field to carry `vgi_rpc.method`,
6//! `vgi_rpc.request_version`, log keys, externalisation pointers,
7//! state tokens, etc. This module hand-rolls the framing layer so the
8//! crate can depend on **stock** arrow-rs from crates.io rather than a
9//! patched fork — the published vgi-rpc crate is therefore directly
10//! installable without any `[patch.crates-io]` directives downstream.
11//!
12//! Internally we delegate column encoding / decoding to
13//! [`arrow_ipc::writer::IpcDataGenerator`] and the
14//! [`arrow_ipc::reader::read_record_batch`] / `read_dictionary`
15//! functions, and only intercept the flatbuffer `Message` wrapper to
16//! attach / extract `custom_metadata`. That keeps the code small and
17//! the on-wire bytes byte-for-byte compatible with arrow-rs's
18//! `StreamWriter`.
19//!
20//! ## DoS guard
21//!
22//! [`StreamReader::new`] pre-validates the schema-message length prefix
23//! against [`MAX_IPC_SCHEMA_BYTES`] *before* allocating; a remote
24//! client cannot trigger a multi-gigabyte alloc by sending a crafted
25//! 4-byte payload. A per-batch message body is bounded twice: an absurd
26//! `bodyLength` is refused outright against [`MAX_IPC_MESSAGE_BYTES`],
27//! and what survives that is buffered as the peer actually delivers it
28//! rather than on the strength of the claim — so the flatbuffer
29//! overshoot the fuzz harness surfaced costs a few MiB and an EOF.
30
31use std::collections::HashMap;
32use std::io::{Read, Write};
33use std::sync::Arc;
34
35use arrow_array::{Array, LargeBinaryArray, RecordBatch};
36use arrow_buffer::Buffer as ArrowBuffer;
37use arrow_ipc::reader as ipc_reader;
38use arrow_ipc::writer::{write_message, DictionaryTracker, IpcDataGenerator, IpcWriteOptions};
39use arrow_ipc::{convert as ipc_convert, root_as_message, MessageHeader};
40use arrow_schema::{DataType, Schema, SchemaRef};
41use flatbuffers::FlatBufferBuilder;
42
43use crate::errors::{Result, RpcError};
44
45/// Per-batch metadata pairs. Order is not preserved across
46/// serialisation; that matches Python's `RecordBatch.custom_metadata`
47/// semantics.
48pub type Metadata = HashMap<String, String>;
49
50/// Internal marker inserted when a fully framed record-batch message carried
51/// invalid UTF-8 in custom metadata. The reader sanitizes only its private
52/// header copy so it can preserve the next stream boundary; request validation
53/// rejects this marker before dispatch.
54pub(crate) const INVALID_UTF8_METADATA_KEY: &str = "\0vgi_rpc.invalid_utf8_metadata";
55
56/// Look up a key in a [`Metadata`] map, returning the value as `&str`.
57#[inline]
58pub fn md_get<'a>(md: &'a Metadata, key: &str) -> Option<&'a str> {
59    md.get(key).map(String::as_str)
60}
61
62/// Maximum permitted size, in bytes, of the schema-message flatbuffer
63/// at the head of an IPC stream. Schemas are typically tens to
64/// hundreds of bytes; 16 MiB is gracious headroom that still rejects
65/// the crafted 4-byte input `[0x1A, 0x2C, 0xF5, 0x2C]` that
66/// `fuzz/wire_stream_reader` discovered would OOM the process by
67/// claiming a ~720 MB schema. Applies to the *schema* message length
68/// prefix on the wire.
69pub const MAX_IPC_SCHEMA_BYTES: usize = 16 * 1024 * 1024;
70
71/// Maximum permitted total size of any per-batch IPC message (header
72/// flatbuffer + body bytes) — the sanity ceiling that refuses the
73/// `bodyLength = 0x4000000100000` overshoot the fuzz harness surfaced.
74///
75/// This used to be 256 MiB, which also made it a hard limit on
76/// *legitimate* payloads: a >2 GiB `large_binary` round-trip is well
77/// within what the Python reference accepts, and the
78/// `large_payload.echo_binary_over_int32_max` conformance test sends
79/// exactly that. Refusing it was a conformance defect, not a defence.
80///
81/// The ceiling no longer carries the anti-OOM job on its own —
82/// `read_message_bytes` grows the body buffer from the bytes that
83/// actually arrive (see `BODY_PREALLOC_LIMIT`), so a crafted length
84/// costs a few MiB and an EOF rather than the amount it claimed.
85/// `u32::MAX` keeps the constant expressible on 32-bit targets, where
86/// it saturates to `usize::MAX` and the allocation guard is the only
87/// one that can meaningfully apply anyway.
88pub const MAX_IPC_MESSAGE_BYTES: usize = u32::MAX as usize;
89
90/// Bytes reserved up front for a message body before any of it has
91/// arrived. Beyond this the buffer grows amortised as the peer actually
92/// delivers, so a header claiming a petabyte cannot turn a 4-byte frame
93/// into a multi-gigabyte allocation.
94///
95/// Growing rather than pre-sizing also keeps the *read* side out of the
96/// trouble [`ChunkedWriter`] fixes on the write side: `impl Read for
97/// &UnixStream` calls `recv(2)` with the length unclamped, exactly as
98/// its `Write` counterpart calls `send(2)`, so handing it a single spare
99/// region past 2 GiB would earn the same `EINVAL`. Doubling from here
100/// means the largest slice ever offered is about half the body — a
101/// 1 GiB read for a 2 GiB message — and never reaches `INT_MAX`.
102const BODY_PREALLOC_LIMIT: usize = 8 * 1024 * 1024;
103
104// ---------------------------------------------------------------------------
105// Writer
106// ---------------------------------------------------------------------------
107
108const CONTINUATION_MARKER: [u8; 4] = [0xFF, 0xFF, 0xFF, 0xFF];
109
110/// Largest slice offered to a single underlying `write` call.
111///
112/// Sits well under `INT_MAX`, which is where the two macOS failure modes
113/// live (see [`ChunkedWriter`]). Slicing is free, so the whole cost of
114/// the clamp is one extra syscall per gigabyte.
115const MAX_WRITE_CHUNK: usize = 1 << 30; // 1 GiB
116
117/// arrow-rs's default IPC alignment. `StreamWriter` always constructs default
118/// write options, so the direct large-binary path below uses the same layout.
119const IPC_ALIGNMENT: usize = 64;
120const IPC_PADDING: [u8; IPC_ALIGNMENT] = [0; IPC_ALIGNMENT];
121
122/// Clamp every `write` to [`MAX_WRITE_CHUNK`] so a large payload
123/// survives the syscall underneath.
124///
125/// Every raw transport writer here is unbuffered on purpose, so one
126/// `write` maps onto one `write(2)` / `send(2)`. That syscall is not
127/// obliged to accept the whole buffer, and above 2 GiB on macOS it
128/// refuses to — in one of two different ways depending on what is
129/// underneath:
130///
131/// * **pipes** return a short count of exactly `INT_MAX` with *no
132///   error*, so a writer that trusts the return value silently drops
133///   the tail and the peer blocks forever waiting for bytes the Arrow
134///   IPC header promised. The symptom is a deadlock, not an exception.
135/// * **sockets** (Unix domain and TCP) fail outright with `EINVAL`.
136///
137/// Both halves are needed. `io::Write::write_all` already loops on the
138/// returned count, and `std`'s file-descriptor writer clamps to
139/// `INT_MAX` for us — but `impl Write for &UnixStream` and `&TcpStream`
140/// do *not* go through it. They call `send(2)` with
141/// `cmp::min(buf.len(), wrlen_t::MAX)`, and `wrlen_t` is `usize` on
142/// unix, so the length reaches the kernel unclamped and a >2 GiB Arrow
143/// IPC body dies with `EINVAL`. Clamping here is what makes the socket
144/// transports behave like the pipe ones.
145///
146/// Deliberately does *not* forward `write_vectored`: the default
147/// implementation routes back through `write`, which is where the clamp
148/// lives.
149struct ChunkedWriter<W: Write> {
150    inner: W,
151    limit: usize,
152}
153
154impl<W: Write> ChunkedWriter<W> {
155    fn new(inner: W) -> Self {
156        Self {
157            inner,
158            limit: MAX_WRITE_CHUNK,
159        }
160    }
161
162    /// Same behaviour with a smaller clamp, so the chunking can be
163    /// exercised without allocating a gigabyte.
164    #[cfg(test)]
165    fn with_limit(inner: W, limit: usize) -> Self {
166        Self { inner, limit }
167    }
168
169    fn get_mut(&mut self) -> &mut W {
170        &mut self.inner
171    }
172}
173
174impl<W: Write> Write for ChunkedWriter<W> {
175    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
176        let end = buf.len().min(self.limit);
177        self.inner.write(&buf[..end])
178    }
179
180    fn flush(&mut self) -> std::io::Result<()> {
181        self.inner.flush()
182    }
183}
184
185/// A streaming IPC writer that supports per-batch custom metadata.
186///
187/// The byte sequence written for a complete stream is:
188///   `SchemaMessage → [DictionaryMessage]* → [RecordBatchMessage]* → EOS(4xFF 0x00)`.
189///
190/// Each call to [`write`](Self::write) emits one record-batch message
191/// (preceded by any newly-needed dictionary messages) with its
192/// `custom_metadata` attached at the IPC Message level.
193///
194/// Every byte the crate puts on a transport passes through here, which
195/// is why the `ChunkedWriter` clamp lives at this level rather than in
196/// each transport: `serve` takes an arbitrary `W`, so a worker that
197/// hands it a raw socket gets the same protection as `serve_unix` does.
198pub struct StreamWriter<W: Write> {
199    writer: ChunkedWriter<W>,
200    schema: SchemaRef,
201    opts: IpcWriteOptions,
202    data_gen: IpcDataGenerator,
203    dict_tracker: DictionaryTracker,
204    finished: bool,
205    /// Reused across `write` calls so the per-batch metadata repack doesn't
206    /// allocate a fresh flatbuffer builder (and its internal vectors) each
207    /// time. `reset()` before each use; the buffer's capacity is retained.
208    fbb: FlatBufferBuilder<'static>,
209}
210
211impl<W: Write> StreamWriter<W> {
212    /// Create a new writer and emit the schema message.
213    pub fn new(writer: W, schema: &Schema) -> Result<Self> {
214        let mut writer = ChunkedWriter::new(writer);
215        let opts = IpcWriteOptions::default();
216        let data_gen = IpcDataGenerator::default();
217        let mut dict_tracker = DictionaryTracker::new(false);
218        let encoded =
219            data_gen.schema_to_bytes_with_dictionary_tracker(schema, &mut dict_tracker, &opts);
220        write_message(&mut writer, encoded, &opts)?;
221        Ok(Self {
222            writer,
223            schema: Arc::new(schema.clone()),
224            opts,
225            data_gen,
226            dict_tracker,
227            finished: false,
228            fbb: FlatBufferBuilder::new(),
229        })
230    }
231
232    /// Write one RecordBatch carrying optional `metadata` as the IPC
233    /// Message-level `custom_metadata` field. Pass `None` to omit the
234    /// field (saves a few bytes per message).
235    pub fn write(&mut self, batch: &RecordBatch, metadata: Option<&Metadata>) -> Result<()> {
236        if self.finished {
237            return Err(RpcError::new("IOError", "writer already finished"));
238        }
239        if write_single_large_binary_direct(
240            &mut self.writer,
241            &mut self.fbb,
242            batch,
243            metadata,
244            &self.opts,
245        )? {
246            return Ok(());
247        }
248        let mut ctx = Default::default();
249        let (dicts, data) = self
250            .data_gen
251            .encode(batch, &mut self.dict_tracker, &self.opts, &mut ctx)
252            .map_err(RpcError::from)?;
253        for d in dicts {
254            write_message(&mut self.writer, d, &self.opts).map_err(RpcError::from)?;
255        }
256        if let Some(md) = metadata.filter(|m| !m.is_empty()) {
257            self.fbb.reset();
258            repack_record_batch_message_with_metadata(&mut self.fbb, &data.ipc_message, md)?;
259            let encoded = arrow_ipc::writer::EncodedData {
260                ipc_message: self.fbb.finished_data().to_vec(),
261                arrow_data: data.arrow_data,
262            };
263            write_message(&mut self.writer, encoded, &self.opts).map_err(RpcError::from)?;
264        } else {
265            write_message(&mut self.writer, data, &self.opts).map_err(RpcError::from)?;
266        }
267        Ok(())
268    }
269
270    /// Return the schema this writer was opened with.
271    pub fn schema(&self) -> SchemaRef {
272        self.schema.clone()
273    }
274
275    /// Write the EOS continuation marker. Idempotent.
276    pub fn finish(&mut self) -> Result<()> {
277        if self.finished {
278            return Ok(());
279        }
280        self.writer.write_all(&CONTINUATION_MARKER)?;
281        self.writer.write_all(&[0u8; 4])?;
282        self.writer.flush()?;
283        self.finished = true;
284        Ok(())
285    }
286
287    /// Flush the underlying writer.
288    pub fn flush(&mut self) -> Result<()> {
289        self.writer.flush()?;
290        Ok(())
291    }
292
293    pub fn get_mut(&mut self) -> &mut W {
294        self.writer.get_mut()
295    }
296}
297
298/// Write the hot-path shape used by large binary unary responses without
299/// copying the value buffer into `EncodedData::arrow_data` first.
300///
301/// `IpcDataGenerator` deliberately produces one contiguous body `Vec`, which
302/// is useful to callers that need encoded bytes but adds a payload-sized copy
303/// when the final destination is already a pipe or socket. A non-null,
304/// one-row, one-column LargeBinary batch has a fixed Arrow IPC layout, so its
305/// small validity and offsets buffers can be framed here while the value
306/// buffer is written directly to the transport.
307fn write_single_large_binary_direct<W: Write>(
308    writer: &mut ChunkedWriter<W>,
309    fbb: &mut FlatBufferBuilder<'static>,
310    batch: &RecordBatch,
311    metadata: Option<&Metadata>,
312    opts: &IpcWriteOptions,
313) -> Result<bool> {
314    use arrow_ipc::{
315        Buffer as IpcBuffer, FieldNode, KeyValue, KeyValueArgs, MessageBuilder, MetadataVersion,
316        RecordBatchBuilder,
317    };
318
319    if batch.num_rows() != 1
320        || batch.num_columns() != 1
321        || batch.schema().field(0).data_type() != &DataType::LargeBinary
322    {
323        return Ok(false);
324    }
325    let Some(array) = batch.column(0).as_any().downcast_ref::<LargeBinaryArray>() else {
326        return Ok(false);
327    };
328    if array.is_null(0) {
329        return Ok(false);
330    }
331
332    let value = array.value(0);
333    let value_len = i64::try_from(value.len())
334        .map_err(|_| RpcError::value_error("large_binary value exceeds i64 offsets"))?;
335    let value_padding = padding_to_alignment(value.len(), IPC_ALIGNMENT);
336    let body_len = 2usize
337        .checked_mul(IPC_ALIGNMENT)
338        .and_then(|prefix| prefix.checked_add(value.len()))
339        .and_then(|length| length.checked_add(value_padding))
340        .ok_or_else(|| RpcError::value_error("large_binary IPC body length overflow"))?;
341    let body_len_i64 = i64::try_from(body_len)
342        .map_err(|_| RpcError::value_error("large_binary IPC body exceeds i64 length"))?;
343
344    fbb.reset();
345    let nodes = fbb.create_vector(&[FieldNode::new(1, 0)]);
346    let buffers = fbb.create_vector(&[
347        IpcBuffer::new(0, 1),
348        IpcBuffer::new(IPC_ALIGNMENT as i64, 16),
349        IpcBuffer::new((2 * IPC_ALIGNMENT) as i64, value_len),
350    ]);
351    let record_batch = {
352        let mut builder = RecordBatchBuilder::new(fbb);
353        builder.add_length(1);
354        builder.add_nodes(nodes);
355        builder.add_buffers(buffers);
356        builder.finish()
357    };
358
359    let custom_metadata = metadata
360        .filter(|entries| !entries.is_empty())
361        .map(|entries| {
362            let entries: Vec<_> = entries
363                .iter()
364                .map(|(key, value)| {
365                    let key = fbb.create_string(key);
366                    let value = fbb.create_string(value);
367                    KeyValue::create(
368                        fbb,
369                        &KeyValueArgs {
370                            key: Some(key),
371                            value: Some(value),
372                        },
373                    )
374                })
375                .collect();
376            fbb.create_vector(&entries)
377        });
378
379    let message = {
380        let mut builder = MessageBuilder::new(fbb);
381        builder.add_version(MetadataVersion::V5);
382        builder.add_header_type(MessageHeader::RecordBatch);
383        builder.add_header(record_batch.as_union_value());
384        builder.add_bodyLength(body_len_i64);
385        if let Some(custom_metadata) = custom_metadata {
386            builder.add_custom_metadata(custom_metadata);
387        }
388        builder.finish()
389    };
390    fbb.finish(message, None);
391    let encoded = arrow_ipc::writer::EncodedData {
392        ipc_message: fbb.finished_data().to_vec(),
393        arrow_data: Vec::new(),
394    };
395    write_message(&mut *writer, encoded, opts).map_err(RpcError::from)?;
396
397    // Validity bitmap (one valid row), then its 64-byte alignment padding.
398    writer.write_all(&[0xff])?;
399    writer.write_all(&IPC_PADDING[..IPC_ALIGNMENT - 1])?;
400
401    // Rebased LargeBinary offsets [0, value_len], then alignment padding.
402    writer.write_all(&0_i64.to_le_bytes())?;
403    writer.write_all(&value_len.to_le_bytes())?;
404    writer.write_all(&IPC_PADDING[..IPC_ALIGNMENT - 16])?;
405
406    // The payload itself is never copied into an intermediate Vec.
407    writer.write_all(value)?;
408    writer.write_all(&IPC_PADDING[..value_padding])?;
409    Ok(true)
410}
411
412#[inline]
413fn padding_to_alignment(length: usize, alignment: usize) -> usize {
414    (alignment - (length % alignment)) % alignment
415}
416
417impl<W: Write> Drop for StreamWriter<W> {
418    fn drop(&mut self) {
419        let _ = self.finish();
420    }
421}
422
423/// Rebuild a Message flatbuffer with `custom_metadata` added,
424/// preserving the embedded RecordBatch header unchanged.
425fn repack_record_batch_message_with_metadata(
426    fbb: &mut FlatBufferBuilder<'static>,
427    msg_bytes: &[u8],
428    metadata: &Metadata,
429) -> Result<()> {
430    use arrow_ipc::{KeyValue, KeyValueArgs, MessageBuilder, RecordBatchBuilder};
431
432    let msg = root_as_message(msg_bytes)
433        .map_err(|e| RpcError::new("IPC", format!("parsing message: {e}")))?;
434    let version = msg.version();
435    let header_type = msg.header_type();
436    let body_length = msg.bodyLength();
437    if header_type != MessageHeader::RecordBatch {
438        return Err(RpcError::new(
439            "IPC",
440            format!("repack expected RecordBatch header, got {header_type:?}"),
441        ));
442    }
443    let rb = msg
444        .header_as_record_batch()
445        .ok_or_else(|| RpcError::new("IPC", "missing RecordBatch header"))?;
446
447    // The caller has already `reset()` the builder. Feed the field-node and
448    // buffer descriptors straight from the source flatbuffer vectors via
449    // `create_vector_from_iter` — no throwaway `Vec` per batch.
450    let src_nodes = rb
451        .nodes()
452        .ok_or_else(|| RpcError::new("IPC", "RecordBatch missing nodes"))?;
453    let nodes_vec = fbb.create_vector_from_iter(src_nodes.iter());
454
455    let src_buffers = rb
456        .buffers()
457        .ok_or_else(|| RpcError::new("IPC", "RecordBatch missing buffers"))?;
458    let buffers_vec = fbb.create_vector_from_iter(src_buffers.iter());
459
460    let variadic_vec = rb
461        .variadicBufferCounts()
462        .map(|v| fbb.create_vector_from_iter(v.iter()));
463
464    let new_rb = {
465        let mut b = RecordBatchBuilder::new(fbb);
466        b.add_length(rb.length());
467        b.add_nodes(nodes_vec);
468        b.add_buffers(buffers_vec);
469        if let Some(v) = variadic_vec {
470            b.add_variadicBufferCounts(v);
471        }
472        // Note: we don't carry compression here; the conformance worker
473        // does not enable IPC batch compression, so this is safe.
474        b.finish()
475    };
476
477    // Build custom_metadata vector. Order matches HashMap iteration —
478    // not stable, but that matches both upstream arrow-ipc and Python
479    // `RecordBatch.custom_metadata` semantics.
480    let kvs: Vec<_> = metadata
481        .iter()
482        .map(|(k, v)| {
483            let k_off = fbb.create_string(k);
484            let v_off = fbb.create_string(v);
485            KeyValue::create(
486                fbb,
487                &KeyValueArgs {
488                    key: Some(k_off),
489                    value: Some(v_off),
490                },
491            )
492        })
493        .collect();
494    let md_vec = fbb.create_vector(&kvs);
495
496    let mut mb = MessageBuilder::new(fbb);
497    mb.add_version(version);
498    mb.add_header_type(header_type);
499    mb.add_header(new_rb.as_union_value());
500    mb.add_bodyLength(body_length);
501    mb.add_custom_metadata(md_vec);
502    let m = mb.finish();
503    fbb.finish(m, None);
504    Ok(())
505}
506
507// ---------------------------------------------------------------------------
508// Reader
509// ---------------------------------------------------------------------------
510
511/// A streaming IPC reader that surfaces per-message custom metadata.
512///
513/// [`read_next`](Self::read_next) returns `Some((batch, metadata))`
514/// for each RecordBatch message and `None` on end-of-stream.
515/// Dictionary and schema messages are consumed transparently.
516pub struct StreamReader<R: Read> {
517    reader: R,
518    schema: SchemaRef,
519    dictionaries: HashMap<i64, arrow_array::ArrayRef>,
520    finished: bool,
521    /// When `Some`, every read batch is rewrapped with this relaxed
522    /// schema before being returned to the caller (used by the
523    /// conformance worker to accept Python's nullable-flag-lying
524    /// `ArrowSerializableDataclass` outputs).
525    relaxed_schema: Option<SchemaRef>,
526}
527
528impl<R: Read> StreamReader<R> {
529    /// Create a new reader and consume the schema message.
530    ///
531    /// The schema-message length prefix is validated against
532    /// [`MAX_IPC_SCHEMA_BYTES`] *before* allocating, so a remote
533    /// client cannot trigger a multi-gigabyte alloc by sending a
534    /// crafted short payload.
535    pub fn new(mut reader: R) -> Result<Self> {
536        let msg = read_message_bytes(&mut reader, MAX_IPC_SCHEMA_BYTES)?
537            .ok_or_else(|| RpcError::new("IPC", "empty IPC stream (no schema)"))?;
538        if msg.had_invalid_utf8 {
539            return Err(RpcError::protocol_error(
540                "Invalid UTF-8 in IPC schema metadata",
541            ));
542        }
543        let msg_fb = root_as_message(&msg.message_bytes)
544            .map_err(|e| RpcError::new("IPC", format!("parse schema message: {e}")))?;
545        if msg_fb.header_type() != MessageHeader::Schema {
546            return Err(RpcError::new(
547                "IPC",
548                format!("expected Schema, got {:?}", msg_fb.header_type()),
549            ));
550        }
551        let ipc_schema = msg_fb
552            .header_as_schema()
553            .ok_or_else(|| RpcError::new("IPC", "bad schema header"))?;
554        // A legitimate Arrow Schema message always carries a `fields` vector
555        // (possibly empty). When it's absent, `fb_to_schema` does
556        // `fb.fields().unwrap()` and panics — under cargo-fuzz's `panic=abort`
557        // that aborts the process before any `catch_unwind` can intercept.
558        // Reject the malformed frame explicitly here. (crates.io arrow tolerates
559        // this; the pinned arrow-rs fork the fuzz harness uses does not.)
560        if ipc_schema.fields().is_none() {
561            return Err(RpcError::new("IPC", "schema message has no fields vector"));
562        }
563        // `fb_to_schema` still `unwrap()`s other optional members while walking
564        // field types; keep the `catch_unwind` net (matching record-batch
565        // decode) so any residual panic becomes a clean `RpcError` in normal
566        // (panic=unwind) builds rather than escaping the reader.
567        let schema = decode_guard("schema message", || ipc_convert::fb_to_schema(ipc_schema))?;
568        Ok(Self {
569            reader,
570            schema: Arc::new(schema),
571            dictionaries: HashMap::new(),
572            finished: false,
573            relaxed_schema: None,
574        })
575    }
576
577    /// Get the schema of the stream (relaxed schema, if relaxation was
578    /// requested).
579    pub fn schema(&self) -> SchemaRef {
580        self.relaxed_schema
581            .clone()
582            .unwrap_or_else(|| self.schema.clone())
583    }
584
585    /// Promote every field in the stream's schema to `nullable = true`,
586    /// recursively (lists, structs, fixed-size lists). Use when a
587    /// producer declares a field non-nullable but legitimately sends
588    /// nulls — e.g. Python's `ArrowSerializableDataclass` for
589    /// `Annotated[T | None, ArrowType(...)]`.
590    pub fn relax_nullability(mut self) -> Self {
591        self.relaxed_schema = Some(Arc::new(relax_schema_nullability(self.schema.as_ref())));
592        self
593    }
594
595    /// Read the next record batch, or `None` on end-of-stream.
596    /// Returns `(batch, metadata)` where `metadata` carries the IPC
597    /// Message-level `custom_metadata` (empty when the producer
598    /// omitted the field).
599    pub fn read_next(&mut self) -> Result<Option<(RecordBatch, Metadata)>> {
600        if self.finished {
601            return Ok(None);
602        }
603        loop {
604            let msg = match read_message_bytes(&mut self.reader, MAX_IPC_MESSAGE_BYTES)? {
605                Some(m) => m,
606                None => {
607                    self.finished = true;
608                    return Ok(None);
609                }
610            };
611            let msg_fb = root_as_message(&msg.message_bytes)
612                .map_err(|e| RpcError::new("IPC", format!("parse message: {e}")))?;
613            let version = msg_fb.version();
614            match msg_fb.header_type() {
615                MessageHeader::DictionaryBatch => {
616                    let dict = msg_fb
617                        .header_as_dictionary_batch()
618                        .ok_or_else(|| RpcError::new("IPC", "bad dictionary header"))?;
619                    let body_buf = ArrowBuffer::from_vec(msg.body);
620                    // Reject buffer descriptors that point outside the
621                    // body *before* handing them to arrow-ipc, which
622                    // would otherwise panic on an out-of-bounds slice.
623                    if let Some(data) = dict.data() {
624                        validate_record_batch_buffers(&data, body_buf.len())?;
625                    }
626                    // arrow-ipc's decoder still has internal invariants
627                    // we don't re-check; `catch_unwind` is the backstop
628                    // that turns any residual panic into a clean error.
629                    decode_guard("dictionary batch", || {
630                        ipc_reader::read_dictionary(
631                            &body_buf,
632                            dict,
633                            self.schema.as_ref(),
634                            &mut self.dictionaries,
635                            &version,
636                        )
637                    })?
638                    .map_err(RpcError::from)?;
639                }
640                MessageHeader::RecordBatch => {
641                    let rb_fb = msg_fb
642                        .header_as_record_batch()
643                        .ok_or_else(|| RpcError::new("IPC", "bad record batch header"))?;
644                    let body_buf = ArrowBuffer::from_vec(msg.body);
645                    validate_record_batch_buffers(&rb_fb, body_buf.len())?;
646                    // When relaxation is in effect, feed the relaxed
647                    // schema directly to `read_record_batch` so its
648                    // validation accepts the legitimate null buffers
649                    // a producer (e.g. Python
650                    // `ArrowSerializableDataclass`) emits for fields
651                    // it declared `nullable=false`.
652                    let decode_schema = self
653                        .relaxed_schema
654                        .clone()
655                        .unwrap_or_else(|| self.schema.clone());
656                    let batch = decode_guard("record batch", || {
657                        ipc_reader::read_record_batch(
658                            &body_buf,
659                            rb_fb,
660                            decode_schema,
661                            &self.dictionaries,
662                            None,
663                            &version,
664                        )
665                    })?
666                    .map_err(RpcError::from)?;
667                    let mut metadata = parse_custom_metadata(&msg_fb);
668                    if msg.had_invalid_utf8 {
669                        metadata.insert(INVALID_UTF8_METADATA_KEY.into(), "true".into());
670                    }
671                    return Ok(Some((batch, metadata)));
672                }
673                MessageHeader::Schema => {
674                    return Err(RpcError::new("IPC", "unexpected schema message mid-stream"));
675                }
676                MessageHeader::NONE => continue,
677                other => {
678                    return Err(RpcError::new(
679                        "IPC",
680                        format!("unsupported message type {other:?}"),
681                    ));
682                }
683            }
684        }
685    }
686
687    /// Drain and discard any remaining messages.
688    pub fn drain(&mut self) -> Result<()> {
689        while self.read_next()?.is_some() {}
690        Ok(())
691    }
692
693    pub fn get_mut(&mut self) -> &mut R {
694        &mut self.reader
695    }
696}
697
698fn parse_custom_metadata(msg: &arrow_ipc::Message) -> Metadata {
699    let Some(md) = msg.custom_metadata() else {
700        return Metadata::new();
701    };
702    // Size the map to the known key count so it doesn't rehash while filling.
703    let mut out = Metadata::with_capacity(md.len());
704    for kv in md.iter() {
705        let k = kv.key().unwrap_or("").to_string();
706        let v = kv.value().unwrap_or("").to_string();
707        out.insert(k, v);
708    }
709    out
710}
711
712/// Validate that every `(offset, length)` buffer descriptor in an IPC
713/// record-batch header references a region wholly inside the message
714/// body. arrow-ipc's column decoders index into the body using these
715/// descriptors verbatim and will panic (slice out-of-bounds / arithmetic
716/// overflow) on a crafted frame whose descriptors are inconsistent with
717/// the body it shipped. Catching that here turns a hostile frame into a
718/// clean `RpcError` instead of a thread panic.
719pub(crate) fn validate_record_batch_buffers(
720    rb: &arrow_ipc::RecordBatch,
721    body_len: usize,
722) -> Result<()> {
723    if let Some(buffers) = rb.buffers() {
724        for buf in buffers.iter() {
725            let offset = buf.offset();
726            let length = buf.length();
727            if offset < 0 || length < 0 {
728                return Err(RpcError::new("IPC", "negative IPC buffer descriptor"));
729            }
730            let end = (offset as u64)
731                .checked_add(length as u64)
732                .ok_or_else(|| RpcError::new("IPC", "IPC buffer descriptor overflows"))?;
733            if end > body_len as u64 {
734                return Err(RpcError::new(
735                    "IPC",
736                    "IPC buffer descriptor exceeds message body",
737                ));
738            }
739        }
740    }
741    Ok(())
742}
743
744/// Run an arrow-ipc decode call, converting any panic into a clean
745/// `RpcError`. The descriptor pre-validation above catches the common
746/// crafted-frame cases; this is the defence-in-depth net for any other
747/// internal arrow-ipc invariant a hostile frame might trip.
748pub(crate) fn decode_guard<T>(what: &str, f: impl FnOnce() -> T) -> Result<T> {
749    std::panic::catch_unwind(std::panic::AssertUnwindSafe(f))
750        .map_err(|_| RpcError::new("IPC", format!("panic decoding {what} (malformed frame)")))
751}
752
753struct RawMessage {
754    message_bytes: Vec<u8>,
755    body: Vec<u8>,
756    had_invalid_utf8: bool,
757}
758
759fn read_exact(r: &mut impl Read, buf: &mut [u8]) -> Result<bool> {
760    let mut read = 0;
761    while read < buf.len() {
762        match r.read(&mut buf[read..]) {
763            Ok(0) => {
764                if read == 0 {
765                    return Ok(false);
766                }
767                return Err(RpcError::new("IOError", "unexpected EOF in IPC message"));
768            }
769            Ok(n) => read += n,
770            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
771            Err(e) => return Err(e.into()),
772        }
773    }
774    Ok(true)
775}
776
777/// Read one IPC message off `r`, capping the header and body at
778/// `max_bytes` so a crafted length prefix or flatbuffer
779/// `bodyLength` cannot trigger an unbounded allocation.
780///
781/// The header flatbuffer has to be buffered whole before it can be
782/// parsed, so its length prefix is checked against `max_bytes` and
783/// allocated outright. The body is not: it is grown from the bytes the
784/// peer actually delivers, so the ceiling can be generous enough for a
785/// legitimate multi-gigabyte batch without a lying `bodyLength` costing
786/// more than [`BODY_PREALLOC_LIMIT`] and an EOF.
787fn read_message_bytes(r: &mut impl Read, max_bytes: usize) -> Result<Option<RawMessage>> {
788    let mut prefix = [0u8; 4];
789    if !read_exact(r, &mut prefix)? {
790        return Ok(None);
791    }
792    let size_bytes = if prefix == CONTINUATION_MARKER {
793        let mut sb = [0u8; 4];
794        if !read_exact(r, &mut sb)? {
795            return Ok(None);
796        }
797        sb
798    } else {
799        prefix
800    };
801    let size = u32::from_le_bytes(size_bytes) as usize;
802    if size == 0 {
803        // EOS
804        return Ok(None);
805    }
806    if size > max_bytes {
807        return Err(RpcError::new(
808            "IPC",
809            format!(
810                "IPC message header length {size} bytes exceeds cap {max_bytes} — \
811                 refusing to allocate before parsing"
812            ),
813        ));
814    }
815    let mut message_bytes = vec![0u8; size];
816    if !read_exact(r, &mut message_bytes)? {
817        return Err(RpcError::new("IOError", "unexpected EOF in message body"));
818    }
819    // Parse just enough to learn the body length, then refuse an absurd
820    // claim outright. Invalid UTF-8 in custom metadata is a semantic request
821    // error, not broken framing: Flatbuffers reports its exact byte range
822    // after validating the surrounding offsets, so replace only that range in
823    // our private header copy and remember to reject it before dispatch.
824    let mut had_invalid_utf8 = false;
825    let msg = loop {
826        match root_as_message(&message_bytes) {
827            Ok(msg) => break msg,
828            Err(flatbuffers::InvalidFlatbuffer::Utf8Error { range, .. })
829                if !range.is_empty() && range.end <= message_bytes.len() =>
830            {
831                message_bytes[range].fill(b'?');
832                had_invalid_utf8 = true;
833            }
834            Err(e) => {
835                return Err(RpcError::new("IPC", format!("parse message header: {e}")));
836            }
837        }
838    };
839    let body_length_signed = msg.bodyLength();
840    if body_length_signed < 0 {
841        return Err(RpcError::new(
842            "IPC",
843            format!("IPC message has negative bodyLength ({body_length_signed})"),
844        ));
845    }
846    // Compare in u64: on a 32-bit target the claim can exceed anything
847    // `usize` can hold, and truncating first would let it wrap under the
848    // cap.
849    if body_length_signed as u64 > max_bytes as u64 {
850        return Err(RpcError::new(
851            "IPC",
852            format!(
853                "IPC message bodyLength {body_length_signed} bytes exceeds cap {max_bytes} — \
854                 refusing to allocate before parsing"
855            ),
856        ));
857    }
858    let body_length = body_length_signed as usize;
859    // Reserve only what a normal batch needs; past that the buffer grows
860    // as the bytes arrive, so the peer pays for the size it claimed
861    // before we do.
862    let mut body = Vec::with_capacity(body_length.min(BODY_PREALLOC_LIMIT));
863    if body_length > 0 {
864        let read = (&mut *r)
865            .take(body_length as u64)
866            .read_to_end(&mut body)
867            .map_err(RpcError::from)?;
868        if read != body_length {
869            return Err(RpcError::new("IOError", "unexpected EOF in message body"));
870        }
871    }
872    Ok(Some(RawMessage {
873        message_bytes,
874        body,
875        had_invalid_utf8,
876    }))
877}
878
879// ---------------------------------------------------------------------------
880// Utilities
881// ---------------------------------------------------------------------------
882
883/// Serialize one record batch as a complete IPC stream
884/// (schema + batch + EOS), with optional custom metadata on the batch.
885pub fn write_one_batch(batch: &RecordBatch, metadata: Option<&Metadata>) -> Result<Vec<u8>> {
886    write_one_batch_as(batch, batch.schema().as_ref(), metadata)
887}
888
889/// Like [`write_one_batch`] but declares `schema` on the stream instead of
890/// the batch's own schema, writing the batch's buffers unchanged.
891///
892/// [`StreamWriter::write`] never reconciles a batch against the schema its
893/// stream was opened with — it encodes the buffers and the reader decodes
894/// them under the declared schema. So a batch that differs from its
895/// enclosing stream's schema only cosmetically (field nullability,
896/// dictionary encoding, schema-level metadata) round-trips invisibly while
897/// it stays inline.
898///
899/// That stops being true the moment the batch is lifted onto a *standalone*
900/// stream, as external-location payloads are: the payload declares its own
901/// schema, and a peer that validates it against the schema it was promised
902/// (the enclosing stream's) sees a hard mismatch over a difference that
903/// never mattered before. Passing the enclosing schema here keeps the two
904/// delivery routes indistinguishable.
905///
906/// Deliberately not a cast: the buffers are emitted as-is, so this cannot
907/// silently do nothing the way an "equivalent schemas" fast path in a cast
908/// helper would, and it cannot change the bytes either.
909pub fn write_one_batch_as(
910    batch: &RecordBatch,
911    schema: &Schema,
912    metadata: Option<&Metadata>,
913) -> Result<Vec<u8>> {
914    let mut buf = Vec::new();
915    {
916        let mut w = StreamWriter::new(&mut buf, schema)?;
917        w.write(batch, metadata)?;
918        w.finish()?;
919    }
920    Ok(buf)
921}
922
923/// Lowercase hex encoding of a byte slice. Internal helper — use the
924/// `hex` crate from your application code.
925// Only the http/external/mtls modules call this; unused in a minimal
926// (macros-only) wasm build, so suppress the conditional dead-code warning.
927#[allow(dead_code)]
928pub(crate) fn bytes_to_hex(bytes: &[u8]) -> String {
929    const HEX: &[u8; 16] = b"0123456789abcdef";
930    let mut out = String::with_capacity(bytes.len() * 2);
931    for b in bytes {
932        out.push(HEX[(b >> 4) as usize] as char);
933        out.push(HEX[(b & 0x0f) as usize] as char);
934    }
935    out
936}
937
938fn relax_field_nullability(f: &arrow_schema::Field) -> arrow_schema::Field {
939    use arrow_schema::DataType;
940    let dt = match f.data_type() {
941        DataType::List(inner) => DataType::List(Arc::new(relax_field_nullability(inner))),
942        DataType::LargeList(inner) => DataType::LargeList(Arc::new(relax_field_nullability(inner))),
943        DataType::FixedSizeList(inner, n) => {
944            DataType::FixedSizeList(Arc::new(relax_field_nullability(inner)), *n)
945        }
946        DataType::Struct(fields) => DataType::Struct(
947            fields
948                .iter()
949                .map(|child| Arc::new(relax_field_nullability(child)))
950                .collect(),
951        ),
952        // Map: leave the entries struct alone (Arrow requires
953        // entries/keys to be non-nullable); leaf nullability inside
954        // the values child is preserved by the original schema.
955        other => other.clone(),
956    };
957    #[allow(deprecated)]
958    let new_field = if let DataType::Dictionary(_, _) = f.data_type() {
959        arrow_schema::Field::new_dict(
960            f.name(),
961            dt,
962            true,
963            f.dict_id().unwrap_or(0),
964            f.dict_is_ordered().unwrap_or(false),
965        )
966    } else {
967        arrow_schema::Field::new(f.name(), dt, true)
968    };
969    new_field.with_metadata(f.metadata().clone())
970}
971
972fn relax_schema_nullability(s: &Schema) -> Schema {
973    let new_fields: Vec<arrow_schema::Field> = s
974        .fields()
975        .iter()
976        .map(|f| relax_field_nullability(f))
977        .collect();
978    Schema::new_with_metadata(new_fields, s.metadata().clone())
979}
980
981/// Build a zero-row `RecordBatch` matching the given schema.
982pub fn empty_batch(schema: &Schema) -> Result<RecordBatch> {
983    use arrow_array::array::new_empty_array;
984    use arrow_array::RecordBatchOptions;
985    let cols: Vec<arrow_array::ArrayRef> = schema
986        .fields()
987        .iter()
988        .map(|f| new_empty_array(f.data_type()))
989        .collect();
990    RecordBatch::try_new_with_options(
991        Arc::new(schema.clone()),
992        cols,
993        &RecordBatchOptions::new().with_row_count(Some(0)),
994    )
995    .map_err(RpcError::from)
996}
997
998#[cfg(test)]
999mod tests {
1000    use super::*;
1001    use arrow_array::{Int64Array, StringArray};
1002    use arrow_schema::{DataType, Field};
1003
1004    /// Records what each `write` call was *offered* and honours a
1005    /// caller-chosen short-count, so both halves of the large-payload
1006    /// contract can be observed: the clamp and the retry.
1007    struct SabotageWriter {
1008        offered: Vec<usize>,
1009        accept: usize,
1010        sink: Vec<u8>,
1011    }
1012
1013    impl Write for SabotageWriter {
1014        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
1015            self.offered.push(buf.len());
1016            let n = buf.len().min(self.accept);
1017            self.sink.extend_from_slice(&buf[..n]);
1018            Ok(n)
1019        }
1020        fn flush(&mut self) -> std::io::Result<()> {
1021            Ok(())
1022        }
1023    }
1024
1025    #[test]
1026    fn chunked_writer_clamps_and_retries() {
1027        // A macOS pipe answers a >2 GiB write with a short count and no
1028        // error; a macOS socket answers it with EINVAL. Surviving both
1029        // needs the clamp *and* the loop, so assert both here rather
1030        // than trusting `write_all` alone.
1031        let mut w = ChunkedWriter::with_limit(
1032            SabotageWriter {
1033                offered: Vec::new(),
1034                accept: 3,
1035                sink: Vec::new(),
1036            },
1037            8,
1038        );
1039        let payload: Vec<u8> = (0..50u8).collect();
1040        w.write_all(&payload).unwrap();
1041        let inner = w.get_mut();
1042        assert!(
1043            inner.offered.iter().all(|n| *n <= 8),
1044            "a write was offered more than the clamp: {:?}",
1045            inner.offered
1046        );
1047        assert_eq!(inner.sink, payload, "short writes lost bytes");
1048    }
1049
1050    #[test]
1051    fn write_chunk_stays_under_int_max() {
1052        // The clamp is only worth anything if it lands below the size at
1053        // which macOS starts rejecting or truncating.
1054        assert!(MAX_WRITE_CHUNK < i32::MAX as usize);
1055    }
1056
1057    #[test]
1058    fn oversized_body_claim_costs_nothing_to_refuse() {
1059        // The ceiling is generous enough for a legitimate multi-gigabyte
1060        // batch, so the flatbuffer overshoot the fuzzer found has to be
1061        // refused by the cap and not by a lucky allocation failure.
1062        assert!(MAX_IPC_MESSAGE_BYTES as u64 > (1u64 << 31) + 1);
1063        assert!((0x4000000100000u64) > MAX_IPC_MESSAGE_BYTES as u64);
1064    }
1065
1066    #[test]
1067    fn roundtrip_with_metadata() {
1068        let schema = Schema::new(vec![
1069            Field::new("idx", DataType::Int64, false),
1070            Field::new("name", DataType::Utf8, false),
1071        ]);
1072        let batch = RecordBatch::try_new(
1073            Arc::new(schema.clone()),
1074            vec![
1075                Arc::new(Int64Array::from(vec![1, 2, 3])) as _,
1076                Arc::new(StringArray::from(vec!["a", "b", "c"])) as _,
1077            ],
1078        )
1079        .unwrap();
1080
1081        let mut buf: Vec<u8> = Vec::new();
1082        {
1083            let mut w = StreamWriter::new(&mut buf, &schema).unwrap();
1084            let mut md = Metadata::new();
1085            md.insert("vgi_rpc.method".into(), "echo_string".into());
1086            w.write(&batch, Some(&md)).unwrap();
1087            w.finish().unwrap();
1088        }
1089
1090        let mut r = StreamReader::new(buf.as_slice()).unwrap();
1091        let (rb, md) = r.read_next().unwrap().expect("batch");
1092        assert_eq!(rb.num_rows(), 3);
1093        assert_eq!(md_get(&md, "vgi_rpc.method"), Some("echo_string"));
1094        assert!(r.read_next().unwrap().is_none());
1095    }
1096
1097    #[test]
1098    fn single_large_binary_direct_write_roundtrips_sliced_value_and_metadata() {
1099        use arrow_array::LargeBinaryArray;
1100
1101        let schema = Schema::new(vec![Field::new("result", DataType::LargeBinary, false)]);
1102        let source = LargeBinaryArray::from_iter_values([
1103            b"discarded prefix".as_slice(),
1104            b"payload retained in the original Arrow allocation".as_slice(),
1105        ]);
1106        let sliced = source.slice(1, 1);
1107        let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(sliced)]).unwrap();
1108        let mut metadata = Metadata::new();
1109        metadata.insert("vgi_rpc.method".into(), "echo_large_binary".into());
1110
1111        let mut encoded = Vec::new();
1112        {
1113            let mut writer = StreamWriter::new(&mut encoded, &schema).unwrap();
1114            writer.write(&batch, Some(&metadata)).unwrap();
1115            writer.finish().unwrap();
1116        }
1117
1118        let mut reader = StreamReader::new(encoded.as_slice()).unwrap();
1119        let (decoded, decoded_metadata) = reader.read_next().unwrap().expect("batch");
1120        let decoded = decoded
1121            .column(0)
1122            .as_any()
1123            .downcast_ref::<LargeBinaryArray>()
1124            .unwrap();
1125        assert_eq!(
1126            decoded.value(0),
1127            b"payload retained in the original Arrow allocation"
1128        );
1129        assert_eq!(
1130            md_get(&decoded_metadata, "vgi_rpc.method"),
1131            Some("echo_large_binary")
1132        );
1133        assert!(reader.read_next().unwrap().is_none());
1134    }
1135
1136    #[test]
1137    fn zero_row_metadata_only() {
1138        let schema = Schema::empty();
1139        let batch = empty_batch(&schema).unwrap();
1140
1141        let mut buf: Vec<u8> = Vec::new();
1142        {
1143            let mut w = StreamWriter::new(&mut buf, &schema).unwrap();
1144            let mut md = Metadata::new();
1145            md.insert("vgi_rpc.log_level".into(), "INFO".into());
1146            w.write(&batch, Some(&md)).unwrap();
1147            w.finish().unwrap();
1148        }
1149        let mut r = StreamReader::new(buf.as_slice()).unwrap();
1150        let (rb, md) = r.read_next().unwrap().expect("batch");
1151        assert_eq!(rb.num_rows(), 0);
1152        assert_eq!(md_get(&md, "vgi_rpc.log_level"), Some("INFO"));
1153    }
1154
1155    #[test]
1156    fn rejects_oversize_schema_length_prefix() {
1157        // The 4-byte payload `[0x1A, 0x2C, 0xF5, 0x2C]` parsed LE
1158        // claims ~720 MB of schema-message body — must be refused
1159        // before any allocation.
1160        let bomb: &[u8] = &[0x1A, 0x2C, 0xF5, 0x2C];
1161        let err = StreamReader::new(bomb).err().expect("must reject");
1162        assert!(
1163            err.message.contains("exceeds cap"),
1164            "unexpected error: {err:?}"
1165        );
1166    }
1167
1168    #[test]
1169    fn rejects_oversize_message_bodylength() {
1170        // Encode a tiny but well-formed schema then send a record-
1171        // batch message whose flatbuffer claims a multi-GB
1172        // `bodyLength` — must be refused before allocating the body.
1173        use arrow_ipc::{Buffer as FbBuffer, FieldNode, MessageBuilder, RecordBatchBuilder};
1174        // Build a real schema first so the schema gate passes.
1175        let schema = Schema::new(vec![Field::new("v", DataType::Int64, false)]);
1176        let mut buf: Vec<u8> = Vec::new();
1177        {
1178            let w = StreamWriter::new(&mut buf, &schema).unwrap();
1179            // Don't write any batches; we'll append a hand-crafted
1180            // malicious message below.
1181            // Drop without finish so EOS is not written.
1182            std::mem::forget(w);
1183        }
1184        // Hand-craft a RecordBatch Message flatbuffer with absurd
1185        // bodyLength.
1186        let mut fbb = FlatBufferBuilder::new();
1187        let nodes_vec = fbb.create_vector(&[FieldNode::new(0, 0)]);
1188        let buffers_vec = fbb.create_vector(&[FbBuffer::new(0, 0)]);
1189        let rb_off = {
1190            let mut b = RecordBatchBuilder::new(&mut fbb);
1191            b.add_length(0);
1192            b.add_nodes(nodes_vec);
1193            b.add_buffers(buffers_vec);
1194            b.finish()
1195        };
1196        let msg_off = {
1197            let mut mb = MessageBuilder::new(&mut fbb);
1198            mb.add_version(arrow_ipc::MetadataVersion::V5);
1199            mb.add_header_type(MessageHeader::RecordBatch);
1200            mb.add_header(rb_off.as_union_value());
1201            mb.add_bodyLength(MAX_IPC_MESSAGE_BYTES as i64 + 1);
1202            mb.finish()
1203        };
1204        fbb.finish(msg_off, None);
1205        let msg_bytes = fbb.finished_data();
1206        // Frame: continuation + 4-byte LE length + flatbuffer body.
1207        buf.extend_from_slice(&CONTINUATION_MARKER);
1208        buf.extend_from_slice(&(msg_bytes.len() as u32).to_le_bytes());
1209        buf.extend_from_slice(msg_bytes);
1210        // No body — but we never get that far; the cap rejects first.
1211
1212        let mut r = StreamReader::new(buf.as_slice()).unwrap();
1213        let err = r.read_next().expect_err("must reject");
1214        assert!(
1215            err.message.contains("bodyLength") && err.message.contains("exceeds cap"),
1216            "unexpected error: {err:?}"
1217        );
1218    }
1219
1220    #[test]
1221    fn malformed_schema_message_is_error_not_panic() {
1222        // Regression for a fuzz-found crash: a structurally-parseable but
1223        // malformed Schema message made arrow-ipc's `fb_to_schema` panic
1224        // (`Option::unwrap()` on `None` in convert.rs) out of `StreamReader::new`
1225        // — aborting the process instead of returning a clean error. The
1226        // schema parse must now be caught like the per-batch decode.
1227        // `crash-cea0477693563377f77c693ca8d3df51ee421811` from
1228        // `fuzz/wire_stream_reader`.
1229        let crash: &[u8] = &[
1230            22, 0, 0, 0, 12, 0, 0, 0, 0, 0, 8, 0, 4, 0, 0, 0, 12, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 134,
1231        ];
1232        // Must not panic; either an Err here or a clean reader is acceptable —
1233        // the contract is "no unwind escapes".
1234        let _ = StreamReader::new(crash);
1235    }
1236
1237    #[test]
1238    fn rejects_buffer_descriptor_past_body() {
1239        // A record-batch message whose body is 8 bytes but whose buffer
1240        // descriptor claims offset 0 / length 1000. arrow-ipc would
1241        // index out of bounds and panic; the descriptor pre-check must
1242        // reject it as a clean `RpcError` first.
1243        use arrow_ipc::{Buffer as FbBuffer, FieldNode, MessageBuilder, RecordBatchBuilder};
1244        let schema = Schema::new(vec![Field::new("v", DataType::Int64, false)]);
1245        let mut buf: Vec<u8> = Vec::new();
1246        {
1247            let w = StreamWriter::new(&mut buf, &schema).unwrap();
1248            std::mem::forget(w);
1249        }
1250        let mut fbb = FlatBufferBuilder::new();
1251        let nodes_vec = fbb.create_vector(&[FieldNode::new(1, 0)]);
1252        // offset 0, length 1000 — far past the 8-byte body below.
1253        let buffers_vec = fbb.create_vector(&[FbBuffer::new(0, 1000)]);
1254        let rb_off = {
1255            let mut b = RecordBatchBuilder::new(&mut fbb);
1256            b.add_length(1);
1257            b.add_nodes(nodes_vec);
1258            b.add_buffers(buffers_vec);
1259            b.finish()
1260        };
1261        let msg_off = {
1262            let mut mb = MessageBuilder::new(&mut fbb);
1263            mb.add_version(arrow_ipc::MetadataVersion::V5);
1264            mb.add_header_type(MessageHeader::RecordBatch);
1265            mb.add_header(rb_off.as_union_value());
1266            mb.add_bodyLength(8);
1267            mb.finish()
1268        };
1269        fbb.finish(msg_off, None);
1270        let msg_bytes = fbb.finished_data().to_vec();
1271        buf.extend_from_slice(&CONTINUATION_MARKER);
1272        buf.extend_from_slice(&(msg_bytes.len() as u32).to_le_bytes());
1273        buf.extend_from_slice(&msg_bytes);
1274        buf.extend_from_slice(&[0u8; 8]); // the 8-byte body
1275
1276        let mut r = StreamReader::new(buf.as_slice()).unwrap();
1277        let err = r.read_next().expect_err("must reject");
1278        assert!(
1279            err.message.contains("buffer descriptor"),
1280            "unexpected error: {err:?}"
1281        );
1282    }
1283}