Skip to content

Commit ed1da23

Browse files
rustyconoverclaude
andcommitted
release: 0.24.0 — stop carrying the attach decrypted in the call state
`VgiCallState` held `plaintext_attach`, the *unwrapped* attach (`uuid(16) || catalog_bytes`), so each turn could shard storage on its UUID without re-opening the auth-scoped seal. Two problems with that. It was redundant: the sealed envelope is already reachable in the same object at `init_call.bind_call.attach_opaque_data`. The call state carried both the ciphertext and its plaintext. And serialized state does not stay sealed at rest. The access log records the *decrypted* state at DEBUG as a deliberate audit artifact (`vgi_rpc.rpc._server`), and on the init turn the recorded response state is `call_state || cursor` — so a decrypted attach base64s catalog credentials into the log. That is exactly what `loggable_attach_options` exists to prevent elsewhere, and attach options routinely carry passwords and tokens. It also sat decrypted in the transport's call-state cache. The comment defending it claimed "the auth-scoped seal can't be reopened on a later, possibly different-auth, produce/finalize turn". That is not true, and it is now a test: continuation turns are ordinary dispatches carrying the caller's identity, so driving a real multi-turn HTTP stream shows the identity present and the envelope re-opening on every turn. The plaintext now lives in a `Transient` field, populated from the sealed envelope per HTTP turn and at construction under pipe (which never calls rehydrate). Re-opening restores a per-turn identity check that the cached copy had quietly skipped. Re-opening cost 42us per turn, ~13% of a turn, which is why the copy existed. So memoize the open on the worker, keyed by `(envelope, caller identity)`: warm unwrap 35-42us -> 1.19us catalog stream 42us -> 1.54us per turn non-catalog stream 0.14us (no envelope; unchanged) The identity half of that key is load-bearing. An envelope's AAD is the only thing binding it to a principal, and a cache hit skips the AEAD — so keyed on the envelope alone (hashed or not), a stolen envelope becomes a bearer token for anyone. Keyed on both, a mismatched caller misses and takes the open, which fails as it should. Only successes are stored, so a rejection is never cached into a later hit. Tests cover cross-principal and anonymous callers against a warm cache, non-poisoning, and the eviction bound. The cache holds plaintext in memory, which is unavoidable — the worker cannot resolve a storage shard without it — and widens nothing that opening the envelope already did. The property being protected is that it never reaches serialized state or logs. Also requires vgi-rpc 0.33.0. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MS3MkkHgE8pA3Aty7BnZKE
1 parent f360ed0 commit ed1da23

5 files changed

Lines changed: 480 additions & 30 deletions

File tree

‎pyproject.toml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[project]
22
name = "vgi-python"
3-
version = "0.23.0"
3+
version = "0.24.0"
44
description = "Vector Gateway Interface - Connect DuckDB to external programs via Apache Arrow"
55
readme = "README.md"
66
keywords = [
@@ -40,7 +40,7 @@ dependencies = [
4040
"pyarrow",
4141
"typer>=0.9",
4242
"platformdirs",
43-
"vgi-rpc>=0.32.0",
43+
"vgi-rpc>=0.33.0",
4444
"httpx>=0.24",
4545
]
4646

‎tests/test_call_state_attach.py‎

Lines changed: 305 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,305 @@
1+
# Copyright 2025, 2026 Query Farm LLC - https://query.farm
2+
3+
"""The stream's call state must not carry the attach in decrypted form.
4+
5+
``attach_opaque_data`` routinely carries catalog credentials. It travels
6+
sealed, but serialized *state* does not stay sealed at rest: the access log
7+
records the decrypted state at DEBUG as an audit artifact
8+
(``vgi_rpc.rpc._server``), and on the init turn the recorded response state is
9+
``call_state || cursor``. So a decrypted attach inside
10+
:class:`~vgi.protocol.VgiCallState` would base64 credentials into the log —
11+
precisely what ``loggable_attach_options`` exists to prevent elsewhere — and
12+
would also sit decrypted in the transport's call-state cache.
13+
14+
The sealed envelope is already reachable at
15+
``init_call.bind_call.attach_opaque_data``, so each turn re-opens that instead.
16+
These tests pin both halves: that the plaintext never serializes, and that
17+
re-opening actually works on turns after the first (the reason the decrypted
18+
copy was there originally).
19+
"""
20+
21+
from __future__ import annotations
22+
23+
import contextlib
24+
from collections.abc import Iterator
25+
from dataclasses import dataclass
26+
from typing import Annotated, Any
27+
28+
import pyarrow as pa
29+
import pytest
30+
from vgi_rpc.rpc import AuthContext, OutputCollector, RpcServer
31+
from vgi_rpc.rpc._common import _current_transport, _TransportContext
32+
from vgi_rpc.utils import ArrowSerializableDataclass
33+
34+
from vgi.arguments import Arg, Arguments
35+
from vgi.invocation import BindResponse
36+
from vgi.protocol import (
37+
BindRequest,
38+
FunctionType,
39+
GlobalInitResponse,
40+
InitRequest,
41+
VgiCallState,
42+
VgiProtocol,
43+
)
44+
from vgi.table_function import (
45+
BindParams,
46+
ProcessParams,
47+
TableFunctionGenerator,
48+
init_single_worker,
49+
)
50+
from vgi.worker import Worker
51+
52+
_KEY = b"\x11" * 32
53+
_SECRET = b"PGPASSWORD=hunter2"
54+
_UUID = b"\xab" * 16
55+
_DOMAIN = "test-domain"
56+
_PRINCIPAL = "alice"
57+
58+
#: Every ``attach_opaque_data`` the probe saw, one entry per produce tick.
59+
_SEEN: list[bytes | None] = []
60+
61+
62+
@contextlib.contextmanager
63+
def _as_principal(domain: str | None, principal: str | None) -> Iterator[None]:
64+
"""Simulate the transport-set auth context for a single request."""
65+
auth = AuthContext(domain=domain, authenticated=principal is not None, principal=principal)
66+
token = _current_transport.set(_TransportContext(auth=auth))
67+
try:
68+
yield
69+
finally:
70+
_current_transport.reset(token)
71+
72+
73+
@dataclass(kw_only=True)
74+
class _ProbeState(ArrowSerializableDataclass):
75+
"""Remaining tick count."""
76+
77+
remaining: int = 0
78+
79+
80+
@dataclass(frozen=True)
81+
class _ProbeArgs:
82+
"""Arguments for the probe."""
83+
84+
count: Annotated[int, Arg(0, doc="ticks", ge=0)]
85+
86+
87+
@init_single_worker
88+
class _AttachProbeFunction(TableFunctionGenerator[_ProbeArgs, _ProbeState]):
89+
"""Records the attach bytes visible on each produce tick.
90+
91+
Attributes:
92+
FunctionArguments: Argument dataclass bound to this function.
93+
94+
"""
95+
96+
FunctionArguments = _ProbeArgs
97+
98+
class Meta:
99+
"""Metadata for the probe."""
100+
101+
name = "attach_probe"
102+
description = "records params.attach_opaque_data per tick"
103+
104+
@classmethod
105+
def on_bind(cls, params: BindParams[_ProbeArgs]) -> BindResponse:
106+
"""A single int64 column."""
107+
return BindResponse(output_schema=pa.schema([("n", pa.int64())]))
108+
109+
@classmethod
110+
def initial_state(cls, params: ProcessParams[_ProbeArgs]) -> _ProbeState:
111+
"""Start with the requested tick count."""
112+
return _ProbeState(remaining=params.args.count)
113+
114+
@classmethod
115+
def process(cls, params: ProcessParams[_ProbeArgs], state: _ProbeState, out: OutputCollector) -> None:
116+
"""Record the attach this turn sees, then emit one row."""
117+
_SEEN.append(bytes(params.attach_opaque_data) if params.attach_opaque_data is not None else None)
118+
if state.remaining <= 0:
119+
out.finish()
120+
return
121+
state.remaining -= 1
122+
out.emit(pa.RecordBatch.from_pydict({"n": [state.remaining]}, schema=params.output_schema))
123+
124+
125+
class _ProbeWorker(Worker):
126+
"""Worker exposing only the attach probe."""
127+
128+
functions = [_AttachProbeFunction]
129+
130+
131+
def _sealed_attach() -> bytes:
132+
"""Seal ``uuid || secret`` for the test principal, as a real attach would be."""
133+
worker = _ProbeWorker(quiet=True)
134+
worker._signing_key = _KEY
135+
with _as_principal(_DOMAIN, _PRINCIPAL):
136+
return bytes(worker._seal_attach(_UUID + _SECRET))
137+
138+
139+
class TestCallStateCarriesNoPlaintextAttach:
140+
"""The serialized call state must never contain the decrypted attach."""
141+
142+
def test_plaintext_attach_is_not_a_serialized_field(self) -> None:
143+
"""``VgiCallState`` has no field holding the decrypted attach."""
144+
import dataclasses
145+
146+
names = {f.name for f in dataclasses.fields(VgiCallState)}
147+
assert names == {"init_call", "init_response"}
148+
149+
def test_secret_absent_from_serialized_call_state(self) -> None:
150+
"""Serializing a call state must not emit the attach plaintext.
151+
152+
The sealed envelope must still be reachable, since re-opening it is
153+
how each turn now obtains the plaintext.
154+
"""
155+
sealed = _sealed_attach()
156+
bind = BindRequest(
157+
function_name="attach_probe",
158+
arguments=Arguments(positional=(pa.scalar(1),)),
159+
function_type=FunctionType.TABLE,
160+
input_schema=pa.schema([]),
161+
attach_opaque_data=sealed,
162+
)
163+
call_state = VgiCallState(
164+
init_call=InitRequest(bind_call=bind, output_schema=pa.schema([("n", pa.int64())])),
165+
init_response=GlobalInitResponse(execution_id=b"exec"),
166+
)
167+
raw = call_state.serialize_to_bytes()
168+
169+
assert _SECRET not in raw, "decrypted attach must not reach the wire or the access log"
170+
assert _UUID not in raw, "the storage-sharding UUID is part of the same plaintext"
171+
assert sealed in raw, "the sealed envelope must remain, since turns re-open it"
172+
173+
174+
class TestAttachSurvivesLaterTurns:
175+
"""Re-opening the seal must work past the init turn, not only at init."""
176+
177+
@pytest.fixture(autouse=True)
178+
def _reset(self) -> Iterator[None]:
179+
_SEEN.clear()
180+
yield
181+
_SEEN.clear()
182+
183+
def test_every_turn_sees_the_attach(self) -> None:
184+
"""A multi-turn HTTP stream sees the catalog bytes on every tick.
185+
186+
This is the property the decrypted copy was there to guarantee: the
187+
original comment held that "the auth-scoped seal can't be reopened on
188+
a later, possibly different-auth, produce/finalize turn". Continuation
189+
turns are ordinary dispatches carrying the caller's identity, so it
190+
can — and a regression would show up here as ``None`` (or a raise)
191+
on every tick after the first.
192+
"""
193+
from vgi_rpc.http import http_connect, make_sync_client
194+
195+
sealed = _sealed_attach()
196+
197+
def authenticate(req: Any) -> AuthContext:
198+
return AuthContext(domain=_DOMAIN, authenticated=True, principal=_PRINCIPAL)
199+
200+
worker = _ProbeWorker(quiet=True)
201+
worker._signing_key = _KEY
202+
client = make_sync_client(
203+
RpcServer(VgiProtocol, worker, enable_describe=False),
204+
token_key=_KEY,
205+
authenticate=authenticate,
206+
)
207+
208+
bind = BindRequest(
209+
function_name="attach_probe",
210+
arguments=Arguments(positional=(pa.scalar(4),)),
211+
function_type=FunctionType.TABLE,
212+
input_schema=pa.schema([]),
213+
attach_opaque_data=sealed,
214+
)
215+
with http_connect(VgiProtocol, client=client, compression_level=None) as proxy: # type: ignore[type-abstract]
216+
resp = proxy.bind(request=bind)
217+
stream = proxy.init(
218+
request=InitRequest(
219+
bind_call=bind,
220+
output_schema=resp.output_schema,
221+
bind_opaque_data=resp.opaque_data,
222+
)
223+
)
224+
rows = sum(ab.batch.num_rows for ab in stream)
225+
stream.close()
226+
227+
assert rows == 4
228+
assert len(_SEEN) > 1, "need more than the init turn for this to prove anything"
229+
# The body sees the catalog bytes with the framework UUID stripped.
230+
assert all(seen == _SECRET for seen in _SEEN), _SEEN
231+
232+
233+
class TestAttachUnwrapCache:
234+
"""Memoizing the AEAD open must not become a way around the AEAD.
235+
236+
Opening an attach is a pure function of (envelope, key, AAD), so repeats
237+
are cacheable — and worth caching, since a catalog stream re-opens the
238+
same envelope every turn. But the AAD is the *only* thing binding an
239+
envelope to a principal, and a cache hit skips it. Keyed on the envelope
240+
alone, anyone holding a stolen envelope would be handed its plaintext.
241+
"""
242+
243+
@staticmethod
244+
def _worker() -> _ProbeWorker:
245+
worker = _ProbeWorker(quiet=True)
246+
worker._signing_key = _KEY
247+
return worker
248+
249+
@staticmethod
250+
def _auth(principal: str | None) -> AuthContext:
251+
return AuthContext(domain=_DOMAIN, authenticated=principal is not None, principal=principal)
252+
253+
def test_same_principal_hits_and_matches(self) -> None:
254+
"""A warm unwrap returns exactly what the cold one did."""
255+
worker = self._worker()
256+
sealed = _sealed_attach()
257+
alice = self._auth(_PRINCIPAL)
258+
259+
cold = worker._unwrap_attach_full_with_auth(sealed, alice)
260+
warm = worker._unwrap_attach_full_with_auth(sealed, alice)
261+
assert cold == _UUID + _SECRET
262+
assert warm == cold
263+
264+
def test_other_principal_rejected_even_after_cache_is_warm(self) -> None:
265+
"""A second principal replaying the envelope must not get a hit.
266+
267+
This is the regression that matters: warming on behalf of one
268+
principal must not turn the envelope into a bearer token for the next.
269+
"""
270+
worker = self._worker()
271+
sealed = _sealed_attach()
272+
273+
assert worker._unwrap_attach_full_with_auth(sealed, self._auth(_PRINCIPAL)) == _UUID + _SECRET
274+
with pytest.raises(ValueError):
275+
worker._unwrap_attach_full_with_auth(sealed, self._auth("bob"))
276+
277+
def test_anonymous_rejected_even_after_cache_is_warm(self) -> None:
278+
"""An unauthenticated caller gets no hit on an authenticated entry."""
279+
worker = self._worker()
280+
sealed = _sealed_attach()
281+
282+
assert worker._unwrap_attach_full_with_auth(sealed, self._auth(_PRINCIPAL)) == _UUID + _SECRET
283+
with pytest.raises(ValueError):
284+
worker._unwrap_attach_full_with_auth(sealed, self._auth(None))
285+
286+
def test_rejection_is_not_cached_as_a_hit(self) -> None:
287+
"""A failed open must not poison the entry for the rightful principal."""
288+
worker = self._worker()
289+
sealed = _sealed_attach()
290+
291+
with pytest.raises(ValueError):
292+
worker._unwrap_attach_full_with_auth(sealed, self._auth("bob"))
293+
assert worker._unwrap_attach_full_with_auth(sealed, self._auth(_PRINCIPAL)) == _UUID + _SECRET
294+
295+
def test_cache_is_bounded(self) -> None:
296+
"""The map cannot grow without limit as attachments come and go."""
297+
from vgi.worker import _AttachUnwrapCache
298+
299+
cache = _AttachUnwrapCache(max_size=4)
300+
for i in range(20):
301+
cache.put((bytes([i]), b"id"), bytes([i]))
302+
assert len(cache._entries) == 4
303+
# Least-recently-used entries evicted; the newest survive.
304+
assert cache.get((bytes([19]), b"id")) == bytes([19])
305+
assert cache.get((bytes([0]), b"id")) is None

‎uv.lock‎

Lines changed: 5 additions & 5 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)