bongo/contrib/testing/smtp-stress-timeout-check.py

#!/usr/bin/env python3
# This program is free software, licensed under the terms of the GNU GPL.
# See the Bongo COPYING file for full details.
# Copyright (c) 2026 Bongo Project contributors

"""Verify load-dependent SMTP timeouts with two real connections."""

from __future__ import annotations

import copy
import json
import os
import select
import socket
import subprocess
import time


ADMIN = os.environ.get("BONGO_TEST_ADMIN", "/usr/bin/bongo-admin")
SYSTEMCTL = os.environ.get("BONGO_TEST_SYSTEMCTL", "/usr/bin/systemctl")
INBOUND_HOST = os.environ.get("BONGO_TEST_INBOUND_HOST", "127.0.0.1")
INBOUND_PORT = int(os.environ.get("BONGO_TEST_INBOUND_PORT", "25"))
INTERNAL_HOST = os.environ.get(
    "BONGO_TEST_INTERNAL_HOST", "172.16.11.190")
INTERNAL_PORT = int(os.environ.get("BONGO_TEST_INTERNAL_PORT", "26"))
SUBMISSION_HOST = os.environ.get(
    "BONGO_TEST_SUBMISSION_HOST", "127.0.0.1")
SUBMISSION_PORT = int(os.environ.get("BONGO_TEST_SUBMISSION_PORT", "587"))
NORMAL_TIMEOUT = int(os.environ.get("BONGO_TEST_NORMAL_TIMEOUT", "5"))
STRESS_TIMEOUT = int(os.environ.get("BONGO_TEST_STRESS_TIMEOUT", "2"))
STARTUP_GRACE = float(os.environ.get("BONGO_TEST_STARTUP_GRACE", "2"))
ALLOW_LIVE = os.environ.get("BONGO_ALLOW_LIVE_SMTP_TEST") == "1"


class StressTimeoutError(RuntimeError):
    """Raised when SMTP does not apply its load-dependent timeout."""


def run(arguments: list[str], *, input_text: str | None = None) -> str:
    completed = subprocess.run(
        arguments,
        input=input_text,
        capture_output=True,
        text=True,
        check=False,
    )
    if completed.returncode:
        raise StressTimeoutError(
            f"{' '.join(arguments)} failed: "
            f"{completed.stderr.strip() or completed.stdout.strip()}")
    return completed.stdout


def admin(*arguments: str, input_text: str | None = None) -> str:
    return run(
        ["sudo", "-n", ADMIN, *arguments],
        input_text=input_text,
    )


def read_configuration() -> dict:
    try:
        configuration = json.loads(admin("__config-read", "smtp"))
    except json.JSONDecodeError as error:
        raise StressTimeoutError(
            "bongo-admin returned invalid SMTP configuration") from error
    if not isinstance(configuration, dict):
        raise StressTimeoutError("SMTP configuration is not an object")
    return configuration


def replace_configuration(configuration: dict) -> None:
    admin(
        "__config-replace",
        "smtp",
        input_text=json.dumps(configuration, separators=(",", ":")),
    )


def read_greeting(connection: socket.socket, name: str) -> None:
    data = bytearray()
    while b"\n" not in data and len(data) <= 8192:
        chunk = connection.recv(4096)
        if not chunk:
            raise StressTimeoutError(
                f"{name} closed before the SMTP greeting")
        data.extend(chunk)
    if not bytes(data).startswith(b"220 "):
        raise StressTimeoutError(
            f"{name} returned an invalid SMTP greeting: {bytes(data)!r}")


def open_connection(host: str, port: int, name: str) -> socket.socket:
    connection = socket.create_connection((host, port), timeout=10)
    connection.settimeout(10)
    try:
        read_greeting(connection, name)
    except BaseException:
        connection.close()
        raise
    return connection


def wait_closed(
    connection: socket.socket,
    started: float,
    expected: int,
    name: str,
) -> float:
    deadline = started + expected + 5
    received = bytearray()
    while time.monotonic() < deadline:
        connection.settimeout(max(0.1, deadline - time.monotonic()))
        try:
            chunk = connection.recv(4096)
        except socket.timeout:
            break
        if not chunk:
            elapsed = time.monotonic() - started
            if elapsed < max(0.5, expected * 0.65):
                raise StressTimeoutError(
                    f"{name} closed too early after {elapsed:.2f}s")
            return elapsed
        received.extend(chunk)
    raise StressTimeoutError(
        f"{name} remained open beyond {expected + 5}s; "
        f"unexpected data={bytes(received)!r}")


def assert_still_open(connection: socket.socket, name: str) -> None:
    readable, _, _ = select.select([connection], [], [], 0)
    if not readable:
        return
    data = connection.recv(4096)
    if not data:
        raise StressTimeoutError(
            f"{name} normal connection received the stress timeout")
    raise StressTimeoutError(
        f"{name} normal connection returned unexpected data: {data!r}")


def wait_ready(host: str, port: int, name: str) -> None:
    deadline = time.monotonic() + 30
    while time.monotonic() < deadline:
        connection: socket.socket | None = None
        try:
            connection = open_connection(host, port, name)
            connection.sendall(b"QUIT\r\n")
            reply = connection.recv(4096)
            if not reply.startswith(b"221 "):
                raise OSError("SMTP did not accept readiness QUIT")
            return
        except (OSError, StressTimeoutError):
            time.sleep(0.1)
        finally:
            if connection is not None:
                connection.close()
    raise StressTimeoutError(
        f"{name} did not become ready at {host}:{port}")


def restart_bongo() -> None:
    run(["sudo", "-n", SYSTEMCTL, "restart", "bongo.service"])
    time.sleep(STARTUP_GRACE)
    for name, host, port in (
        ("inbound", INBOUND_HOST, INBOUND_PORT),
        ("internal", INTERNAL_HOST, INTERNAL_PORT),
        ("submission", SUBMISSION_HOST, SUBMISSION_PORT),
    ):
        wait_ready(host, port, name)
    if run(
        ["sudo", "-n", SYSTEMCTL, "is-active", "bongo.service"]
    ).strip() != "active":
        raise StressTimeoutError("bongo.service is not active")


def check_endpoint(name: str, host: str, port: int) -> tuple[float, float]:
    normal: socket.socket | None = None
    stressed: socket.socket | None = None
    try:
        normal_started = time.monotonic()
        normal = open_connection(host, port, f"{name} normal")
        time.sleep(0.25)

        stress_started = time.monotonic()
        stressed = open_connection(host, port, f"{name} stressed")
        stress_elapsed = wait_closed(
            stressed, stress_started, STRESS_TIMEOUT, f"{name} stressed")
        stressed.close()
        stressed = None

        assert_still_open(normal, f"{name} normal")
        normal_elapsed = wait_closed(
            normal, normal_started, NORMAL_TIMEOUT, f"{name} normal")
        normal.close()
        normal = None

        if stress_elapsed >= normal_elapsed:
            raise StressTimeoutError(
                f"{name} stress timeout was not shorter than normal")
        return normal_elapsed, stress_elapsed
    finally:
        if stressed is not None:
            stressed.close()
        if normal is not None:
            normal.close()


def main() -> int:
    if not ALLOW_LIVE:
        raise StressTimeoutError(
            "set BONGO_ALLOW_LIVE_SMTP_TEST=1 for the live stress test")
    if not 1 <= STRESS_TIMEOUT < NORMAL_TIMEOUT <= 15:
        raise StressTimeoutError(
            "require 1 <= stress timeout < normal timeout <= 15 seconds")

    original = read_configuration()
    test_configuration = copy.deepcopy(original)
    test_configuration.update({
        "socket_timeout": NORMAL_TIMEOUT,
        "stress_socket_timeout": STRESS_TIMEOUT,
        "stress_load_percent": 100,
        "max_thread_load": 2,
    })
    results: list[str] = []
    test_error: BaseException | None = None
    restore_error: BaseException | None = None
    try:
        replace_configuration(test_configuration)
        restart_bongo()
        for name, host, port in (
            ("inbound", INBOUND_HOST, INBOUND_PORT),
            ("internal", INTERNAL_HOST, INTERNAL_PORT),
            ("submission", SUBMISSION_HOST, SUBMISSION_PORT),
        ):
            normal, stressed = check_endpoint(name, host, port)
            results.append(
                f"{name}=normal:{normal:.2f}s,stress:{stressed:.2f}s")
    except BaseException as error:
        test_error = error
    finally:
        try:
            replace_configuration(original)
            restart_bongo()
            if read_configuration() != original:
                raise StressTimeoutError(
                    "SMTP configuration was not restored exactly")
        except BaseException as error:
            restore_error = error

    if restore_error is not None:
        if test_error is not None:
            raise StressTimeoutError(
                f"test failed ({test_error}) and restoration also failed "
                f"({restore_error})")
        raise StressTimeoutError(
            f"failed to restore SMTP configuration: {restore_error}")
    if test_error is not None:
        raise test_error

    print(
        "SMTP-33 PASS "
        f"maximum-connections=2 stress-threshold=100% "
        f"normal-timeout={NORMAL_TIMEOUT}s stress-timeout={STRESS_TIMEOUT}s "
        f"sessions={';'.join(results)} "
        "config-restored=yes service=active"
    )
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (
        OSError,
        StressTimeoutError,
        subprocess.SubprocessError,
        ValueError,
    ) as error:
        print(f"SMTP-33 FAIL: {type(error).__name__}: {error}")
        raise SystemExit(1)