From 1dadd03d0e64e65be1b50212ca2f1aa385bdd722 Mon Sep 17 00:00:00 2001 From: CMGS Date: Mon, 5 Oct 2026 14:20:22 +0800 Subject: [PATCH] sdk/python: send the first frames with the agent upgrade request over plain HTTP A relay dial waited for the 101 before writing its first frame, one extra round trip on every dial: the proto probe, every call after the idle window, every call with keep-alive off. sandboxd replays bytes buffered behind the request, so over plain HTTP the request now leaves with the first frame and the 101 is read on the first recv, under the dial timeout or the caller's tighter socket timeout. HTTPS keeps the old order, since a TLS edge may be a proxy that drops early bytes. A connection that closes before the upgrade reply is now a ProtocolError, so a run deadline that cuts a pending upgrade surfaces as SandboxTimeout. The relay fakes set TCP_NODELAY as sandboxd does; without it Nagle held a reply behind the 101 and the fake's reset dropped it. --- sdk/python/cocoonsandbox/conn.py | 120 ++++++++++++++++++----------- sdk/python/tests/conftest.py | 1 + sdk/python/tests/test_hardening.py | 9 ++- sdk/python/tests/test_upgrade.py | 91 ++++++++++++++++++++++ 4 files changed, 175 insertions(+), 46 deletions(-) create mode 100644 sdk/python/tests/test_upgrade.py diff --git a/sdk/python/cocoonsandbox/conn.py b/sdk/python/cocoonsandbox/conn.py index 0dc4a264..da2514bb 100644 --- a/sdk/python/cocoonsandbox/conn.py +++ b/sdk/python/cocoonsandbox/conn.py @@ -9,7 +9,7 @@ import threading import time import urllib.parse -from collections.abc import Iterator +from collections.abc import Callable, Iterator from typing import Any, BinaryIO, Protocol, TypeVar from .errors import APIError, ProtocolError, SandboxError, SandboxTimeout, SilkdError @@ -33,15 +33,22 @@ def close(self) -> None: ... class Conn(_Closeable): - """A live frame stream to one sandbox's silkd, via the owner node.""" + """A frame stream to one sandbox's silkd; a plain-HTTP upgrade rides the first send, its 101 the first recv.""" - def __init__(self, sock: socket.socket, reader: BinaryIO) -> None: + def __init__( + self, sock: socket.socket, reader: BinaryIO, request: bytes = b"", upgrade_timeout: float | None = None + ) -> None: self._sock = sock self._reader = reader + self._request = request + self._upgrade_timeout = upgrade_timeout def send(self, op: str, **fields: object) -> None: + frame = encode_request(op, **fields) + if self._request: + frame, self._request = self._request + frame, b"" try: - self._sock.sendall(encode_request(op, **fields)) + self._sock.sendall(frame) except OSError as exc: raise ProtocolError(f"write failed: {exc}") from exc @@ -55,6 +62,8 @@ def abort(self) -> None: def recv(self) -> dict[str, Any]: """Reads one frame or raises a typed protocol or guest error.""" + if self._upgrade_timeout is not None: + self._finish_upgrade(self._upgrade_timeout) try: line = self._reader.readline(MAX_FRAME + 1) except (OSError, ValueError) as exc: @@ -94,6 +103,22 @@ def close(self) -> None: finally: self._sock.close() + def _finish_upgrade(self, timeout: float) -> None: + self._upgrade_timeout = None + caller = self._sock.gettimeout() + limit = timeout if caller is None else min(caller, timeout) + deadline = time.monotonic() + limit + try: + _read_upgrade( + self._reader, lambda: self._sock.settimeout(remaining_timeout(limit, deadline, "agent upgrade")) + ) + self._sock.settimeout(caller) + except Exception as exc: + self.close() + if isinstance(exc, OSError) and not isinstance(exc, SandboxError): + raise _handshake_error("agent upgrade", exc) from exc + raise + class ConnPool: """Parks a handle's idle relay connections between calls; one sweeper timer closes them as they expire.""" @@ -172,49 +197,31 @@ def dial_agent( except OSError as exc: raise _handshake_error(f"dial {addr}", exc) from exc sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + request = ( + f"GET /v1/sandboxes/{sandbox_id}/agent HTTP/1.1\r\n" + f"Host: {endpoint.netloc}\r\n" + "Connection: Upgrade\r\n" + "Upgrade: silkd\r\n" + f"Authorization: Bearer {token}\r\n" + "\r\n" + ).encode() reader = None try: - if endpoint.scheme == "https": - context = ssl_context or ssl.create_default_context() - sock = context.wrap_socket(sock, server_hostname=host, do_handshake_on_connect=False) - sock.settimeout(remaining_timeout(timeout, deadline)) - try: - sock.do_handshake() - except OSError as exc: - raise _handshake_error(f"tls handshake {addr}", exc) from exc - request = ( - f"GET /v1/sandboxes/{sandbox_id}/agent HTTP/1.1\r\n" - f"Host: {endpoint.netloc}\r\n" - "Connection: Upgrade\r\n" - "Upgrade: silkd\r\n" - f"Authorization: Bearer {token}\r\n" - "\r\n" - ) + if endpoint.scheme != "https": + sock.settimeout(None) + return Conn(sock, sock.makefile("rb"), request, timeout) + context = ssl_context or ssl.create_default_context() + sock = context.wrap_socket(sock, server_hostname=host, do_handshake_on_connect=False) sock.settimeout(remaining_timeout(timeout, deadline)) - sock.sendall(request.encode()) - reader = sock.makefile("rb") + try: + sock.do_handshake() + except OSError as exc: + raise _handshake_error(f"tls handshake {addr}", exc) from exc + # a TLS edge may be a proxy that drops bytes sent ahead of the 101 sock.settimeout(remaining_timeout(timeout, deadline)) - status = reader.readline(1024).decode(errors="replace") - parts = status.split(" ", 2) - code = int(parts[1]) if len(parts) > 1 and parts[1].isdigit() else 0 - body_len = 0 - while True: - sock.settimeout(remaining_timeout(timeout, deadline)) - header = reader.readline(4096) - if header in (b"\r\n", b"\n", b""): - break - name, _, value = header.decode(errors="replace").partition(":") - if name.strip().lower() == "content-length": - try: - body_len = int(value.strip()) - if body_len < 0: - raise ValueError("negative content-length") - except ValueError as exc: - raise ProtocolError("invalid content-length in upgrade reply") from exc - if code != 101: - sock.settimeout(remaining_timeout(timeout, deadline)) - body = reader.read(min(body_len, MAX_FRAME)).decode(errors="replace") if body_len else "" - raise APIError("agent upgrade", code, body.strip() or status.strip()) + sock.sendall(request) + reader = sock.makefile("rb") + _read_upgrade(reader, lambda: sock.settimeout(remaining_timeout(timeout, deadline))) remaining_timeout(timeout, deadline) sock.settimeout(None) return Conn(sock, reader) @@ -241,3 +248,30 @@ def _handshake_error(what: str, exc: OSError) -> SandboxError: if isinstance(exc, (TimeoutError, socket.timeout)): return SandboxTimeout(f"{what} timed out") return ProtocolError(f"{what}: {exc}") + + +def _read_upgrade(reader: BinaryIO, arm: Callable[[], None]) -> None: + arm() + status = reader.readline(1024).decode(errors="replace") + if not status: + raise ProtocolError("connection closed before the upgrade reply") + parts = status.split(" ", 2) + code = int(parts[1]) if len(parts) > 1 and parts[1].isdigit() else 0 + body_len = 0 + while True: + arm() + header = reader.readline(4096) + if header in (b"\r\n", b"\n", b""): + break + name, _, value = header.decode(errors="replace").partition(":") + if name.strip().lower() == "content-length": + try: + body_len = int(value.strip()) + if body_len < 0: + raise ValueError("negative content-length") + except ValueError as exc: + raise ProtocolError("invalid content-length in upgrade reply") from exc + if code != 101: + arm() + body = reader.read(min(body_len, MAX_FRAME)).decode(errors="replace") if body_len else "" + raise APIError("agent upgrade", code, body.strip() or status.strip()) diff --git a/sdk/python/tests/conftest.py b/sdk/python/tests/conftest.py index 3fcc0fd5..2023a274 100644 --- a/sdk/python/tests/conftest.py +++ b/sdk/python/tests/conftest.py @@ -54,6 +54,7 @@ def sandbox_at(addr: str, **client_kwargs) -> Sandbox: def accept_upgrade(conn: socket.socket) -> BinaryIO: + conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) reader = conn.makefile("rb") while reader.readline() not in (b"\r\n", b""): pass diff --git a/sdk/python/tests/test_hardening.py b/sdk/python/tests/test_hardening.py index c0275a35..1f4bbaa6 100644 --- a/sdk/python/tests/test_hardening.py +++ b/sdk/python/tests/test_hardening.py @@ -80,7 +80,9 @@ def test_raw_replies_surface_as_typed_errors(raw_reply, response, call, raises, addr = raw_reply(response) with pytest.raises(raises) as exc_info: if call == "dial": - dial_agent(addr, "sb_1", "tok", 0.5) + with dial_agent(addr, "sb_1", "tok", 0.5) as conn: + conn.send("info") + conn.recv() else: Client(addr).info() if status is not None: @@ -192,8 +194,9 @@ def slow_headers() -> None: threading.Thread(target=slow_headers, daemon=True).start() addr = f"127.0.0.1:{server.getsockname()[1]}" started = time.monotonic() - with pytest.raises(SandboxTimeout): - dial_agent(addr, "sb_x", "tok", timeout=0.2) + with pytest.raises(SandboxTimeout), dial_agent(addr, "sb_x", "tok", timeout=0.2) as conn: + conn.send("info") + conn.recv() assert time.monotonic() - started < 0.6, ( "each header arrived under the per-read timeout; only a whole-dial budget stops this" ) diff --git a/sdk/python/tests/test_upgrade.py b/sdk/python/tests/test_upgrade.py new file mode 100644 index 00000000..29980f6f --- /dev/null +++ b/sdk/python/tests/test_upgrade.py @@ -0,0 +1,91 @@ +"""The relay upgrade: over plain HTTP the request leaves with the first frame and +the 101 is read on the first recv; a refusal surfaces as an APIError.""" + +import contextlib +import json +import socket +import threading +import time + +import pytest +from conftest import sandbox_at + +from cocoonsandbox import APIError, SandboxTimeout +from cocoonsandbox.frames import KEEP_ALIVE_PROTO + +REPLIES = { + "info": {"type": "info", "version": "fake", "proto": KEEP_ALIVE_PROTO, "uptime_secs": 0, "procs": 0}, + "fs_mkdir": {"type": "done"}, +} +REJECT = b"HTTP/1.1 404 Not Found\r\nContent-Type: text/plain\r\nContent-Length: 15\r\n\r\nunknown sandbox" + + +class EarlyFrameAgent: + """Answers the 101 only after a frame has arrived behind the request, then serves frames until the peer leaves.""" + + def __init__(self, reply: bytes = b"HTTP/1.1 101 Switching Protocols\r\n\r\n") -> None: + self.reply = reply + self.early: list[str] = [] + self._server = socket.create_server(("127.0.0.1", 0)) + self.addr = f"127.0.0.1:{self._server.getsockname()[1]}" + threading.Thread(target=self._serve, daemon=True).start() + + def stop(self) -> None: + self._server.close() + + def _serve(self) -> None: + conn, _ = self._server.accept() + with conn, contextlib.suppress(OSError): + reader = conn.makefile("rb") + while reader.readline() not in (b"\r\n", b""): + pass + conn.settimeout(1) + op = json.loads(reader.readline())["op"] + self.early.append(op) + conn.settimeout(None) + conn.sendall(self.reply) + while op: + conn.sendall(json.dumps(REPLIES[op]).encode() + b"\n") + line = reader.readline() + op = json.loads(line)["op"] if line else "" + + +@pytest.mark.parametrize(("keep_alive", "early"), [(0, ["fs_mkdir"]), (30, ["info"])], ids=["no probe", "proto probe"]) +def test_first_frame_rides_ahead_of_the_101(keep_alive, early): + agent = EarlyFrameAgent() + sb = sandbox_at(agent.addr, timeout=5, keep_alive=keep_alive) + try: + sb.mkdir("/w") + finally: + sb._pool.drain() + agent.stop() + assert agent.early == early, agent.early + + +def test_rejected_upgrade_is_an_api_error(): + agent = EarlyFrameAgent(REJECT) + sb = sandbox_at(agent.addr, timeout=5, keep_alive=0) + try: + with pytest.raises(APIError) as refused: + sb.mkdir("/w") + finally: + agent.stop() + assert (refused.value.status, refused.value.message) == (404, "unknown sandbox"), refused.value + assert agent.early == ["fs_mkdir"], agent.early + + +def test_run_timeout_cuts_an_unanswered_upgrade(): + server = socket.create_server(("127.0.0.1", 0)) + addr = f"127.0.0.1:{server.getsockname()[1]}" + held: list[socket.socket] = [] + threading.Thread(target=lambda: held.append(server.accept()[0]), daemon=True).start() + sb = sandbox_at(addr, timeout=30, keep_alive=0) + started = time.monotonic() + try: + with pytest.raises(SandboxTimeout): + sb.run(["true"], timeout=0.2) + finally: + server.close() + for conn in held: + conn.close() + assert time.monotonic() - started < 2, "the run deadline must cut a hung upgrade under a long client timeout"