sdk/python: send the first frames with the agent upgrade request over plain HTTP

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())

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

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"
     )

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"