bongo/contrib/testing/smtp-pipelining-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 SMTP PIPELINING response order and transaction recovery."""

from __future__ import annotations

import os
import re
import socket
import subprocess
import time

from bongo.store.StoreClient import StoreClient


HOST = os.environ.get("BONGO_TEST_HOST", "127.0.0.1")
PORT = int(os.environ.get("BONGO_TEST_SMTP_PORT", "25"))
STORE_PORT = int(os.environ.get("BONGO_TEST_STORE_PORT", "689"))
QUEUE_TOOL = os.environ.get(
    "BONGO_TEST_QUEUE_TOOL", "/usr/bin/bongo-queuetool"
)
USER = os.environ.get("BONGO_TEST_USER", "test1")
DOMAIN = os.environ.get("BONGO_TEST_DOMAIN", "bongo.test")
PASSWORD = os.environ.get("BONGO_TEST_PASSWORD", "")
TIMEOUT = float(os.environ.get("BONGO_TEST_TIMEOUT", "30"))
ALLOW_LIVE = os.environ.get("BONGO_ALLOW_LIVE_SMTP_TEST") == "1"
TOKEN = f"smtp12-{os.getpid()}-{int(time.time())}"
TOKEN_PREFIX = b"smtp12-"
RESPONSE = re.compile(rb"^([0-9]{3})([- ])(.*)\r\n$")
QUEUE_ID = re.compile(r"^[0-9A-Fa-f]{3}-[0-9A-Fa-f]+$")


class SMTP12Error(RuntimeError):
    """Raised when a pipelined SMTP session loses ordering or state."""


def run_queue(
    *arguments: str, binary: bool = False, check: bool = True
) -> subprocess.CompletedProcess:
    completed = subprocess.run(
        [
            "sudo",
            "-n",
            "-u",
            "bongo",
            QUEUE_TOOL,
            *arguments,
        ],
        capture_output=True,
        text=not binary,
        check=False,
    )
    if check and completed.returncode:
        stdout = completed.stdout
        stderr = completed.stderr
        if binary:
            stdout = stdout.decode("utf-8", "replace")
            stderr = stderr.decode("utf-8", "replace")
        raise SMTP12Error(
            f"{QUEUE_TOOL} {' '.join(arguments)} failed: "
            f"{stderr.strip() or stdout.strip()}"
        )
    return completed


class SMTPConnection:
    def __init__(self) -> None:
        self.socket = socket.create_connection((HOST, PORT), timeout=TIMEOUT)
        self.socket.settimeout(TIMEOUT)
        self.reader = self.socket.makefile("rb")

    def close(self) -> None:
        try:
            self.reader.close()
        finally:
            self.socket.close()

    def response(self) -> tuple[int, list[bytes]]:
        lines: list[bytes] = []
        response_code: int | None = None
        while True:
            line = self.reader.readline(8194)
            if not line:
                raise SMTP12Error("connection closed while awaiting response")
            if len(line) > 8192:
                raise SMTP12Error("SMTP response line exceeded 8192 bytes")
            match = RESPONSE.match(line)
            if match is None:
                raise SMTP12Error(f"malformed SMTP response: {line!r}")
            code = int(match.group(1))
            if response_code is None:
                response_code = code
            elif code != response_code:
                raise SMTP12Error(
                    f"multiline response changed status: {line!r}"
                )
            lines.append(match.group(3))
            if match.group(2) == b" ":
                return code, lines

    def require(self, expected: int, context: str) -> list[bytes]:
        code, lines = self.response()
        if code != expected:
            raise SMTP12Error(
                f"{context}: expected SMTP {expected}, got {code}: {lines!r}"
            )
        return lines

    def command(self, command: str, expected: int) -> list[bytes]:
        self.socket.sendall(command.encode("ascii") + b"\r\n")
        return self.require(expected, command)

    def pipeline(
        self,
        commands: list[str],
        expected: list[int],
        context: str,
    ) -> list[list[bytes]]:
        if len(commands) != len(expected):
            raise SMTP12Error("internal pipeline expectation mismatch")
        wire = "".join(f"{command}\r\n" for command in commands)
        self.socket.sendall(wire.encode("ascii"))
        return [
            self.require(code, f"{context} response {index + 1}")
            for index, code in enumerate(expected)
        ]


def queue_token_ids() -> list[str]:
    matches: list[str] = []
    for line in run_queue("list").stdout.splitlines():
        fields = line.split()
        if not fields or QUEUE_ID.fullmatch(fields[0]) is None:
            continue
        message = run_queue(
            "message", fields[0], binary=True, check=False
        )
        if message.returncode == 0 and TOKEN_PREFIX in message.stdout:
            matches.append(fields[0])
    return matches


def open_store() -> StoreClient:
    return StoreClient(
        USER,
        USER,
        authPassword=PASSWORD,
        host=HOST,
        port=STORE_PORT,
    )


def store_token_documents(store: StoreClient) -> list[str]:
    identifiers = [entry.uid for entry in store.List("/mail/INBOX")]
    matches: list[str] = []
    for identifier in identifiers:
        try:
            data = store.Read(identifier)
        except Exception:
            continue
        if TOKEN_PREFIX in data:
            matches.append(identifier)
    return matches


def cleanup(store: StoreClient) -> None:
    deadline = time.monotonic() + 20
    quiet_since: float | None = None
    while time.monotonic() < deadline:
        found = False
        for document in store_token_documents(store):
            found = True
            store.Delete(document)
        for queue_id in queue_token_ids():
            found = True
            run_queue("delete", queue_id, check=False)
        if found:
            quiet_since = None
        elif quiet_since is None:
            quiet_since = time.monotonic()
        elif time.monotonic() - quiet_since >= 1:
            return
        time.sleep(0.25)


def wait_for_delivery(store: StoreClient) -> str:
    deadline = time.monotonic() + TIMEOUT
    while time.monotonic() < deadline:
        documents = store_token_documents(store)
        if len(documents) > 1:
            raise SMTP12Error(
                f"pipelined message delivered more than once: {documents!r}"
            )
        if documents:
            return documents[0]
        time.sleep(0.25)
    raise SMTP12Error("pipelined message did not reach the Store")


def message() -> bytes:
    return (
        f"From: SMTP-12 <sender@example.test>\r\n"
        f"To: {USER}@{DOMAIN}\r\n"
        f"Subject: SMTP-12 pipelining {TOKEN}\r\n"
        f"Message-ID: <{TOKEN}@{DOMAIN}>\r\n"
        f"X-Bongo-Test: {TOKEN}\r\n"
        "\r\n"
        f"Pipelined delivery fixture {TOKEN}.\r\n"
    ).encode("ascii")


def main() -> int:
    if not ALLOW_LIVE:
        raise SMTP12Error(
            "set BONGO_ALLOW_LIVE_SMTP_TEST=1 for the live delivery test"
        )
    if not PASSWORD:
        raise SMTP12Error("BONGO_TEST_PASSWORD is required")

    store = open_store()
    try:
        cleanup(store)
        client = SMTPConnection()
        try:
            client.require(220, "greeting")
            capabilities = client.command(
                "EHLO smtp12-client.example", 250
            )
            if b"PIPELINING" not in capabilities:
                raise SMTP12Error(
                    f"EHLO did not advertise PIPELINING: {capabilities!r}"
                )

            client.pipeline(
                [
                    "MAIL FROM:<sender@example.test>",
                    f"RCPT TO:<{USER}@{DOMAIN}>",
                    "RSET",
                    "NOOP",
                ],
                [250, 250, 250, 250],
                "ordinary",
            )

            client.pipeline(
                [
                    "MAIL FROM:<sender@example.test>",
                    "RCPT TO:<victim@recipient.example>",
                    "DATA",
                ],
                [250, 554, 503],
                "rejected recipient",
            )
            client.pipeline(
                ["RSET", "NOOP"],
                [250, 250],
                "rejected recipient recovery",
            )

            client.pipeline(
                [
                    "MAIL FROM:<sender@example.test>",
                    "BOGUS",
                    f"RCPT TO:<{USER}@{DOMAIN}>",
                    "RSET",
                ],
                [250, 500, 250, 250],
                "unknown command",
            )

            client.pipeline(
                [
                    "MAIL FROM:<sender@example.test>",
                    "RCPT TO:<victim@recipient.example>",
                    f"RCPT TO:<{USER}@{DOMAIN}>",
                    "DATA",
                ],
                [250, 554, 250, 354],
                "mixed recipients and DATA",
            )
            client.socket.sendall(
                message()
                + b".\r\n"
                + b"NOOP\r\nRSET\r\nQUIT\r\n"
            )
            client.require(250, "message completion")
            client.require(250, "post-DATA pipelined NOOP")
            client.require(250, "post-DATA pipelined RSET")
            client.require(221, "post-DATA pipelined QUIT")
            client.socket.settimeout(2)
            if client.reader.read(1) != b"":
                raise SMTP12Error("server did not close after pipelined QUIT")
        finally:
            client.close()

        document = wait_for_delivery(store)
        store.Delete(document)
        cleanup(store)
        if store_token_documents(store) or queue_token_ids():
            raise SMTP12Error("test content remains after cleanup")
    finally:
        try:
            cleanup(store)
        finally:
            store.Quit()

    print(
        "SMTP-12 PASS advertised=yes ordinary-order=250/250/250/250 "
        "rejection-order=250/554/503/250/250 "
        "unknown-command=500/recovered mixed-recipients=554/250 "
        "data-sync=354 post-data=250/250/250/221 "
        "duplicate=no queue-clean=yes store-clean=yes"
    )
    return 0


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