|
| 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 |
0 commit comments