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"