From e1ca7fe79f0902c608089320d8524f372b5efdfc Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Mon, 26 Oct 2026 15:06:11 -0400 Subject: [PATCH 1/3] fix: validate complete typed IPC values through byte codecs --- docs/api/serialization.md | 15 ++ tests/test_single_record_framing.py | 261 ++++++++++++++++++++++++++++ vgi_rpc/rpc/_wire.py | 13 +- vgi_rpc/utils.py | 42 ++++- 4 files changed, 321 insertions(+), 10 deletions(-) create mode 100644 tests/test_single_record_framing.py diff --git a/docs/api/serialization.md b/docs/api/serialization.md index ddd6ea8..fdea49c 100644 --- a/docs/api/serialization.md +++ b/docs/api/serialization.md @@ -36,6 +36,21 @@ assert m == m2 These dataclasses work directly as RPC parameters and return types. They're also the base class for [`StreamState`](streaming.md) — any field you add to a stream state is automatically serialized between calls. +RPC parameters and unary results are decoded through the class's +`deserialize_from_bytes(data, ipc_validation)` method. The default decoder +requires a complete single-batch IPC value: physical EOF and the legacy or +modern Arrow EOS marker are accepted; extra messages, trailing bytes and +truncated endings are rejected. These checks also apply when batch validation +is disabled. The low-level `deserialize_record_batch` helper continues to accept +empty bytes and schema-only streams as empty nested values. + +A subclass can override `deserialize_from_bytes` to apply a raw byte limit or +validate a custom envelope before decoding. The override receives the original +bytes and validation level and is responsible for complete envelope validation. +Calling `super().deserialize_from_bytes(data, ipc_validation)` retains the +default framing checks. This does not add a global size limit or change the +framing of streaming RPC inputs and results. + ### Type mappings | Python type | Arrow type | diff --git a/tests/test_single_record_framing.py b/tests/test_single_record_framing.py new file mode 100644 index 0000000..c9ee9e7 --- /dev/null +++ b/tests/test_single_record_framing.py @@ -0,0 +1,261 @@ +# © Copyright 2025-2026, Query.Farm LLC - https://query.farm +# SPDX-License-Identifier: Apache-2.0 + +"""Bounded typed IPC values must consume their complete single-record envelope.""" + +from __future__ import annotations + +from dataclasses import dataclass +from io import BytesIO +from typing import ClassVar, Protocol, Self + +import falcon.testing +import pyarrow as pa +import pyarrow.ipc as ipc +import pytest + +from vgi_rpc.http import make_wsgi_app +from vgi_rpc.rpc import PipeTransport, RpcConnection, RpcError, RpcServer, rpc_methods +from vgi_rpc.rpc._wire import _deserialize_value +from vgi_rpc.utils import ( + ArrowSerializableDataclass, + IPCError, + IpcValidation, + deserialize_record_batch, + read_single_record_batch, + serialize_record_batch_bytes, +) +from vgi_rpc.wire import write_request, write_unary_result + +_EOS = b"\xff\xff\xff\xff\0\0\0\0" + + +@dataclass(frozen=True) +class _Request(ArrowSerializableDataclass): + """A minimal typed request whose first batch would otherwise be accepted.""" + + value: int + + +class _FramingProtocol(Protocol): + """Expose the bounded value through public HTTP and pipe dispatch.""" + + protocol_name: ClassVar[str] = "test.Framing" + + def accept(self, request: _Request) -> int: + """Return the accepted request value.""" + ... + + +class _ResultProtocol(Protocol): + """Exercise the same typed decoder on the receiving client.""" + + def retrieve(self) -> _Request: + """Return a typed record.""" + ... + + +class _FramingImpl: + """Record whether malformed inputs reached application code.""" + + def __init__(self) -> None: + """Start with no dispatched calls.""" + self.calls: list[int] = [] + + def accept(self, request: _Request) -> int: + """Record and return the accepted value.""" + self.calls.append(request.value) + return request.value + + +def _payload(*, batches: int = 1) -> bytes: + """Construct the inner typed value without involving the RPC decoder.""" + batch = pa.record_batch([[42]], schema=_Request.ARROW_SCHEMA) + out = BytesIO() + with ipc.new_stream(out, batch.schema) as writer: + for _ in range(batches): + writer.write_batch(batch, custom_metadata={b"tag": b"preserved"}) + return out.getvalue() + + +def _ending(data: bytes, ending: str) -> bytes: + """Replace a modern EOS with another supported Arrow stream ending.""" + assert data.endswith(_EOS) + return data[:-8] + {"modern": _EOS, "legacy": b"\0\0\0\0", "eof": b""}[ending] + + +_MALFORMED = ( + "extra_batch", + "concatenated_stream", + "trailing_byte", + "legacy_trailing_byte", + "repeated_eos", + "truncated_schema", + "truncated_body", + "truncated_next_message", + *(f"truncated_eos_{length}" for length in range(1, 8)), + *(f"truncated_legacy_eos_{length}" for length in range(1, 4)), +) + + +def _malformed(kind: str) -> bytes: + """Build independent corruptions of a valid single-batch envelope.""" + data = _payload() + if kind.startswith("truncated_eos_"): + return data[:-8] + _EOS[: int(kind.rsplit("_", 1)[1])] + if kind.startswith("truncated_legacy_eos_"): + return data[:-8] + b"\0" * int(kind.rsplit("_", 1)[1]) + return { + "extra_batch": _payload(batches=2), + "concatenated_stream": data + data, + "trailing_byte": data + b"x", + "legacy_trailing_byte": _ending(data, "legacy") + b"x", + "repeated_eos": data + _EOS, + "truncated_schema": data[:12], + "truncated_body": data[:-9], + "truncated_next_message": data[:-8] + b"\xff\xff\xff\xff\x08\0\0\0x", + }[kind] + + +@pytest.mark.parametrize("validation", list(IpcValidation)) +@pytest.mark.parametrize("ending", ["modern", "legacy", "eof"]) +def test_single_batch_accepts_arrow_endings(ending: str, validation: IpcValidation) -> None: + """Complete valid streams preserve values and record custom metadata.""" + data = _ending(_payload(), ending) + batch, metadata = deserialize_record_batch(data, validation) + assert batch.to_pylist() == [{"value": 42}] + assert metadata is not None and metadata[b"tag"] == b"preserved" + assert _deserialize_value(data, _Request, validation) == _Request(42) + + +@pytest.mark.parametrize("validation", list(IpcValidation)) +@pytest.mark.parametrize("kind", _MALFORMED) +def test_single_batch_rejects_incomplete_or_surplus_framing(kind: str, validation: IpcValidation) -> None: + """Disabling buffer validation must never disable envelope framing checks.""" + data = _malformed(kind) + with pytest.raises((IPCError, pa.ArrowException, ValueError)): + deserialize_record_batch(data, validation) + with pytest.raises((TypeError, ValueError)): + _deserialize_value(data, _Request, validation) + + +@pytest.mark.parametrize("ending", ["modern", "legacy", "eof"]) +def test_schema_only_and_empty_utility_values_remain_supported(ending: str) -> None: + """The generic helper retains the cross-language empty-value convention.""" + empty, metadata = deserialize_record_batch(b"") + assert empty.num_rows == empty.num_columns == 0 + assert metadata is None + data = _ending(_payload(batches=0), ending) + batch, metadata = deserialize_record_batch(data) + assert batch.num_rows == 0 and batch.schema == _Request.ARROW_SCHEMA + assert metadata is None + with pytest.raises((IPCError, pa.ArrowException, ValueError)): + deserialize_record_batch(data + b"x") + + +def test_dictionary_messages_do_not_count_as_extra_batches() -> None: + """A dictionary and its single record batch remain a valid typed value.""" + dictionary = pa.array(["red", "green", "red"]).dictionary_encode() + original = pa.record_batch([dictionary], names=["color"]) + result, _ = deserialize_record_batch(serialize_record_batch_bytes(original)) + assert result.equals(original) + + +def test_legacy_format_messages_and_eos_are_supported() -> None: + """Legacy prefixes on schema and batch messages remain interoperable.""" + batch = pa.record_batch([[42]], schema=_Request.ARROW_SCHEMA) + out = BytesIO() + with ipc.new_stream(out, batch.schema, options=ipc.IpcWriteOptions(use_legacy_format=True)) as writer: + writer.write_batch(batch) + data = out.getvalue() + assert data[:4] != b"\xff\xff\xff\xff" + assert _deserialize_value(data, _Request) == _Request(42) + + +@pytest.mark.parametrize("validation", list(IpcValidation)) +def test_parameter_decoder_honors_bytes_override_before_arrow( + validation: IpcValidation, monkeypatch: pytest.MonkeyPatch +) -> None: + """A request type can reject its raw byte size before any Arrow decoding.""" + seen: list[IpcValidation] = [] + + @dataclass(frozen=True) + class BoundedRequest(_Request): + """Apply an application-specific raw cap using the existing bytes hook.""" + + @classmethod + def deserialize_from_bytes(cls, data: bytes, ipc_validation: IpcValidation = IpcValidation.FULL) -> Self: + """Check size before delegating to the generic Arrow decoder.""" + seen.append(ipc_validation) + if len(data) > 1024: + raise ValueError("request exceeds raw byte limit") + return super().deserialize_from_bytes(data, ipc_validation) + + def unexpected_decode(*args: object, **kwargs: object) -> None: + """Fail immediately if the raw bound is checked too late.""" + pytest.fail("Arrow decoder ran before the request's raw byte limit") + + monkeypatch.setattr(ipc, "open_stream", unexpected_decode) + with pytest.raises(ValueError, match="raw byte limit"): + _deserialize_value(b"x" * 1025, BoundedRequest, validation) + assert seen == [validation] + + +def test_stream_reader_leaves_following_transport_frame_available() -> None: + """Strict bounded-value decoding must not consume the next pipe RPC frame.""" + stream = BytesIO(_payload() + _payload()) + first, _ = read_single_record_batch(stream) + second, _ = read_single_record_batch(stream) + assert first.equals(second) + assert stream.tell() == len(stream.getvalue()) + + +def _request(data: bytes) -> bytes: + """Wrap arbitrary IPC bytes in an otherwise valid public RPC request.""" + schema = rpc_methods(_FramingProtocol)["accept"].params_schema + return write_request("accept", schema, {"request": data}, protocol=_FramingProtocol.protocol_name) + + +@pytest.mark.parametrize("kind", _MALFORMED) +def test_pipe_rejects_bad_typed_value_then_dispatches_next_request(kind: str) -> None: + """Reject before the handler and preserve the next request on the same pipe.""" + impl = _FramingImpl() + server = RpcServer(_FramingProtocol, impl) + output = BytesIO() + transport = PipeTransport(BytesIO(_request(_malformed(kind)) + _request(_payload())), output) + server.serve_one(transport) + assert impl.calls == [] + server.serve_one(transport) + assert impl.calls == [42] + output.seek(0) + with ipc.open_stream(output) as errors: + first = list(errors.iter_batches_with_custom_metadata()) + assert any(md is not None and md.get(b"vgi_rpc.log_level") == b"EXCEPTION" for _, md in first) + with ipc.open_stream(output) as results: + assert results.read_next_batch().column("result")[0].as_py() == 42 + + +@pytest.mark.parametrize("kind", _MALFORMED) +def test_http_rejects_bad_typed_value_as_client_error(kind: str) -> None: + """Malformed wrapped IPC is HTTP 400, never handler execution or HTTP 500.""" + impl = _FramingImpl() + client = falcon.testing.TestClient(make_wsgi_app(RpcServer(_FramingProtocol, impl), token_key=b"framing-test")) + headers = {"content-type": "application/vnd.apache.arrow.stream"} + path = "/test.Framing/accept" + bad = client.simulate_post(path, body=_request(_malformed(kind)), headers=headers) + assert bad.status_code == 400 + assert impl.calls == [] + good = client.simulate_post(path, body=_request(_payload()), headers=headers) + assert good.status_code == 200 + assert impl.calls == [42] + + +@pytest.mark.parametrize("kind", ["extra_batch", "trailing_byte", "truncated_eos_4"]) +def test_client_classifies_malformed_typed_result_as_transport_error(kind: str) -> None: + """Typed result corruption keeps the client's established RpcError contract.""" + schema = rpc_methods(_ResultProtocol)["retrieve"].result_schema + response = write_unary_result(schema, _malformed(kind)) + transport = PipeTransport(BytesIO(response), BytesIO()) + with RpcConnection(_ResultProtocol, transport) as client, pytest.raises(RpcError) as caught: + client.retrieve() + assert caught.value.error_type == "TransportError" diff --git a/vgi_rpc/rpc/_wire.py b/vgi_rpc/rpc/_wire.py index 5c86a2d..806ddd0 100644 --- a/vgi_rpc/rpc/_wire.py +++ b/vgi_rpc/rpc/_wire.py @@ -70,6 +70,7 @@ from vgi_rpc.shm import ShmSegment, is_shm_pointer_batch, maybe_write_to_shm, resolve_shm_batch from vgi_rpc.utils import ( ArrowSerializableDataclass, + IPCError, IpcValidation, ValidatedReader, _annotation_details, @@ -798,9 +799,15 @@ def _deserialize_value(value: object, type_hint: object, ipc_validation: IpcVali if isinstance(base, type) and issubclass(base, ArrowSerializableDataclass): if not isinstance(value, bytes): raise TypeError(f"Expected bytes for {base.__name__} deserialization, got {type(value).__name__}") - reader = ValidatedReader(ipc.open_stream(value), ipc_validation) - batch, metadata = reader.read_next_batch_with_custom_metadata() - return base.deserialize_from_batch(batch, metadata, ipc_validation=ipc_validation) + # Use the same bytes hook as nested values. A custom codec can bound + # the raw payload before Arrow decoding, and the default codec checks + # the complete single-record envelope instead of accepting its prefix. + try: + return base.deserialize_from_bytes(value, ipc_validation) + except (IPCError, pa.ArrowException) as exc: + # HTTP treats ArrowInvalid as a bad request, and clients preserve + # their existing TransportError handling for malformed unary results. + raise pa.ArrowInvalid(f"Invalid {base.__name__} IPC value: {exc}") from exc if isinstance(base, type) and issubclass(base, Enum) and not isinstance(value, str): raise TypeError(f"Expected str for {base.__name__} deserialization, got {type(value).__name__}") return ArrowSerializableDataclass._convert_value_for_deserialization(value, type_hint, ipc_validation) diff --git a/vgi_rpc/utils.py b/vgi_rpc/utils.py index 4c614e4..8ca1ced 100644 --- a/vgi_rpc/utils.py +++ b/vgi_rpc/utils.py @@ -380,7 +380,12 @@ def deserialize_record_batch( data: bytes, ipc_validation: IpcValidation = IpcValidation.FULL, ) -> tuple[pa.RecordBatch, pa.KeyValueMetadata | None]: - """Deserialize bytes back to a RecordBatch with custom metadata. + """Deserialize a complete single-batch IPC value with custom metadata. + + The batch may end at physical EOF or at a legacy/modern Arrow EOS marker. + Additional messages, truncated endings and trailing bytes are rejected. + Empty bytes and schema-only streams remain valid empty nested values. + This bounded-value check does not apply to concatenated transport streams. Args: data: Bytes containing a serialized RecordBatch in Arrow IPC stream format. @@ -391,7 +396,7 @@ def deserialize_record_batch( if no custom metadata was attached to the batch. Raises: - IPCError: If no batches are found in a (non-empty) IPC stream. + IPCError: If the value has extra messages or invalid trailing framing. """ # A 0-length buffer means "empty" — a lenient producer (the Rust/Go/Java SDKs, @@ -401,10 +406,13 @@ def deserialize_record_batch( # choking in open_stream ("was null or length 0"). if not data: return pa.record_batch([], schema=pa.schema([])), None - # Hand pyarrow the bytes directly: wrapping them in a BufferReader first - # costs an extra object and sends pyarrow down its path-like probe - # (_stringify_path) before it gets to the buffer. Measured 7.6us -> 5.7us. - with ValidatedReader(ipc.open_stream(data, options=IPC_READ_OPTIONS), ipc_validation) as reader: + # Keep the byte position: a stream reader's StopIteration alone does not + # prove that a bounded value ended cleanly or that all bytes were consumed. + with ( + pa.BufferReader(data) as source, + ValidatedReader(ipc.open_stream(source, options=IPC_READ_OPTIONS), ipc_validation) as reader, + ): + after_schema = source.tell() try: batch, custom_metadata = reader.read_next_batch_with_custom_metadata() except StopIteration: @@ -412,12 +420,28 @@ def deserialize_record_batch( # (the Rust/Go/Java SDKs, tolerated by the DuckDB C++ client) encode an # empty value — e.g. a scan function with no arguments. Return a 0-row # batch of the declared schema instead of raising. + _check_record_batch_end(data, after_schema) schema = reader.schema return empty_batch(schema), None + _check_record_batch_end(data, source.tell()) return batch, custom_metadata +def _check_record_batch_end(data: bytes, position: int) -> None: + """Require EOF or exactly one complete Arrow EOS after a bounded value.""" + remaining = len(data) - position + if remaining == 0: + return + # Check the length before slicing, so malformed oversized tails are not + # copied and a second record batch is never decoded merely to reject it. + if remaining == 4 and data[position:] == b"\x00\x00\x00\x00": + return + if remaining == 8 and data[position:] == b"\xff\xff\xff\xff\x00\x00\x00\x00": + return + raise IPCError("Expected a single complete record batch with no trailing data") + + def read_single_record_batch( stream: IOBase, context: str = "batch", @@ -1625,6 +1649,10 @@ def _convert_value_for_deserialization( def deserialize_from_bytes(cls, data: bytes, ipc_validation: IpcValidation = IpcValidation.FULL) -> Self: """Deserialize an instance from Arrow IPC bytes. + RPC parameters and unary results use this hook, including subclass + overrides. A custom codec can reject oversized raw bytes before Arrow + decoding; it is responsible for validating its complete envelope. + Args: data: Arrow IPC stream bytes containing a single-row RecordBatch. ipc_validation: Validation level for the deserialized batch. @@ -1634,7 +1662,7 @@ def deserialize_from_bytes(cls, data: bytes, ipc_validation: IpcValidation = Ipc Raises: ValueError: If the batch is invalid (wrong row count or missing fields). - IPCError: If the IPC stream is malformed or truncated. + IPCError: If the IPC value contains extra messages or invalid trailing framing. TypeError: If a field value has an unexpected type during conversion. KeyError: If an Enum name cannot be resolved. From 99b3090cade6057ae9639fbebde1b3135e336535 Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Tue, 27 Oct 2026 14:29:21 -0400 Subject: [PATCH 2/3] fix: classify truncated in-memory IPC bodies as malformed values --- tests/test_single_record_framing.py | 2 +- vgi_rpc/utils.py | 42 ++++++++++++++++------------- 2 files changed, 25 insertions(+), 19 deletions(-) diff --git a/tests/test_single_record_framing.py b/tests/test_single_record_framing.py index c9ee9e7..e73ed1e 100644 --- a/tests/test_single_record_framing.py +++ b/tests/test_single_record_framing.py @@ -250,7 +250,7 @@ def test_http_rejects_bad_typed_value_as_client_error(kind: str) -> None: assert impl.calls == [42] -@pytest.mark.parametrize("kind", ["extra_batch", "trailing_byte", "truncated_eos_4"]) +@pytest.mark.parametrize("kind", ["extra_batch", "trailing_byte", "truncated_eos_4", "truncated_body"]) def test_client_classifies_malformed_typed_result_as_transport_error(kind: str) -> None: """Typed result corruption keeps the client's established RpcError contract.""" schema = rpc_methods(_ResultProtocol)["retrieve"].result_schema diff --git a/vgi_rpc/utils.py b/vgi_rpc/utils.py index 8ca1ced..c2aa1f5 100644 --- a/vgi_rpc/utils.py +++ b/vgi_rpc/utils.py @@ -408,24 +408,30 @@ def deserialize_record_batch( return pa.record_batch([], schema=pa.schema([])), None # Keep the byte position: a stream reader's StopIteration alone does not # prove that a bounded value ended cleanly or that all bytes were consumed. - with ( - pa.BufferReader(data) as source, - ValidatedReader(ipc.open_stream(source, options=IPC_READ_OPTIONS), ipc_validation) as reader, - ): - after_schema = source.tell() - try: - batch, custom_metadata = reader.read_next_batch_with_custom_metadata() - except StopIteration: - # A schema-only stream (no batch message) is how lenient producers - # (the Rust/Go/Java SDKs, tolerated by the DuckDB C++ client) encode an - # empty value — e.g. a scan function with no arguments. Return a 0-row - # batch of the declared schema instead of raising. - _check_record_batch_end(data, after_schema) - schema = reader.schema - return empty_batch(schema), None - - _check_record_batch_end(data, source.tell()) - return batch, custom_metadata + try: + with ( + pa.BufferReader(data) as source, + ValidatedReader(ipc.open_stream(source, options=IPC_READ_OPTIONS), ipc_validation) as reader, + ): + after_schema = source.tell() + try: + batch, custom_metadata = reader.read_next_batch_with_custom_metadata() + except StopIteration: + # A schema-only stream (no batch message) is how lenient producers + # (the Rust/Go/Java SDKs, tolerated by the DuckDB C++ client) encode an + # empty value — e.g. a scan function with no arguments. Return a 0-row + # batch of the declared schema instead of raising. + _check_record_batch_end(data, after_schema) + schema = reader.schema + return empty_batch(schema), None + + _check_record_batch_end(data, source.tell()) + return batch, custom_metadata + except OSError as exc: + # Arrow reports a truncated message body as a plain OSError. This + # reader only consumes in-memory bytes, so this is malformed IPC, not + # an external I/O failure. Keep custom byte-codec exceptions untouched. + raise IPCError(f"Error reading record batch IPC value: {exc}") from exc def _check_record_batch_end(data: bytes, position: int) -> None: From da0b513a9d569f2a2b72290b138632b3f3e1e9e8 Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Tue, 27 Oct 2026 15:07:46 -0400 Subject: [PATCH 3/3] release: prepare 0.52.3 with complete typed IPC validation --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c6a0edb..e671e49 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "vgi-rpc" -version = "0.52.2" +version = "0.52.3" description = "Vector Gateway Interface - RPC framework based on Apache Arrow" readme = "README.md" requires-python = ">=3.13" diff --git a/uv.lock b/uv.lock index 1e331ff..ebca0a4 100644 --- a/uv.lock +++ b/uv.lock @@ -3047,7 +3047,7 @@ wheels = [ [[package]] name = "vgi-rpc" -version = "0.52.2" +version = "0.52.3" source = { editable = "." } dependencies = [ { name = "docstring-parser" },