myxce/src/xceadmin/tests/test_wire.py

"""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 a
self-signed cert/key generated fresh per test session (see conftest.py's
``tls_certs`` fixture, same generation approach ``deploy-ha.sh`` uses) --
so these tests exercise the real trust model (no CA, ``SSL_VERIFY_NONE``,
no client certificate -- matching ``netio.c``'s client-side
``create_ssl_context(NULL, NULL)``), 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


def _server_context(cert_file: Path, key_file: Path) -> 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, cert_file: Path, key_file: Path):
        self._handler = handler
        self._cert_file = cert_file
        self._key_file = key_file
        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(self._cert_file, self._key_file).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(tls_certs):
    """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, tls_certs.cert_file, tls_certs.key_file)
        servers.append(srv)
        return srv

    yield _make
    for srv in servers:
        srv.close()


# ---------------------------------------------------------------------------
# connect()
# ---------------------------------------------------------------------------


def test_connect_establishes_real_tls_session(echo_server):
    srv = echo_server()
    sock = wire.connect("127.0.0.1", srv.port, 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, 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, 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(tls_certs):
    """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, tls_certs.cert_file, tls_certs.key_file)
    try:
        sock = wire.connect("127.0.0.1", srv.port, 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, 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, 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(tls_certs):
    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, tls_certs.cert_file, tls_certs.key_file)
    try:
        sock = wire.connect("127.0.0.1", srv.port, 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(tls_certs):
    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, tls_certs.cert_file, tls_certs.key_file)
    try:
        sock = wire.connect("127.0.0.1", srv.port, 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(tls_certs):
    def _claim_huge_length(conn: ssl.SSLSocket) -> None:
        conn.sendall(struct.pack("=Q", wire._MAX_FRAME_BYTES + 1))

    srv = _TlsEchoServer(_claim_huge_length, tls_certs.cert_file, tls_certs.key_file)
    try:
        sock = wire.connect("127.0.0.1", srv.port, 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, 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, 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()