diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d577182..d879247 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -161,7 +161,7 @@ jobs: env: ARBITER_GO_CLIENT_SMOKE_OUTDIR: .ci/go-client-smoke GOCACHE: ${{ runner.temp }}/go-build - run: nox -s unit -- server/tests/integration/test_credential_peer.py server/tests/integration/test_provisional_peer.py server/tests/integration/test_pairing_peer.py server/tests/integration/test_enrollment_cli_peer.py server/tests/integration/test_control_tls_peer.py server/tests/integration/test_enrollment_receipt_peer.py server/tests/integration/test_activation_restart_peer.py server/tests/integration/test_installed_enrollment.py + run: nox -s unit -- server/tests/integration/test_credential_peer.py server/tests/integration/test_provisional_peer.py server/tests/integration/test_pairing_peer.py server/tests/integration/test_enrollment_cli_peer.py server/tests/integration/test_control_tls_peer.py server/tests/integration/test_authentication_peer.py server/tests/integration/test_enrollment_receipt_peer.py server/tests/integration/test_activation_restart_peer.py server/tests/integration/test_installed_enrollment.py - name: Verify native local-control rejection across OS users if: ${{ startsWith(matrix.platform_id, 'linux-') || startsWith(matrix.platform_id, 'macos-') }} diff --git a/client/go-cli/internal/identity/authentication_peer_test.go b/client/go-cli/internal/identity/authentication_peer_test.go new file mode 100644 index 0000000..db87ab1 --- /dev/null +++ b/client/go-cli/internal/identity/authentication_peer_test.go @@ -0,0 +1,107 @@ +package identity + +import ( + "bytes" + "crypto/tls" + "encoding/json" + "fmt" + "io" + "os" + "testing" + "time" + + "github.com/gorilla/websocket" +) + +// A bounded synthetic peer adapter can batch two TLS records or split records +// across WS messages. Production still uses its existing controlSocket. +type streamPeerSocket struct { + *controlSocket + batch, fragment bool + pending []byte +} + +func (c *streamPeerSocket) Write(p []byte) (int, error) { + if c.batch { + if len(c.pending)+len(p) > 65536 { + return 0, io.ErrShortBuffer + } + c.pending = append(c.pending, p...) + return len(p), nil + } + if !c.fragment { + return c.controlSocket.Write(p) + } + n := len(p) + for len(p) > 0 { + size := min(2049, len(p)) + if err := c.WriteMessage(websocket.BinaryMessage, p[:size]); err != nil { + return n - len(p), err + } + p = p[size:] + } + return n, nil +} + +func TestPythonAuthenticationStreamPeer(t *testing.T) { + path := os.Getenv("ARBITER_TEST_AUTH_STREAM") + if path == "" { + t.Skip("private Python stream peer only") + } + var fixture struct { + URL, ServerID string + Authority []byte + Fragment bool + } + raw, err := os.ReadFile(path) + if err != nil || json.Unmarshal(raw, &fixture) != nil { + t.Fatal("invalid synthetic peer fixture") + } + dialer := websocket.Dialer{HandshakeTimeout: 5 * time.Second, + TLSClientConfig: &tls.Config{InsecureSkipVerify: true}} + ws, _, err := dialer.Dial(fixture.URL, nil) + if err != nil { + t.Fatal(err) + } + defer ws.Close() + socket := &streamPeerSocket{controlSocket: &controlSocket{Conn: ws}, fragment: fixture.Fragment} + if err := socket.SetDeadline(time.Now().Add(8 * time.Second)); err != nil { + t.Fatal(err) + } + conn := tls.Client(socket, &tls.Config{MinVersion: tls.VersionTLS13, MaxVersion: tls.VersionTLS13, + NextProtos: []string{controlALPN}, SessionTicketsDisabled: true, InsecureSkipVerify: true, + VerifyConnection: func(state tls.ConnectionState) error { + if state.Version != tls.VersionTLS13 || state.NegotiatedProtocol != controlALPN || + state.DidResume || len(state.PeerCertificates) != 1 { + return errControl + } + _, err := VerifyCredential(state.PeerCertificates[0].Raw, + Identity{Role: ServerTLS, SubjectID: fixture.ServerID}, + &Issuer{DER: fixture.Authority, ID: fixture.ServerID}, time.Now().UTC()) + return err + }}) + if err := conn.Handshake(); err != nil { + t.Fatal(err) + } + payload := bytes.Repeat([]byte("bounded-stream!!"), 2048) // 32 KiB / two chunks. + for round := range 16 { + payload[0] = byte(round) + socket.batch = !fixture.Fragment + for offset := 0; offset < len(payload); offset += 16384 { + if _, err := conn.Write(payload[offset : offset+16384]); err != nil { + t.Fatal(err) + } + } + if socket.batch { + if err := ws.WriteMessage(websocket.BinaryMessage, socket.pending); err != nil { + t.Fatal(err) + } + socket.pending = nil + } + got := make([]byte, len(payload)) + if _, err := io.ReadFull(conn, got); err != nil || !bytes.Equal(got, payload) { + t.Fatalf("stream round %d: %v", round, err) + } + } + fmt.Println("ARBITER_STREAM_PEER=pass") +} diff --git a/docs/authentication-cli.md b/docs/authentication-cli.md new file mode 100644 index 0000000..088e779 --- /dev/null +++ b/docs/authentication-cli.md @@ -0,0 +1,7 @@ +# Native authentication delivery status + +The shared bounded TLS stream is an internal transport primitive. Its private +Python/Go interoperability peer does not enable ordinary authenticated API +requests or add a user command. Existing enrollment and fixed control behavior +remain available. The approved ordinary native API and enrollment selection +are delivered by the following cuts in the authentication implementation plan. diff --git a/server/news/+auth-3-1-1.feature.md b/server/news/+auth-3-1-1.feature.md new file mode 100644 index 0000000..59944e0 --- /dev/null +++ b/server/news/+auth-3-1-1.feature.md @@ -0,0 +1,4 @@ +Extract the shared bounded native TLS stream pump and carrier credit accounting, +with private Python/Go streaming interoperability tests. Existing enrollment and +fixed control behavior remain unchanged; ordinary authenticated API access is +not enabled by this internal transport step. diff --git a/server/src/arbiter_server/control_tls.py b/server/src/arbiter_server/control_tls.py index 4f107e3..3255bf4 100644 --- a/server/src/arbiter_server/control_tls.py +++ b/server/src/arbiter_server/control_tls.py @@ -2,15 +2,14 @@ from __future__ import annotations -from collections.abc import Callable from pathlib import Path -from typing import Any from cryptography import x509 from OpenSSL import SSL from .pairing_carrier import PairingProtocol from .pairing_identity import load_pairing_identity +from .native_stream import TLSStream ALPN = b"arbiter-control-v1" MAX_CONTROL = 65536 @@ -42,7 +41,7 @@ def select(connection: SSL.Connection, protocols: list[bytes]) -> bytes: return connection -class ControlTLS: +class ControlTLS(TLSStream): """One bounded exchange over the existing carrier and crypto ownership. Every BIO/crypto call is serialized through the carrier's shared worker pool. @@ -50,9 +49,6 @@ class ControlTLS: data is drained before accepting more input; no background read-ahead exists. """ - def __init__(self, carrier: PairingProtocol, connection: SSL.Connection) -> None: - self.carrier, self.connection = carrier, connection - @classmethod async def accept(cls, carrier: PairingProtocol, root: Path) -> ControlTLS: assert carrier.lease is not None @@ -68,53 +64,6 @@ async def accept(cls, carrier: PairingProtocol, root: Path) -> ControlTLS: carrier.begin_control() return stream - async def _call(self, operation: Callable[[], Any]) -> Any: - carrier, connection = self.carrier, self.connection - assert carrier.lease is not None - first = True - result: Any = None - needs_input = False - while True: - - def step() -> tuple[Any, bool, bytes | None]: - operation_result: Any = None - operation_needs_input = False - if first: - try: - operation_result = operation() - except SSL.WantReadError: - operation_needs_input = True - try: - # Leave headroom for the encrypted WebSocket frame and - # the transport's own write-buffer copy. - chunk = connection.bio_read(max(1, carrier.frame_limit // 2)) - except SSL.WantReadError: - chunk = None - carrier.update_control_tls_pending(connection.pending()) - return operation_result, operation_needs_input, chunk - - step_result, step_needs_input, chunk = await carrier.responder._run( - carrier.lease, step - ) - if first: - result, needs_input = step_result, step_needs_input - first = False - if chunk is not None: - await carrier.send(chunk) - if not needs_input: - if chunk is None: - return result - continue - if chunk is not None: - continue - data = await carrier.receive() - carrier.reserve_control_tls_input(len(data)) - await carrier.responder._run( - carrier.lease, lambda: connection.bio_write(data) - ) - del data # BIO owns the input before the next carrier read. - first = True - async def receive(self) -> bytes: async def exact(size: int) -> bytes: result = bytearray() @@ -138,12 +87,9 @@ async def send(self, message: bytes) -> None: raise ValueError("Control exceeded limit") await self._call(lambda: self.connection.send(len(message).to_bytes(4, "big"))) view = memoryview(message) - try: - while view: - written = await self._call(lambda: self.connection.send(view)) - view = view[written:] - finally: - view.release() + while view: + written = await self._call(lambda: self.connection.send(view)) + view = view[written:] def exporter(self) -> bytes: return self.connection.export_keying_material( diff --git a/server/src/arbiter_server/native_stream.py b/server/src/arbiter_server/native_stream.py new file mode 100644 index 0000000..11f0bd6 --- /dev/null +++ b/server/src/arbiter_server/native_stream.py @@ -0,0 +1,195 @@ +"""Shared private carrier credits and serialized memory-BIO stream pump. + +This module does not select credentials, authorize callers or expose a route. +Stream consumers must consume each returned chunk before requesting another, +and charge any retained application copies to their own bounded owner. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +from OpenSSL import SSL + +if TYPE_CHECKING: + from .pairing_carrier import PairingProtocol + +CHUNK = 16384 +RECORD_RESERVE = 36 * 1024 + + +class CarrierCredits: + """One carrier ledger; fixed control remains cumulative, streams move.""" + + def __init__(self, maximum: int, lock: Any) -> None: + self.maximum = maximum + self.lock = lock + self.fed = self.pending = self.queued = 0 + self.output = self.application = self.adapter = 0 + self.streaming = False + + def check(self, extra: int = 0) -> None: + if ( + extra < 0 + or self.fed + + self.pending + + 2 * self.queued + + self.output + + self.application + + self.adapter + + extra + > self.maximum + ): + raise ValueError("Control TLS pending data exceeded limit") + + def reserve_input(self, size: int) -> None: + with self.lock: + if size < 0 or ( + not self.streaming and self.fed + size > max(1, self.maximum // 2) + ): + raise ValueError("Control TLS input exceeded limit") + charge = size * (2 if self.streaming else 1) + self.check(charge) + self.fed += charge + + def update_pending(self, size: int, *, wants_input: bool = False) -> None: + with self.lock: + if size < 0: + raise ValueError("Control TLS pending data exceeded limit") + self.pending = size + if self.streaming and wants_input and not size: + # pending() cannot reveal an incomplete encrypted record. Never + # discard its conservative record-and-copy reserve between feeds. + self.fed = RECORD_RESERVE + self.check() + + +class TLSStream: + """One in-flight operation, on the carrier's existing bounded crypto pool. + + Credits are not cleared by cancellation: the actual worker still owns its + inputs and lease. The closed carrier is discarded only after that owner + ends. Outgoing BIO data drains before any new input or application write. + """ + + def __init__(self, carrier: PairingProtocol, connection: SSL.Connection) -> None: + self.carrier, self.connection = carrier, connection + self._calling = False + self._input = b"" + self._offset = 0 + + async def _call( + self, operation: Callable[[], Any], *, application: int | None = None + ) -> Any: + with self.carrier.credits.lock: + if self._calling or ( + self.carrier.lease is not None and self.carrier.lease._busy + ): + raise ValueError("Concurrent TLS operation") + if application is not None: + self.carrier.credits.application = application + self.carrier.credits.check() + self._calling = True + try: + return await self._pump(operation) + finally: + self._calling = False + + async def _pump(self, operation: Callable[[], Any]) -> Any: + carrier, connection = self.carrier, self.connection + assert carrier.lease is not None + first, drained = True, False + result: Any = None + wants: str | None = None + while True: + limit = max(1, min(CHUNK, carrier.frame_limit // 2)) + with carrier.credits.lock: + if carrier.credits.streaming: + carrier.credits.output = 3 * limit + assert carrier.transport is not None + carrier.credits.check(carrier.transport.get_write_buffer_size()) + + def step() -> tuple[Any, str | None, bytes | None]: + value: Any = None + need: str | None = None + if first: + try: + value = operation() + except SSL.WantReadError: + need = "read" + except SSL.WantWriteError: + need = "write" + try: + chunk = connection.bio_read(limit) + except SSL.WantReadError: + chunk = None + carrier.credits.update_pending( + connection.pending(), wants_input=need == "read" + ) + return value, need, chunk + + value, need, chunk = await carrier.responder._run(carrier.lease, step) + if first: + result, wants = value, need + first = False + if chunk is not None: + drained = True + await carrier.send(chunk) + del chunk + continue + carrier.credits.output = 0 + if wants is None: + return result + if wants == "write": + if not drained: + raise ValueError("TLS write made no progress") + first, drained = True, False + continue + await self._feed() + first, drained = True, False + + async def _feed(self) -> None: + carrier = self.carrier + assert carrier.lease is not None + if not self._input: + self._input = await carrier.receive() + self._offset = 0 + carrier.credits.adapter = len(self._input) + carrier.credits.check() + part = self._input[self._offset : self._offset + CHUNK] + carrier.reserve_control_tls_input(len(part)) + written = await carrier.responder._run( + carrier.lease, lambda: self.connection.bio_write(part) + ) + if not isinstance(written, int) or not 0 < written <= len(part): + raise ValueError("TLS input made no progress") + self._offset += written + if self._offset == len(self._input): + self._input = b"" # Drop the adapter owner before another read. + carrier.credits.adapter = 0 + + async def read(self, maximum: int = CHUNK) -> bytes: + if not 1 <= maximum <= CHUNK: + raise ValueError("TLS application chunk exceeded limit") + return await self._call( + lambda: self.connection.recv(maximum), application=2 * maximum + ) + + async def write(self, data: bytes) -> None: + if not 1 <= len(data) <= CHUNK: + raise ValueError("TLS application chunk exceeded limit") + view = memoryview(data) + while view: + written = await self._call( + lambda: self.connection.send(view), application=2 * len(data) + ) + if not isinstance(written, int) or not 0 < written <= len(view): + raise ValueError("TLS output made no progress") + view = view[written:] + + async def close(self) -> None: + try: + await self._call(self.connection.shutdown) + finally: + self.carrier.abort() diff --git a/server/src/arbiter_server/pairing_carrier.py b/server/src/arbiter_server/pairing_carrier.py index a0b4116..f16bf81 100644 --- a/server/src/arbiter_server/pairing_carrier.py +++ b/server/src/arbiter_server/pairing_carrier.py @@ -15,6 +15,7 @@ from wsproto.utilities import LocalProtocolError, RemoteProtocolError from .limits import NativeTransportLimitsConfig +from .native_stream import CarrierCredits, RECORD_RESERVE from .pairing import PairingLease, PairingResponder, PairingSession, ReplayedHandshake PAIRING_PATH = "/api/v1/enrollment/requests/" @@ -98,9 +99,7 @@ def __init__( # TLS keeps a copy of accepted ciphertext while it may also retain # decoded plaintext. Keep a conservative half-budget owner for that # input so both native sides fit inside the inclusive carrier budget. - self._control_tls_fed = 0 - self._control_tls_pending = 0 - self._control_queued = 0 + self.credits = CarrierCredits(limits.max_buffered_bytes, responder._lock) self.frame_limit = min( limits.max_noise_frame_bytes, limits.max_websocket_message_bytes, @@ -173,20 +172,25 @@ def data_received(self, data: bytes) -> None: if len(data) > self.limits.max_buffered_bytes: self.abort() return - if self.received > self.limits.max_handshake_bytes or ( - not self.accepted and self.received > self.limits.max_headers_bytes - ): + if ( + not self.credits.streaming + and self.received > self.limits.max_handshake_bytes + ) or (not self.accepted and self.received > self.limits.max_headers_bytes): self.abort() return try: for offset in range(0, len(data), chunk): if self.closed: break - self.conn.receive_data(data[offset : offset + chunk]) - for event in self.conn.events(): - if self.closed: - break - self._event(event) + with self.credits.lock: + if self.credits.streaming: + incoming = min(chunk, len(data) - offset) + self.credits.check(3 * (len(self.buffer) + incoming)) + self.conn.receive_data(data[offset : offset + chunk]) + for event in self.conn.events(): + if self.closed: + break + self._event(event) except (RemoteProtocolError, LocalProtocolError, ValueError): self.abort() @@ -223,15 +227,7 @@ def _event(self, event: events.Event) -> None: return frame_size = len(self.buffer) + len(event.data) if self.control: - if ( - self._control_tls_fed - + self._control_tls_pending - + 2 * self._control_queued - + 3 * frame_size - > self.limits.max_buffered_bytes - ): - self.abort() - return + self.credits.check(3 * frame_size) self.buffer.extend(event.data) if event.message_finished: pending = self.queue.get_nowait() if not self.queue.empty() else b"" @@ -239,7 +235,12 @@ def _event(self, event: events.Event) -> None: frame_size = len(self.buffer) projected_live = 2 * pending_size + 3 * frame_size if ( - projected_live + self._control_tls_fed + self._control_tls_pending + projected_live + + self.credits.fed + + self.credits.pending + + self.credits.output + + self.credits.application + + self.credits.adapter > self.limits.max_buffered_bytes ): # pending + buffer + bytes(buffer) + pending+buffer remain @@ -249,7 +250,7 @@ def _event(self, event: events.Event) -> None: return combined = pending + bytes(self.buffer) if self.control: - self._control_queued = len(combined) + self.credits.queued = len(combined) self.queue.put_nowait(combined) self.buffer.clear() assert self.transport is not None @@ -273,12 +274,18 @@ def _event(self, event: events.Event) -> None: def _write(self, event: events.Event) -> None: if self.closed or self.transport is None: raise ValueError("Pairing carrier is closed") + if self.credits.streaming: + payload = getattr(event, "data", getattr(event, "payload", b"")) + self.credits.check( + self.transport.get_write_buffer_size() + len(payload) + 14 + ) wire = self.conn.send(event) if ( - self.sent + len(wire) > self.limits.max_handshake_bytes - or self.transport.get_write_buffer_size() + len(wire) - > self.limits.max_buffered_bytes - ): + not self.credits.streaming + and self.sent + len(wire) > self.limits.max_handshake_bytes + ) or self.transport.get_write_buffer_size() + len( + wire + ) > self.limits.max_buffered_bytes: raise ValueError("Pairing carrier byte limit exceeded") self.sent += len(wire) self.transport.write(wire) @@ -306,7 +313,7 @@ async def receive(self) -> bytes: else: result = await self.queue.get() if self.control: - self._control_queued -= len(result) + self.credits.queued -= len(result) if not self.control: self.transport.resume_reading() return result @@ -323,7 +330,7 @@ def begin_control(self) -> None: if not self.control or self.control_started_once or self.closed: raise ValueError("Control is unavailable") assert self.lease is not None and self.timer is not None - if self._control_queued > self._control_tls_budget: + if self.credits.queued > self._control_tls_budget: raise ValueError("Control TLS input exceeded limit") with self.responder._lock: self.responder._check(self.lease) @@ -333,23 +340,24 @@ def begin_control(self) -> None: self.timer.cancel() self.timer = self.loop.call_later(10.0, cast(Callable[..., object], self.abort)) + def begin_stream(self) -> None: + """Private stream consumer; no public route is registered here.""" + self.begin_control() + self.credits.streaming = True + self.credits.fed = max(self.credits.fed, RECORD_RESERVE) + self.credits.check() + @property def _control_tls_budget(self) -> int: return max(1, self.limits.max_buffered_bytes // 2) def reserve_control_tls_input(self, size: int) -> None: - if ( - not self.control - or size < 0 - or self._control_tls_fed + size > self._control_tls_budget - ): + if not self.control: raise ValueError("Control TLS input exceeded limit") - self._control_tls_fed += size + self.credits.reserve_input(size) def update_control_tls_pending(self, size: int) -> None: - if size < 0 or self._control_tls_fed + size > self.limits.max_buffered_bytes: - raise ValueError("Control TLS pending data exceeded limit") - self._control_tls_pending = size + self.credits.update_pending(size) async def _exchange(self, selector: str | None) -> None: assert self.lease is not None diff --git a/server/tests/integration/test_authentication_peer.py b/server/tests/integration/test_authentication_peer.py new file mode 100644 index 0000000..dd23bd9 --- /dev/null +++ b/server/tests/integration/test_authentication_peer.py @@ -0,0 +1,137 @@ +"""Private real stream peer; ordinary authentication/routing is not enabled.""" + +import asyncio +import base64 +import json +import os +from pathlib import Path +import shutil +import subprocess +from typing import cast + +import pytest +from starlette.applications import Starlette + +from arbiter_server import control_tls +from arbiter_server.authority import initialize_authority +from arbiter_server.main import _generate_self_signed_tls_certificate +from arbiter_server.native_stream import CHUNK, TLSStream +from arbiter_server.pairing_carrier import protocol_factory +from arbiter_server.pairing_identity import load_pairing_identity +from server.tests.unit.test_artifact_http_admission import live_server +from server.tests.unit.test_native_http_protocol import _client +from server.tests.unit.conftest import private_windows_test_parent, tmp_path + + +@pytest.mark.parametrize("fragment", [False, True]) +@pytest.mark.parametrize("outer_tls", [False, True]) +def test_private_python_go_stream( + tmp_path, private_windows_test_parent, fragment, outer_tls +): + if shutil.which("go") is None: + pytest.skip("go is not installed") + root = tmp_path / "server" / "authority" + initialize_authority(root, recovery_file=tmp_path / "operator.json") + material = load_pairing_identity(root) + assert material is not None + leaf, authority = material + app = cast(Starlette, _client(tmp_path)._app) + completed = [] + server_batches = [] + + async def handler(carrier): + # Only this test factory replaces the fixed-control consumer. The + # production route and credential policy remain unchanged. + connection = await carrier.responder._run( + carrier.lease, lambda: control_tls._connection(root) + ) + stream = TLSStream(carrier, connection) + await stream._call(connection.do_handshake) + carrier.begin_stream() + original_send = carrier.send + pending = [] + + async def send(data): + if not fragment: + # A bounded test relay holds multiple application records, including + # every separately drained BIO chunk, for one outbound message. + assert sum(map(len, pending)) + len(data) <= carrier.frame_limit + pending.append(data) + return + for offset in range(0, len(data), 1375): + await original_send(data[offset : offset + 1375]) + + carrier.send = send + total = 0 + batch_writes = 0 + # Exceed the fixed control's lifetime budget in both directions. + while total < CHUNK * 32: + data = await stream.read() + assert data + total += len(data) + await stream.write(data) + batch_writes += 1 + if not fragment and total % (CHUNK * 2) == 0: + assert batch_writes >= 2 + assert len(pending) >= 2 + await original_send(b"".join(pending)) + pending.clear() + server_batches.append(batch_writes) + batch_writes = 0 + carrier.credits.check() + await asyncio.sleep(0) # Allow queued callbacks/backpressure. + assert not pending + completed.append(total) + await stream.close() + + app.state.native_websocket_protocol = protocol_factory( + app.state.pairing_responder, + app.state.pairing_responder._limits, + continue_control=handler, + ) + kwargs = {} + if outer_tls: + cert, key = tmp_path / "outer.pem", tmp_path / "outer.key" + _generate_self_signed_tls_certificate( + cert_path=cert, key_path=key, names=("127.0.0.1",) + ) + kwargs = {"ssl_certfile": str(cert), "ssl_keyfile": str(key)} + with live_server( + app, + ws=app.state.native_websocket_protocol, + http=app.state.native_http_protocol, + **kwargs, + ) as port: + fixture = tmp_path / "public-stream-peer.json" + fixture.write_text( + json.dumps( + { + "URL": f"{'wss' if outer_tls else 'ws'}://127.0.0.1:{port}/api/v1/native/control", + "Fragment": fragment, + "ServerID": leaf.identity.subject_id, + "Authority": base64.b64encode(authority).decode("ascii"), + } + ) + ) + result = subprocess.run( + [ + "go", + "test", + "-count=1", + "-v", + "-run", + "^TestPythonAuthenticationStreamPeer$", + "./internal/identity", + ], + cwd=Path(__file__).parents[3] / "client/go-cli", + env={**os.environ, "ARBITER_TEST_AUTH_STREAM": str(fixture)}, + capture_output=True, + text=True, + timeout=60, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert result.stdout.splitlines().count("ARBITER_STREAM_PEER=pass") == 1 + assert completed == [CHUNK * 32] + if not fragment: + assert len(server_batches) == 16 and min(server_batches) >= 2 + assert not app.state.pairing_responder._leases diff --git a/server/tests/unit/test_control_tls.py b/server/tests/unit/test_control_tls.py index 7107edb..aae355a 100644 --- a/server/tests/unit/test_control_tls.py +++ b/server/tests/unit/test_control_tls.py @@ -135,6 +135,72 @@ async def settle(): asyncio.run(run()) +def test_cancelled_control_payload_retains_export_until_worker_exits(): + from cffi import FFI + + entered, release = threading.Event(), threading.Event() + ffi = FFI() + + class Connection: + calls = 0 + + def send(self, data): + self.calls += 1 + if self.calls == 1: + return len(data) + exported = ffi.from_buffer(data) + entered.set() + try: + assert release.wait(5) + return len(data) + finally: + del exported + + def bio_read(self, size): + raise control_tls.SSL.WantReadError() + + def pending(self): + return 0 + + async def run(): + async def waiting(_carrier): + await asyncio.Event().wait() + + async def wait_for_entered(): + for _ in range(2000): + if entered.is_set(): + return + await asyncio.sleep(0.001) + raise AssertionError("worker did not enter") + + owner, carrier, _, _ = setup(waiting, max_connections=1) + stream = control_tls.ControlTLS( + carrier, cast(control_tls.SSL.Connection, Connection()) + ) + task = asyncio.create_task(stream.send(b"payload")) + try: + await wait_for_entered() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + carrier.connection_lost(None) + assert len(owner._leases) == 1 + release.set() + + async def settled(): + while owner._leases: + await asyncio.sleep(0.001) + + await asyncio.wait_for(settled(), 2) + finally: + release.set() + carrier.connection_lost(None) + owner.close() + await asyncio.sleep(0) + + asyncio.run(run()) + + def test_real_tls_disconnect_retains_capacity_and_health( tmp_path, private_windows_test_parent, monkeypatch ): @@ -322,7 +388,7 @@ async def run(): with pytest.raises(ValueError, match="Control TLS input|Pairing carrier"): await stream.receive() assert carrier.closed - assert carrier._control_tls_fed <= carrier.limits.max_buffered_bytes // 2 + assert carrier.credits.fed <= carrier.limits.max_buffered_bytes // 2 finally: carrier.connection_lost(None) owner.close() @@ -365,10 +431,24 @@ async def run(): root, 262144, input_chunk_limit=16406, coalesce_input=coalesce_input ) try: + snapshots = [] + original_feed = stream._feed + + async def observed_feed(): + await original_feed() + snapshots.append( + (len(stream._input), stream._offset, carrier.credits.adapter) + ) + + stream._feed = observed_feed payload = b"r" * 65536 _client_send(client, len(payload).to_bytes(4, "big") + payload) assert await asyncio.wait_for(stream.receive(), 3) == payload - assert carrier._control_tls_fed <= carrier.limits.max_buffered_bytes // 2 + assert any( + remaining > offset > 0 and adapter == remaining + for remaining, offset, adapter in snapshots + ) + assert carrier.credits.fed <= carrier.limits.max_buffered_bytes // 2 finally: carrier.connection_lost(None) owner.close() @@ -385,8 +465,13 @@ def _live_tls_bytes() -> int: while frame is not None: name = frame.f_code.co_name filename = frame.f_code.co_filename.replace("\\", "/") - if (filename.endswith("/control_tls.py") and name in {"send", "_call"}) or ( - filename.endswith("/pairing_carrier.py") and name in {"send", "_write"} + if ( + (filename.endswith("/control_tls.py") and name == "send") + or ( + filename.endswith("/native_stream.py") + and name in {"_call", "_pump", "_feed"} + ) + or (filename.endswith("/pairing_carrier.py") and name in {"send", "_write"}) ): for value in frame.f_locals.values(): if isinstance(value, (bytes, bytearray)) and value: @@ -414,6 +499,8 @@ def _compiled_live_tls_bytes(filename: str, function_name: str) -> int: [ (str(Path("arbiter_server") / "control_tls.py"), "send"), (str(PureWindowsPath("arbiter_server") / "control_tls.py"), "send"), + (str(Path("arbiter_server") / "native_stream.py"), "_pump"), + (str(PureWindowsPath("arbiter_server") / "native_stream.py"), "_pump"), (str(Path("arbiter_server") / "pairing_carrier.py"), "_write"), (str(PureWindowsPath("arbiter_server") / "pairing_carrier.py"), "_write"), ], @@ -498,7 +585,7 @@ async def receive(): assert carrier.closed assert not carrier.control_started_once assert offset <= carrier.limits.max_buffered_bytes - assert carrier._control_tls_fed <= carrier.limits.max_buffered_bytes // 2 + assert carrier.credits.fed <= carrier.limits.max_buffered_bytes // 2 finally: carrier.connection_lost(None) owner.close() @@ -514,8 +601,8 @@ async def waiting(_carrier): owner, carrier, transport, ws = setup(waiting, max_buffered_bytes=8192) try: - carrier._control_tls_fed = 2048 - carrier._control_tls_pending = 1024 + carrier.credits.fed = 2048 + carrier.credits.pending = 1024 carrier.data_received(ws.send(events.BytesMessage(data=b"x" * 1710))) assert carrier.closed assert transport.closed diff --git a/server/tests/unit/test_native_stream.py b/server/tests/unit/test_native_stream.py new file mode 100644 index 0000000..602da84 --- /dev/null +++ b/server/tests/unit/test_native_stream.py @@ -0,0 +1,293 @@ +"""Shared pump progress and credits, including actual cancelled worker owners.""" + +import asyncio +import threading +import time +from typing import cast + +from OpenSSL import SSL +import pytest + +from arbiter_server.native_stream import CHUNK, RECORD_RESERVE, TLSStream +from server.tests.unit.test_control_carrier import setup + + +class Connection: + def __init__(self): + self.output = [] + self.received = bytearray() + self.attempts = 0 + + def bio_read(self, size): + if self.output: + return self.output.pop(0) + raise SSL.WantReadError() + + def pending(self): + return 0 + + def bio_write(self, data): + self.received.extend(data) + return len(data) + + def send(self, data): + return len(data) + + def recv(self, size): + raise SSL.WantReadError() + + def shutdown(self): + return None + + +def test_want_write_drains_before_retry_and_read_feeds_bounded_chunks(): + async def run(): + async def waiting(carrier): + await asyncio.Event().wait() + + owner, carrier, _, _ = setup(waiting) + connection = Connection() + stream = TLSStream(carrier, cast(SSL.Connection, connection)) + carrier.begin_stream() + sends = [] + + async def capture(wire): + sends.append(wire) + + carrier.send = capture + + async def receive(): + return b"x" * (CHUNK + 1) + + carrier.receive = receive + + def operation(): + connection.attempts += 1 + if connection.attempts == 1: + connection.output.append(b"drain first") + raise SSL.WantWriteError() + assert sends == [b"drain first"] + if len(connection.received) < CHUNK + 1: + raise SSL.WantReadError() + return b"complete" + + try: + assert await stream._call(operation) == b"complete" + assert connection.received == b"x" * (CHUNK + 1) + assert carrier.credits.fed >= RECORD_RESERVE + assert carrier.credits.adapter == 0 + assert connection.attempts == 4 + finally: + carrier.connection_lost(None) + owner.close() + await asyncio.sleep(0) + + asyncio.run(run()) + + +@pytest.mark.parametrize("failure", ["write-stall", "zero-write", "overrun"]) +def test_no_progress_and_pending_overrun_are_finite(failure): + async def run(): + async def waiting(carrier): + await asyncio.Event().wait() + + owner, carrier, _, _ = setup(waiting) + connection = Connection() + stream = TLSStream(carrier, cast(SSL.Connection, connection)) + carrier.begin_stream() + + def write(data): + if failure == "write-stall": + raise SSL.WantWriteError() + return 0 + + connection.send = write + try: + if failure == "overrun": + carrier.credits.pending = carrier.limits.max_buffered_bytes + with pytest.raises(ValueError, match="progress|limit"): + await asyncio.wait_for(stream.write(b"payload"), 1) + finally: + carrier.connection_lost(None) + owner.close() + await asyncio.sleep(0) + + asyncio.run(run()) + + +def test_cancelled_worker_retains_credits_and_lease_until_actual_exit(): + entered, release = threading.Event(), threading.Event() + from cffi import FFI + + ffi = FFI() + + async def run(): + async def waiting(carrier): + await asyncio.Event().wait() + + async def wait_for_entered(): + for _ in range(2000): + if entered.is_set(): + return + await asyncio.sleep(0.001) + raise AssertionError("worker did not enter") + + owner, carrier, _, _ = setup(waiting, max_connections=1) + connection = Connection() + stream = TLSStream(carrier, cast(SSL.Connection, connection)) + carrier.begin_stream() + + def write(data): + exported = ffi.from_buffer(data) + entered.set() + try: + assert release.wait(5) + return min(3, len(data)) + finally: + del exported + + connection.send = write + task = asyncio.create_task(stream.write(b"payload")) + try: + await wait_for_entered() + held = carrier.credits.application + with pytest.raises(ValueError, match="Concurrent"): + await stream.read(1) + assert carrier.credits.application == held == 14 + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + with pytest.raises(ValueError, match="Concurrent"): + await stream.read(1) + carrier.connection_lost(None) + assert len(owner._leases) == 1 + assert carrier.credits.application == held + release.set() + + async def settled(): + while owner._leases: + await asyncio.sleep(0.001) + + await asyncio.wait_for(settled(), 2) + lease = owner.reserve() + lease.close() + finally: + release.set() + carrier.connection_lost(None) + owner.close() + await asyncio.sleep(0) + + asyncio.run(run()) + + +@pytest.mark.parametrize("blocked", [False, True]) +def test_partial_writes_and_transport_credit_preflight(blocked): + async def run(): + async def waiting(carrier): + await asyncio.Event().wait() + + owner, carrier, transport, _ = setup(waiting) + connection = Connection() + stream = TLSStream(carrier, cast(SSL.Connection, connection)) + carrier.begin_stream() + written = bytearray() + threads = [] + + def send(data): + threads.append(threading.get_ident()) + written.extend(data[:3]) + return min(3, len(data)) + + connection.send = send + connection.shutdown = lambda: threads.append(threading.get_ident()) + try: + if blocked: + transport.buffered = carrier.limits.max_buffered_bytes - 1 + with pytest.raises(ValueError, match="limit"): + await stream.write(b"payload") + assert not written # Reject before another crypto allocation. + else: + await stream.write(b"payload") + assert written == b"payload" + await stream.close() + assert carrier.closed + assert threads and threading.get_ident() not in threads + finally: + carrier.connection_lost(None) + owner.close() + await asyncio.sleep(0) + + asyncio.run(run()) + + +@pytest.mark.parametrize("direction", ["read", "write"]) +def test_abort_wakes_a_stream_blocked_on_either_direction(direction): + async def run(): + async def waiting(carrier): + await asyncio.Event().wait() + + owner, carrier, _, _ = setup(waiting) + connection = Connection() + stream = TLSStream(carrier, cast(SSL.Connection, connection)) + carrier.begin_stream() + + connection.output = [b"response"] if direction == "write" else [] + if direction == "write": + carrier.pause_writing() + + async def transfer(): + if direction == "write": + await stream.write(b"request") + else: + await stream.read() + + task = asyncio.create_task(transfer()) + try: + await asyncio.sleep(0.02) + assert not task.done() + carrier.abort() + with pytest.raises(ValueError, match="closed"): + await asyncio.wait_for(task, 1) + finally: + carrier.connection_lost(None) + owner.close() + await asyncio.sleep(0) + + asyncio.run(run()) + + +def test_competing_input_owners_cannot_overbook_credits(): + async def run(): + async def waiting(carrier): + await asyncio.Event().wait() + + owner, carrier, _, _ = setup(waiting, max_buffered_bytes=65536) + carrier.begin_stream() + original = carrier.credits.check + barrier = threading.Barrier(2) + results = [] + + def slow_check(extra=0): + original(extra) + time.sleep(0.02) # Force overlap between preflight and reservation. + + carrier.credits.check = slow_check + + def reserve(): + barrier.wait(timeout=1) + try: + carrier.credits.reserve_input(10000) + results.append("reserved") + except ValueError: + results.append("full") + + try: + await asyncio.gather(asyncio.to_thread(reserve), asyncio.to_thread(reserve)) + assert sorted(results) == ["full", "reserved"] + original() + finally: + carrier.connection_lost(None) + owner.close() + await asyncio.sleep(0) + + asyncio.run(run())