path: root/debian/tests/python/ucspi_test/__init__.py

"""Run a couple of UCSPI client and server tests."""

from __future__ import annotations

import abc
import dataclasses
import pathlib
import shlex
import subprocess  # noqa: S404
import sys
import typing


if typing.TYPE_CHECKING:
    import socket
    from collections.abc import Callable
    from typing import Any, Final


VERSION: Final = "0.2.0"

MSG_RESP_HELLO: Final = "a01 hello\n"
MSG_RESP_BYE: Final = "a02 bye\n"


@dataclasses.dataclass(frozen=True)
class Config:
    """Runtime configuration for the UCSPI test runner."""

    bindir: pathlib.Path
    proto: str
    utf8_env: dict[str, str]


@dataclasses.dataclass
class RunnerError(Exception):
    """An error that occurred while preparing for or running the tests."""


@dataclasses.dataclass
class HandlerMismatchError(RunnerError):
    """The test framework attempted to add a different handler."""

    proto: str
    current: type[Runner]
    runner: type[Runner]

    def __str__(self) -> str:
        """Provide a human-readable error message."""
        return (
            f"Handler mismatch for the {self.proto!r} protocol: "
            "had {self.current!r}, now {self.runner!r}"
        )


@dataclasses.dataclass
class SocketAddressLengthError(RunnerError):
    """An address with an unexpected length was specified."""

    proto: str
    addr: Any

    def __str__(self) -> str:
        """Provide a human-readable error message."""
        return f"{self.proto}.get_connected_socket(): unexpected address length for {self.addr!r}"


class Runner(abc.ABC):
    """A helper class for running tests for a single UCSPI protocol."""

    _cfg: Config
    _proto: str

    def __init__(self, cfg: Config, proto: str) -> None:
        """Store the configuration object."""
        self._cfg = cfg
        self._proto = proto

    @property
    def cfg(self) -> Config:
        """Get the configuration for this runner."""
        return self._cfg

    @property
    def proto(self) -> str:
        """Get the name of the UCSPI protocol to test."""
        return self._proto

    @property
    def supports_remote_info(self) -> bool:
        """Whether the client and server support the -R command-line option."""
        return True

    @property
    def logs_to_stdout(self) -> bool:
        """Whether verbose output goes to the standard output stream (argh)."""
        return False

    @abc.abstractmethod
    def find_listening_address(self) -> list[str]:
        """Find an available protocol-specific address to listen on."""
        raise NotImplementedError

    @abc.abstractmethod
    def get_listening_socket(self, addr: list[str]) -> socket.socket:
        """Start listening on the specified address."""
        raise NotImplementedError(repr(addr))

    @abc.abstractmethod
    def get_connected_socket(self, addr: list[str]) -> socket.socket:
        """Connect to the specified address."""
        raise NotImplementedError(repr(addr))

    @abc.abstractmethod
    def format_local_addr(self, addr: list[str]) -> str:
        """Format an address returned by accept(), etc."""
        raise NotImplementedError(repr(addr))

    @abc.abstractmethod
    def format_remote_addr(self, addr: Any) -> str:  # noqa: ANN401
        """Format an address returned by accept(), etc."""
        raise NotImplementedError(repr(addr))

    def __repr__(self) -> str:
        """Provide a Python-esque representation of the helper."""
        return f"{type(self).__name__}(cfg={self._cfg!r}, proto={self._proto!r})"


HANDLERS: dict[str, type[Runner]] = {}


def with_listener(
    runner: Runner,
    addr: list[str],
    tag: str,
    callback: Callable[[socket.socket], None],
) -> None:
    """Get a listener socket from the runner, invoke the callback."""
    proto: Final = runner.proto
    saddr: Final = runner.format_local_addr(addr)
    print(f"- {tag}: setting up a listening socket at {saddr}")
    try:
        listener: Final = runner.get_listening_socket(addr)
    except RunnerError as err:
        sys.exit(f"{tag}: could not get a {proto} listening socket: {err}")

    try:
        callback(listener)
    finally:
        try:
            listener.close()
        except OSError as err:
            print(f"{tag}: could not close the listening socket at {saddr}: {err}", file=sys.stderr)


def with_connect(
    runner: Runner,
    addr: list[str],
    tag: str,
    callback: Callable[[socket.socket], None],
) -> None:
    """Ask the runner to connect to the specified address, invoke the callback."""
    proto: Final = runner.proto
    saddr: Final = runner.format_local_addr(addr)
    print(f"- {tag}: setting up a socket connected to {saddr}")
    try:
        conn: Final = runner.get_connected_socket(addr)
    except RunnerError as err:
        sys.exit(f"{tag}: could not get a {proto} connected socket: {err}")

    try:
        callback(conn)
    finally:
        try:
            conn.close()
        except OSError as err:
            print(f"{tag}: could not close the socket connected to {saddr}: {err}", file=sys.stderr)


def with_child(
    runner: Runner,
    tag: str,
    cmd: list[str],
    callback: Callable[[subprocess.Popen[str]], None],
    stderr: int | None = None,
) -> None:
    """Spawn a child process, invoke the callback."""
    print(f"- {tag}: starting {shlex.join(cmd)}")
    try:
        with subprocess.Popen(
            cmd,  # noqa: S603
            bufsize=0,
            encoding="UTF-8",
            env=runner.cfg.utf8_env,
            stdout=subprocess.PIPE,
            stderr=stderr,
        ) as child:
            try:
                callback(child)
            finally:
                if child.poll() is None:
                    ptag: Final = f"the {child.pid} ({cmd[0]}) process"
                    print(f"- {tag}: {ptag} is still active, killing it")
                    try:
                        child.kill()
                    except OSError as err:
                        print(
                            f"{tag}: could not send a kill signal to {ptag}: {err}",
                            file=sys.stderr,
                        )
                    print(f"- {tag}: waiting for {ptag} to go away")
                    try:
                        child.wait()
                    except OSError as err:
                        print(f"{tag}: could not wait for {ptag} to end: {err}", file=sys.stderr)
    except (OSError, subprocess.CalledProcessError) as err:
        sys.exit(f"Could not spawn {shlex.join(cmd)}: {err}")


def test_local_spew(runner: Runner, addr: list[str], tag: str, cmd: list[str]) -> None:
    """Test a client program against our listening socket, unidirectional transfer."""

    def run(listener: socket.socket, catproc: subprocess.Popen[str]) -> None:  # noqa: PLR0912
        """Run the spew test itself."""
        saddr: Final = runner.format_local_addr(addr)
        print(f"- started the client as pid {catproc.pid}")
        print(f"- waiting for a connection at {saddr}")
        try:
            conn, rem_addr = listener.accept()
        except OSError as err:
            sys.exit(f"Could not accept an incoming connection at {saddr}: {err}")
        print(
            f"- accepted a connection at fd {conn.fileno()} from "
            f"{runner.format_remote_addr(rem_addr)}",
        )

        msg: Final = MSG_RESP_HELLO + MSG_RESP_BYE
        try:
            conn.sendall(msg.encode("UTF-8"))
        except OSError as err:
            sys.exit(f"Could not send a message on the {conn} connection at {saddr}: {err}")

        print("- closing the incoming connection")
        try:
            conn.close()
        except OSError as err:
            sys.exit(f"Could not close the {conn} connection at {saddr}: {err}")

        try:
            output, _ = catproc.communicate()
        except OSError as err:
            sys.exit(f"Could not read the output of the {cmd[0]} process: {err}")
        if not isinstance(output, str):
            raise TypeError(repr(output))

        try:
            res: Final = catproc.wait()
        except OSError as err:
            sys.exit(f"Could not wait for the {cmd[0]} process to end: {err}")
        print(f"- client exit code: {res}; output: {output!r}")
        if res != 0:
            sys.exit(f"The {cmd[0]} program exited with non-zero code {res}")

        if output != msg:
            sys.exit(f"Expected {msg!r} as {cmd[0]} output, got {catproc.stdout!r}")

    with_listener(
        runner,
        addr,
        tag,
        lambda listener: with_child(runner, tag, cmd, lambda catproc: run(listener, catproc)),
    )
    print(f"- {cmd[0]} seems fine")


def test_local_cat(runner: Runner, addr: list[str]) -> None:
    """Test the {proto}cat program."""
    proto: Final = runner.proto
    print(f"\n=== Testing {proto}cat")
    if runner.cfg.bindir != pathlib.Path("/usr/bin"):
        print(
            f"- test skipped, {proto}cat will probably not find {runner.cfg.bindir}/{proto}client",
        )
        return

    catpath: Final = runner.cfg.bindir / f"{proto}cat"
    print(f"- will spawn {proto}cat at {catpath} in a while")
    test_local_spew(runner, addr, "test_local_cat", [str(catpath), *addr])


def remote_opt(runner: Runner) -> list[str]:
    """Add the -R command-line option if the client and server support it."""
    return ["-R"] if runner.supports_remote_info else []


def test_local_client_spew(runner: Runner, addr: list[str]) -> None:
    """Test {proto}client against our own listening socket."""
    proto: Final = runner.proto
    print(f"\n=== Testing {proto}client against our own listening socket")

    clipath: Final = runner.cfg.bindir / f"{proto}client"
    print(f"- will spawn {proto}client at {clipath} in a while")
    test_local_spew(
        runner,
        addr,
        "test_local_client_spew",
        [str(clipath), *remote_opt(runner), *addr, "sh", "-c", "set -e; exec <&6; exec cat"],
    )


def test_server_local(runner: Runner, addr: list[str]) -> None:  # noqa: C901,PLR0915
    """Test {proto}server against our own client socket."""

    def run(srvproc: subprocess.Popen[str], conn: socket.socket) -> None:  # noqa: PLR0912
        """Run the test itself."""
        try:
            rem_addr: Final = conn.getpeername()
        except OSError as err:
            sys.exit(f"getpeername() failed for {conn!r}: {err}")
        raddr: Final = runner.format_remote_addr(rem_addr)
        print(f"- got a connection at fd {conn.fileno()} to {raddr}")

        print("- reading all the data we can")
        data = b""
        while True:
            try:
                chunk = conn.recv(4096)
            except OSError as err:
                sys.exit(f"Could not read from the {conn!r} socket: {err}")
            print(f"- read {len(chunk)} bytes from the socket")
            if not chunk:
                break
            data += chunk
        print(f"- read a total of {len(data)} bytes from the socket")

        try:
            output: Final = data.decode("UTF-8")
        except ValueError as err:
            sys.exit(f"Could not decode {data!r} as valid UTF-8: {err}")
        if output != message:
            sys.exit(f"Expected {message!r}, got {output!r}")

        if srvproc.poll() is not None:
            sys.exit(
                f"Did not expect the {proto}server process to have exited: "
                f"code {srvproc.returncode}",
            )
        print(f"- sending a SIGTERM signal to the {proto}server process")
        try:
            srvproc.terminate()
        except OSError as err:
            sys.exit(
                f"Could not send a SIGTERM signal to the {proto}server process at "
                f"{srvproc.pid}: {err}",
            )
        print(f"- waiting for the {proto}server process to exit")
        res: Final = srvproc.wait()
        if res != 0:
            sys.exit(f"The {proto}server process exited with a non-zero code {res}")

    def wait_and_connect(srvproc: subprocess.Popen[str]) -> None:
        """Wait for the "ready" message from tcpserver."""
        stream: Final = srvproc.stdout if runner.logs_to_stdout else srvproc.stderr
        assert stream is not None, repr(srvproc)  # noqa: S101  # mypy needs this
        print(f"- awaiting the first 'status' line from {proto}server")
        line: Final = stream.readline()
        if "server: status: 0/" not in line:
            sys.exit(
                f"Unexpected first line from {proto}server: expected 'status: 0/N', got {line!r}",
            )
        with_connect(runner, addr, "test_server_local", lambda conn: run(srvproc, conn))

    proto: Final = runner.proto
    print(f"\n=== Testing {proto}server against our own listening socket")

    srvpath: Final = runner.cfg.bindir / f"{proto}server"
    print(f"- will spawn {proto}server at {srvpath} in a while")
    message: Final = MSG_RESP_HELLO + MSG_RESP_BYE
    with_child(
        runner,
        "test_server_local",
        [
            "stdbuf",
            "-oL",
            "-eL",
            "--",
            str(srvpath),
            "-v",
            *remote_opt(runner),
            *addr,
            "printf",
            "--",
            message.replace("\n", "\\n"),
        ],
        wait_and_connect,
        stderr=None if runner.logs_to_stdout else subprocess.PIPE,
    )

    print(f"- test_server_local: {srvpath} seems fine")


def test_server_client_spew(runner: Runner, addr: list[str]) -> None:  # noqa: C901
    """Test {proto}server against {proto}client, unidirectional data transfer."""

    def run(srvproc: subprocess.Popen[str], cliproc: subprocess.Popen[str]) -> None:
        """Read the data received by the client."""
        try:
            output, _ = cliproc.communicate()
        except OSError as err:
            sys.exit(f"Could not read the output of the {proto}client process: {err}")
        res_cli: Final = cliproc.poll()
        print(f"- client exit code {res_cli}; output {output!r}")
        if res_cli is None:
            sys.exit(f"Expected the {proto}client process to be done by now")
        if res_cli != 0:
            sys.exit(f"The {proto}client process exited with a non-zero code {res_cli}")
        if output != message:
            sys.exit(f"Expected {message!r}, got {output!r}")

        if srvproc.poll() is not None:
            sys.exit(
                f"Did not expect the {proto}server process to have exited: "
                f"code {srvproc.returncode}",
            )
        print(f"- sending a SIGTERM signal to the {proto}server process")
        try:
            srvproc.terminate()
        except OSError as err:
            sys.exit(
                f"Could not send a SIGTERM signal to the {proto}server process at "
                f"{srvproc.pid}: {err}",
            )
        print(f"- waiting for the {proto}server process to exit")
        res_srv: Final = srvproc.wait()
        if res_srv != 0:
            sys.exit(f"The {proto}server process exited with a non-zero code {res_srv}")

    def wait_and_connect(srvproc: subprocess.Popen[str]) -> None:
        """Wait for the "ready" message from tcpserver."""
        stream: Final = srvproc.stdout if runner.logs_to_stdout else srvproc.stderr
        assert stream is not None, repr(srvproc)  # noqa: S101  # mypy needs this
        print(f"- awaiting the first 'status' line from {proto}server")
        line: Final = stream.readline()
        if "server: status: 0/" not in line:
            sys.exit(
                f"Unexpected first line from {proto}server: expected 'status: 0/N', got {line!r}",
            )
        with_child(
            runner,
            "test_server_client_spew",
            [str(clipath), *remote_opt(runner), *addr, "sh", "-c", "set -e; exec <&6; exec cat"],
            lambda cliproc: run(srvproc, cliproc),
        )

    proto: Final = runner.proto
    print(f"\n=== Testing {proto}server against {proto}client")

    srvpath: Final = runner.cfg.bindir / f"{proto}server"
    print(f"- will spawn {proto}server at {srvpath} in a while")
    clipath: Final = runner.cfg.bindir / f"{proto}client"
    print(f"- will spawn {proto}client at {clipath} in a while")
    message: Final = MSG_RESP_HELLO + MSG_RESP_BYE
    with_child(
        runner,
        "test_server_client_spew",
        [
            "stdbuf",
            "-oL",
            "-eL",
            "--",
            str(srvpath),
            "-v",
            *remote_opt(runner),
            *addr,
            "printf",
            "--",
            message.replace("\n", "\\n"),
        ],
        wait_and_connect,
        stderr=None if runner.logs_to_stdout else subprocess.PIPE,
    )

    print(f"- {srvpath} seems fine")


def run_test(runner: Runner) -> None:
    """Run a couple of UCSPI tests."""
    addr: Final = runner.find_listening_address()

    test_local_cat(runner, addr)
    test_local_client_spew(runner, addr)

    test_server_local(runner, addr)
    test_server_client_spew(runner, addr)

    print(f"\n=== The tests for {runner.cfg.proto} passed")


def add_handler(proto: str, runner: type[Runner]) -> None:
    """Add a UCSPI protocol test runner."""
    current: Final = HANDLERS.get(proto)
    if current is None:
        HANDLERS[proto] = runner
    elif current != runner:
        raise HandlerMismatchError(proto, current, runner)


def run_test_handler(cfg: Config) -> None:
    """Parse command-line arguments, run the tests."""
    hprot: Final = HANDLERS.get(cfg.proto)
    if hprot is None:
        sys.exit(f"Don't know how to test the {cfg.proto!r} UCSPI protocol")

    run_test(hprot(cfg, cfg.proto))