From 77ee831e7a3729f5506be40cfcf12d98e0ea7e5a Mon Sep 17 00:00:00 2001 From: Rongsong Shen Date: Wed, 22 Jul 2026 07:47:51 +0000 Subject: [PATCH] feat(admin-config-tool): task 2.4 - xcecommon.wire TLS connection + 8-byte native-order length-prefixed framing to the controller, verified byte-for-byte against request.c's own framing and myapp.c's client TLS posture (SSL_VERIFY_NONE, no CN check). Opcode-agnostic: moves raw bytes only, so the same send_request() serves both OP_upgradeConfig and ProbeConfigTag's direct-struct variant. Connection always closes on error, including handshake timeout. --- .kiro/specs/admin-config-tool/tasks.md | 3 +- src/xceadmin/src/xcecommon/wire.py | 158 +++++++++++++ src/xceadmin/tests/test_wire.py | 307 +++++++++++++++++++++++++ 3 files changed, 467 insertions(+), 1 deletion(-) create mode 100644 src/xceadmin/src/xcecommon/wire.py create mode 100644 src/xceadmin/tests/test_wire.py diff --git a/.kiro/specs/admin-config-tool/tasks.md b/.kiro/specs/admin-config-tool/tasks.md index 5dd0127b..7789769f 100644 --- a/.kiro/specs/admin-config-tool/tasks.md +++ b/.kiro/specs/admin-config-tool/tasks.md @@ -48,7 +48,7 @@ - _Boundary: xcecommon.mapping_ - _Depends: 2.2_ -- [ ] 2.4 (P) Establish the TLS connection and wire framing to the controller +- [x] 2.4 (P) Establish the TLS connection and wire framing to the controller - A single request can be sent and a single reply received over a TLS connection using the cluster's existing certificate trust model - Message framing matches the controller's own byte-for-byte, verified against known-good byte sequences - The connection always closes cleanly, including on error or interruption @@ -149,3 +149,4 @@ - **2.1**: `check_undefined_refs` deliberately does not parse resource names out of `res_req`/`gpu_req` free-form LSF-style expression strings (e.g. `"mem>1000 rusage[mem=1000]"`) — design.md's `xcecommon.validation` component scopes this check to queue→host-group and limit→group/`resname` only. `limits.yaml`'s `projects` field and `scaler.yaml` name collisions are likewise out of scope — no `projects.yaml` catalog exists among the 7 config types, and Requirement 4.3 only enumerates queue/host/resource/group/limit names. `src/configs/hosts.yaml` has 2 pre-existing undefined refs (resource `mg`, host `h3`) that now surface as warnings — expected, not a regression; fixing the example fixture is outside this task's boundary. - **2.2**: The real `controller.capnp` names its reply struct `RequestReply`, not `Reply` as design.md's postcondition states — an inexact design pseudo-type, corrected empirically (verified directly against `src/schema/controller.capnp`, not assumed). pycapnp 2.0.0/capnproto 1.2.0 abort the whole process (uncaught C++ exception, not a Python exception) on a missing cross-file import or malformed `.capnp` content passed to `capnp.load()` — a bare `try/except` cannot catch this. `xcecommon.schema.load_schema()` works around it by pre-validating in an isolated subprocess before ever calling `capnp.load()` in the caller's own process, converting any abort into a clean `SchemaNotFoundError`. Any future code that calls `capnp.load()` directly (bypassing this loader) reintroduces the crash risk — always go through `xcecommon.schema.load_schema()`. - **2.3**: `map_all_sections()` (a `SECTION_MAP`-driven orchestrator, not explicitly named in design.md's Service Interface but consistent with its Contracts) is the single call `xceadmin.build` (task 3.1) needs — it treats a missing YAML file identically to a present-but-empty section (byte-for-byte), satisfying Requirement 1.5 structurally. `target.init(field, 0)` and never touching an empty/absent list field produce *different* serialized bytes for the same logical "nothing here" value (verified empirically) — the mapper always picks "never touch" for byte-determinism (9.1). Calling `.to_bytes()` twice on the same capnp builder triggers a spurious "message already written" `UserWarning` even when no field was mutated between calls — capture the bytes once and reuse them rather than calling `.to_bytes()` again. +- **2.4**: `errors.py`'s exception subclasses (task 1.2) all take `message` as an explicit first constructor argument in addition to their structured fields (e.g. `ConnectionTimeoutError(message, host, port)`, `ReplyDecodeError(message, reason)`) — design.md's abbreviated `Service Interface` type stubs (`class ConnectionTimeoutError(XceAdminError): host: str; port: int`) read as if only the structured fields were needed; always check the real constructor in `errors.py`, not the design doc's shorthand. `xcecommon.wire` is deliberately opcode-agnostic — `send_request`/`send_frame`/`recv_frame` move raw `bytes` only and never import a capnp type; callers own `.to_bytes()`/`.from_bytes()`. Reused `ha/data/xce.crt`/`xce.key` (the repo's existing dev cert pair, same one `deploy-ha.sh` already copies for controller TLS) as the test fixture for both the TLS test server and the client `cert_dir` — no synthetic cert generation needed. diff --git a/src/xceadmin/src/xcecommon/wire.py b/src/xceadmin/src/xcecommon/wire.py new file mode 100644 index 00000000..3480f630 --- /dev/null +++ b/src/xceadmin/src/xcecommon/wire.py @@ -0,0 +1,158 @@ +"""TLS connection and wire framing for xcecommon/xceadmin (design.md, +``xcecommon.wire``; Requirements 2.1, 2.3, 2.4, 5.2, 6.1, 7.1, 7.3, 7.6). + +This module is opcode-agnostic: it moves raw ``bytes`` across a TLS socket +and never decodes capnp content itself. That is deliberate -- the same +:func:`send_request` primitive serves both the ``Request``/``Reply`` +union (``OP_upgradeConfig``) and ``ProbeConfigTag``'s direct-struct +variant (``config push --dry-run``), which have unrelated capnp shapes. +Callers own ``.to_bytes()``/``.from_bytes()``. + +Wire compatibility (verified against the real C implementation, not +assumed): + +- TLS posture mirrors the controller's own outbound client context exactly + (``myapp.c:103-109``): ``SSL_CTX_new(TLS_client_method())`` + + ``SSL_CTX_set_verify(..., SSL_VERIFY_NONE, NULL)`` -- no CN/peer-cert + validation anywhere in this trust model to diverge from. +- Framing matches ``request.c``'s ``xce_replconfig_req_ssl`` + (``src/lib/request.c:925-934``): an 8-byte ``uint64_t`` length header + written via a raw ``SSL_write(ssl, &hdr, sizeof(hdr))`` -- host/native + byte order, no ``htole64``/``htobe64`` anywhere in the C framing path. + ``struct.Struct("=Q")`` reproduces this exactly. + +Never pass an ``SSLSocket`` (or its ``fileno()``) into a pycapnp +read/write call -- pycapnp's ``.write()``/``.read()`` operate on the raw +fd directly, which would bypass TLS encryption entirely. Encode/decode +happens in-memory (``to_bytes()``/``from_bytes()``) against the plaintext +buffers this module's :func:`send_frame`/:func:`recv_frame` assemble. +""" + +from __future__ import annotations + +import socket +import ssl +import struct + +from xcecommon.errors import ConnectionTimeoutError, ReplyDecodeError + +#: Native-order uint64_t length prefix -- matches request.c's +#: ``uint64_t hdr = (uint64_t)blen;`` raw-memory write exactly (no network +#: byte order conversion exists anywhere in the C framing path). +_HEADER = struct.Struct("=Q") + +#: Mirrors request.c's own sanity bound on an incoming frame length +#: (``xce_replconfig_reply_ssl``, ``src/lib/request.c:1010``) -- a frame +#: claiming to be larger than this is corrupt/hostile, not merely large. +_MAX_FRAME_BYTES = 64 * 1024 * 1024 + +_CERT_FILENAME = "xce.crt" +_KEY_FILENAME = "xce.key" + +#: Timeout values below this are almost certainly a test deliberately +#: forcing a fast timeout, not a production caller -- kept as a constant +#: only so ``connect``'s default is named, not to gate behavior. +_DEFAULT_TIMEOUT_S = 30.0 + + +def _peername_or(sock: socket.socket, default: tuple[str, int]) -> tuple[str, int]: + try: + host, port = sock.getpeername()[:2] + return host, port + except OSError: + return default + + +def connect( + host: str, port: int, cert_dir: str, timeout_s: float = _DEFAULT_TIMEOUT_S +) -> ssl.SSLSocket: + """Open a TLS connection to a controller node. + + Loads ``xce.crt``/``xce.key`` from ``cert_dir`` (2.3, 2.4, 7.1, 7.6) and + establishes trust exactly as the controller's own outbound TLS does -- + ``check_hostname=False``, ``verify_mode=ssl.CERT_NONE`` (no separate + client CA flag or cert-issuance step, matching ``xce_test`` and + ``deploy-ha.sh``'s existing trust model). + + Raises :class:`ConnectionTimeoutError` if the TCP connect or the TLS + handshake does not complete within ``timeout_s``. A refused connection + (``ConnectionRefusedError``) or other OS-level failure propagates + unwrapped -- the CLI boundary (not this module) is responsible for + turning that into an operator-facing message (Requirement 7.1). + """ + context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + context.check_hostname = False + context.verify_mode = ssl.CERT_NONE + context.load_cert_chain( + certfile=f"{cert_dir}/{_CERT_FILENAME}", keyfile=f"{cert_dir}/{_KEY_FILENAME}" + ) + + try: + raw_sock = socket.create_connection((host, port), timeout=timeout_s) + except TimeoutError: + raise ConnectionTimeoutError(f"timeout connecting to {host}:{port}", host, port) from None + + try: + return context.wrap_socket(raw_sock, server_hostname=None) + except TimeoutError: + raw_sock.close() + raise ConnectionTimeoutError(f"timeout connecting to {host}:{port}", host, port) from None + except Exception: + raw_sock.close() + raise + + +def send_frame(sock: ssl.SSLSocket, payload: bytes) -> None: + """Send one length-prefixed frame: an 8-byte native-order ``uint64_t`` + size header followed by ``payload``, byte-for-byte matching the + controller's own C framing (module docstring). + """ + try: + sock.sendall(_HEADER.pack(len(payload))) + if payload: + sock.sendall(payload) + except TimeoutError: + host, port = _peername_or(sock, (str(sock), 0)) + raise ConnectionTimeoutError(f"timeout connecting to {host}:{port}", host, port) from None + + +def _recv_exact(sock: ssl.SSLSocket, nbytes: int) -> bytes: + chunks: list[bytes] = [] + remaining = nbytes + while remaining > 0: + try: + chunk = sock.recv(remaining) + except TimeoutError: + host, port = _peername_or(sock, (str(sock), 0)) + raise ConnectionTimeoutError(f"timeout connecting to {host}:{port}", host, port) from None + if not chunk: + received = nbytes - remaining + reason = f"connection closed after {received} of {nbytes} expected bytes" + raise ReplyDecodeError(f"invalid reply: {reason}", reason) + chunks.append(chunk) + remaining -= len(chunk) + return b"".join(chunks) + + +def recv_frame(sock: ssl.SSLSocket) -> bytes: + """Receive one length-prefixed frame written by :func:`send_frame`. + + Raises :class:`ReplyDecodeError` if the connection closes before a full + header or a full payload arrives (truncated), or if the declared + length is implausibly large (corrupt) -- Requirement 7.3. + """ + header = _recv_exact(sock, _HEADER.size) + (length,) = _HEADER.unpack(header) + if length > _MAX_FRAME_BYTES: + reason = f"frame length {length} exceeds max {_MAX_FRAME_BYTES}" + raise ReplyDecodeError(f"invalid reply: {reason}", reason) + return _recv_exact(sock, length) + + +def send_request(sock: ssl.SSLSocket, request_bytes: bytes) -> bytes: + """Frame and send ``request_bytes``, then receive and unframe exactly + one reply -- one-shot, no retry, no connection pooling (Invariant: a + single CLI command, not a daemon). + """ + send_frame(sock, request_bytes) + return recv_frame(sock) diff --git a/src/xceadmin/tests/test_wire.py b/src/xceadmin/tests/test_wire.py new file mode 100644 index 00000000..6b6644c5 --- /dev/null +++ b/src/xceadmin/tests/test_wire.py @@ -0,0 +1,307 @@ +"""Unit tests for ``xcecommon.wire`` (design.md, ``xcecommon.wire``; +Requirements 2.1, 2.3, 2.4, 5.2, 6.1, 7.1, 7.3, 7.6). + +Uses a real local TLS loopback server (threading) wrapped with this +repo's existing ``ha/data/xce.crt``/``xce.key`` -- the same cert pair +``deploy-ha.sh`` already copies for controller-to-controller TLS -- so +these tests exercise the real trust model (no CA, ``SSL_VERIFY_NONE``), +not a synthetic substitute. +""" + +from __future__ import annotations + +import socket +import ssl +import struct +import threading +from pathlib import Path + +import pytest + +from xcecommon import wire +from xcecommon.errors import ConnectionTimeoutError, ReplyDecodeError + +REPO_HA_DATA_DIR = Path(__file__).resolve().parents[3] / "ha" / "data" +CERT_FILE = REPO_HA_DATA_DIR / "xce.crt" +KEY_FILE = REPO_HA_DATA_DIR / "xce.key" + +pytestmark = pytest.mark.skipif( + not (CERT_FILE.is_file() and KEY_FILE.is_file()), + reason="ha/data/xce.crt or ha/data/xce.key not present", +) + + +def _server_context() -> ssl.SSLContext: + context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.load_cert_chain(certfile=str(CERT_FILE), keyfile=str(KEY_FILE)) + return context + + +class _TlsEchoServer: + """Accepts exactly one TLS connection, hands the raw accepted socket + to ``handler`` on a background thread, then closes it. ``handler`` + receives the wrapped :class:`ssl.SSLSocket`. + """ + + def __init__(self, handler): + self._handler = handler + self._listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self._listener.bind(("127.0.0.1", 0)) + self._listener.listen(1) + self.port = self._listener.getsockname()[1] + self._thread = threading.Thread(target=self._run, daemon=True) + self._thread.start() + + def _run(self) -> None: + try: + raw_conn, _ = self._listener.accept() + except OSError: + return + try: + tls_conn = _server_context().wrap_socket(raw_conn, server_side=True) + except (ssl.SSLError, OSError): + raw_conn.close() + return + try: + self._handler(tls_conn) + except OSError: + pass + finally: + tls_conn.close() + + def close(self) -> None: + self._listener.close() + self._thread.join(timeout=5) + + +@pytest.fixture +def echo_server(): + """A server that reads exactly one frame and writes it straight back.""" + servers: list[_TlsEchoServer] = [] + + def _handler(conn: ssl.SSLSocket) -> None: + # Loops so a single connection can serve more than one + # send_request() call -- send_request itself is one-shot per call, + # but nothing in this module prevents a caller reusing the socket. + while True: + try: + payload = wire.recv_frame(conn) + except ReplyDecodeError: + return # client closed the connection; done + wire.send_frame(conn, payload) + + def _make() -> _TlsEchoServer: + srv = _TlsEchoServer(_handler) + servers.append(srv) + return srv + + yield _make + for srv in servers: + srv.close() + + +CERT_DIR = str(REPO_HA_DATA_DIR) + + +# --------------------------------------------------------------------------- +# connect() +# --------------------------------------------------------------------------- + + +def test_connect_establishes_real_tls_session(echo_server): + srv = echo_server() + sock = wire.connect("127.0.0.1", srv.port, CERT_DIR, timeout_s=5.0) + try: + assert isinstance(sock, ssl.SSLSocket) + assert sock.getpeername()[1] == srv.port + finally: + sock.close() + + +def test_connect_refused_propagates_unwrapped(): + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + listener.bind(("127.0.0.1", 0)) + port = listener.getsockname()[1] + listener.close() # nothing listens on this port now + + with pytest.raises(ConnectionRefusedError): + wire.connect("127.0.0.1", port, CERT_DIR, timeout_s=5.0) + + +def test_connect_handshake_timeout_raises_connection_timeout_error_and_closes_cleanly(): + """The server accepts the TCP connection but never completes a TLS + handshake -- the client's wrap_socket() must time out, raise + ConnectionTimeoutError, and leave no dangling socket (2.4's "connection + always closes cleanly, including on error"). + """ + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + port = listener.getsockname()[1] + accepted: list[socket.socket] = [] + + def _accept_and_stall() -> None: + try: + conn, _ = listener.accept() + accepted.append(conn) + except OSError: + pass + + thread = threading.Thread(target=_accept_and_stall, daemon=True) + thread.start() + try: + with pytest.raises(ConnectionTimeoutError) as exc_info: + wire.connect("127.0.0.1", port, CERT_DIR, timeout_s=0.3) + assert exc_info.value.host == "127.0.0.1" + assert exc_info.value.port == port + + thread.join(timeout=5) + assert accepted, "server never observed the TCP connection" + # The client sent a real ClientHello before its read timed out, so + # the server side has bytes already buffered -- drain them and + # confirm the stream then hits EOF (the client closed its raw + # socket on timeout), rather than hanging or erroring oddly. + accepted[0].settimeout(5.0) + while accepted[0].recv(4096): + pass + finally: + listener.close() + for conn in accepted: + conn.close() + + +# --------------------------------------------------------------------------- +# send_frame() / recv_frame() -- byte-for-byte framing +# --------------------------------------------------------------------------- + + +def test_send_frame_writes_exact_known_good_byte_sequence(echo_server): + """Requirement: framing matches the controller's own byte-for-byte, + verified against a known-good byte sequence (not just a round-trip). + """ + captured = bytearray() + + def _capture_handler(conn: ssl.SSLSocket) -> None: + # Read exactly what send_frame wrote, raw, without going through + # recv_frame -- this is what proves the header shape independently + # of this module's own decoder. + while len(captured) < 8 + 4: + chunk = conn.recv(4096) + if not chunk: + break + captured.extend(chunk) + + srv = _TlsEchoServer(_capture_handler) + try: + sock = wire.connect("127.0.0.1", srv.port, CERT_DIR, timeout_s=5.0) + try: + wire.send_frame(sock, b"xce!") + finally: + sock.close() + finally: + srv.close() + + # 8-byte native-order uint64_t length header (4), then the 4 payload + # bytes -- byte-for-byte, matching struct.pack("=Q", 4) + b"xce!". + expected_header = struct.pack("=Q", 4) + assert bytes(captured) == expected_header + b"xce!" + + +def test_recv_frame_round_trips_arbitrary_payload(echo_server): + srv = echo_server() + sock = wire.connect("127.0.0.1", srv.port, CERT_DIR, timeout_s=5.0) + try: + payload = b"\x00\x01\xffcapnp-shaped-bytes" * 100 + wire.send_frame(sock, payload) + assert wire.recv_frame(sock) == payload + finally: + sock.close() + + +def test_recv_frame_round_trips_empty_payload(echo_server): + srv = echo_server() + sock = wire.connect("127.0.0.1", srv.port, CERT_DIR, timeout_s=5.0) + try: + wire.send_frame(sock, b"") + assert wire.recv_frame(sock) == b"" + finally: + sock.close() + + +def test_recv_frame_raises_reply_decode_error_on_truncated_header(): + def _close_after_partial_header(conn: ssl.SSLSocket) -> None: + conn.sendall(b"\x01\x02\x03") # 3 of 8 header bytes, then hang up + + srv = _TlsEchoServer(_close_after_partial_header) + try: + sock = wire.connect("127.0.0.1", srv.port, CERT_DIR, timeout_s=5.0) + try: + with pytest.raises(ReplyDecodeError): + wire.recv_frame(sock) + finally: + sock.close() + finally: + srv.close() + + +def test_recv_frame_raises_reply_decode_error_on_truncated_payload(): + def _promise_more_than_sent(conn: ssl.SSLSocket) -> None: + conn.sendall(struct.pack("=Q", 100)) + conn.sendall(b"only ten!") # 9 bytes, not 100 -- then hang up + + srv = _TlsEchoServer(_promise_more_than_sent) + try: + sock = wire.connect("127.0.0.1", srv.port, CERT_DIR, timeout_s=5.0) + try: + with pytest.raises(ReplyDecodeError): + wire.recv_frame(sock) + finally: + sock.close() + finally: + srv.close() + + +def test_recv_frame_raises_reply_decode_error_on_implausible_length(): + def _claim_huge_length(conn: ssl.SSLSocket) -> None: + conn.sendall(struct.pack("=Q", wire._MAX_FRAME_BYTES + 1)) + + srv = _TlsEchoServer(_claim_huge_length) + try: + sock = wire.connect("127.0.0.1", srv.port, CERT_DIR, timeout_s=5.0) + try: + with pytest.raises(ReplyDecodeError): + wire.recv_frame(sock) + finally: + sock.close() + finally: + srv.close() + + +# --------------------------------------------------------------------------- +# send_request() +# --------------------------------------------------------------------------- + + +def test_send_request_returns_exact_reply_bytes(echo_server): + srv = echo_server() + sock = wire.connect("127.0.0.1", srv.port, CERT_DIR, timeout_s=5.0) + try: + reply = wire.send_request(sock, b"pretend-capnp-request-bytes") + assert reply == b"pretend-capnp-request-bytes" + finally: + sock.close() + + +def test_send_request_is_opcode_agnostic_works_for_any_byte_shape(echo_server): + """send_request never inspects payload content -- it must serve both + the Request/Reply union and ProbeConfigTag's direct-struct variant + identically, since both are just bytes to this module. + """ + srv = echo_server() + sock = wire.connect("127.0.0.1", srv.port, CERT_DIR, timeout_s=5.0) + try: + for payload in (b"", b"\x00", b"short", b"x" * 10_000): + reply = wire.send_request(sock, payload) + assert reply == payload + finally: + sock.close()