bongo/contrib/testing/smtp-proxy-protocol-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 PROXY v1/v2 framing and trusted-source enforcement."""

from __future__ import annotations

import copy
import importlib.util
import json
import os
from pathlib import Path
import socket
import struct
import subprocess
import sys
import time


HOST = os.environ.get("BONGO_TEST_SMTP24_HOST", "127.0.0.1")
PORT = int(os.environ.get("BONGO_TEST_SMTP24_PORT", "10026"))
HAPROXY_PORT = int(
    os.environ.get("BONGO_TEST_SMTP24_HAPROXY_PORT", "11025")
)
TIMEOUT = float(os.environ.get("BONGO_TEST_TIMEOUT", "10"))
PROXY_TIMEOUT = 7.0
TRUSTED_SOURCE = "127.0.0.1/32"
UNTRUSTED_SOURCE = "192.0.2.0/24"
V1_CLIENT = "198.51.100.24"
V2_CLIENT = "203.0.113.25"


class SMTP24Error(RuntimeError):
    """Raised when the SMTP PROXY protocol contract is violated."""


def load_helpers():
    path = Path(__file__).with_name("smtp-forward-metadata-check.py")
    specification = importlib.util.spec_from_file_location(
        "bongo_smtp24_helpers", path
    )
    if specification is None or specification.loader is None:
        raise SMTP24Error(f"cannot load SMTP fixture helpers from {path}")
    module = importlib.util.module_from_spec(specification)
    sys.modules[specification.name] = module
    specification.loader.exec_module(module)
    return module


SMTP22 = load_helpers()


def proxy_v1(source: str = V1_CLIENT) -> bytes:
    return (
        f"PROXY TCP4 {source} {HOST} 43100 {PORT}\r\n"
    ).encode("ascii")


def proxy_v2(source: str = V2_CLIENT) -> bytes:
    payload = (
        socket.inet_aton(source)
        + socket.inet_aton(HOST)
        + struct.pack("!HH", 43101, PORT)
    )
    return (
        b"\r\n\r\n\x00\r\nQUIT\n"
        + bytes((0x21, 0x11))
        + struct.pack("!H", len(payload))
        + payload
    )


def proxy_configuration(
    original: dict,
    *,
    proxy_networks: list[str],
    relay_networks: list[str],
) -> dict:
    configuration = copy.deepcopy(original)
    configuration.update(
        {
            "internal_relay_enabled": True,
            "internal_relay_bind_address": HOST,
            "internal_relay_port": PORT,
            "internal_relay_networks": relay_networks,
            "proxy_protocol_enabled": True,
            "proxy_protocol_networks": proxy_networks,
        }
    )
    return configuration


class SMTPConnection:
    def __init__(self, host: str = HOST, port: int = PORT) -> 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 send(self, data: bytes) -> None:
        self.socket.sendall(data)

    def response(self) -> tuple[int, list[bytes]]:
        lines: list[bytes] = []
        code = -1
        while True:
            line = self.reader.readline(8194)
            if not line:
                raise SMTP24Error("connection closed before SMTP response")
            if len(line) > 8192 or len(line) < 5:
                raise SMTP24Error(f"invalid SMTP response line: {line!r}")
            try:
                current = int(line[:3])
            except ValueError as error:
                raise SMTP24Error(
                    f"non-numeric SMTP response: {line!r}"
                ) from error
            if line[3:4] not in {b"-", b" "}:
                raise SMTP24Error(f"invalid SMTP response framing: {line!r}")
            if code < 0:
                code = current
            elif current != code:
                raise SMTP24Error(f"mixed SMTP response codes: {lines + [line]!r}")
            lines.append(line)
            if line[3:4] == b" ":
                return code, lines


def smtp_dialog(
    header: bytes | None,
    *,
    host: str = HOST,
    port: int = PORT,
) -> None:
    connection = SMTPConnection(host, port)
    try:
        if header is not None:
            connection.send(header)
        code, _ = connection.response()
        if code != 220:
            raise SMTP24Error(f"expected greeting 220, received {code}")
        connection.send(b"EHLO smtp24.bongo.test\r\n")
        code, _ = connection.response()
        if code != 250:
            raise SMTP24Error(f"expected EHLO 250, received {code}")
        connection.send(b"QUIT\r\n")
        code, _ = connection.response()
        if code != 221:
            raise SMTP24Error(f"expected QUIT 221, received {code}")
    finally:
        connection.close()


def expect_silent_close(
    initial: bytes | None,
    *,
    host: str = HOST,
    port: int = PORT,
) -> None:
    connection = socket.create_connection((host, port), timeout=TIMEOUT)
    try:
        connection.settimeout(PROXY_TIMEOUT)
        if initial is not None:
            connection.sendall(initial)
        try:
            received = connection.recv(1024)
        except ConnectionResetError:
            received = b""
        except socket.timeout as error:
            raise SMTP24Error(
                "trusted proxy connection remained open past PROXY timeout"
            ) from error
        if received:
            raise SMTP24Error(
                f"invalid PROXY handshake received SMTP data: {received!r}"
            )
    finally:
        connection.close()


def wait_ready(header: bytes | None) -> None:
    deadline = time.monotonic() + 30
    last_error: BaseException | None = None
    while time.monotonic() < deadline:
        try:
            smtp_dialog(header)
            return
        except (OSError, SMTP24Error) as error:
            last_error = error
            time.sleep(0.2)
    raise SMTP24Error(
        f"internal SMTP did not become ready at {HOST}:{PORT}: {last_error}"
    )


def restart(header: bytes | None) -> None:
    SMTP22.SMTP07.run(
        SMTP22.privileged(
            [SMTP22.SMTP07.SYSTEMCTL, "restart", "bongo.service"]
        )
    )
    wait_ready(header)
    active = SMTP22.SMTP07.run(
        SMTP22.privileged(
            [SMTP22.SMTP07.SYSTEMCTL, "is-active", "bongo.service"]
        )
    ).stdout.strip()
    if active != "active":
        raise SMTP24Error(f"bongo.service is {active or 'not active'}")


def check_untrusted_forgery() -> None:
    connection = SMTPConnection()
    try:
        code, _ = connection.response()
        if code != 220:
            raise SMTP24Error(
                f"native untrusted client did not receive greeting: {code}"
            )
        connection.send(proxy_v1())
        code, _ = connection.response()
        if code < 500 or code > 599:
            raise SMTP24Error(
                f"unsolicited PROXY line was not rejected as SMTP: {code}"
            )
        connection.send(b"EHLO native.smtp24.bongo.test\r\n")
        code, _ = connection.response()
        if code != 250:
            raise SMTP24Error(
                "unsolicited PROXY line altered or terminated native session"
            )
        connection.send(b"QUIT\r\n")
        code, _ = connection.response()
        if code != 221:
            raise SMTP24Error("native SMTP session did not close cleanly")
    finally:
        connection.close()


def haproxy_configuration(public_port: int) -> str:
    return f"""\
global
    maxconn 16

defaults
    mode tcp
    timeout connect 5s
    timeout client 15s
    timeout server 15s

frontend fe_bongo_smtp24
    bind {HOST}:{HAPROXY_PORT}
    default_backend be_bongo_smtp24

backend be_bongo_smtp24
    server bongo {HOST}:{public_port} send-proxy-v2
"""


def check_haproxy_public_listener(public_port: int) -> None:
    configuration = Path(
        f"/tmp/bongo-smtp24-haproxy-{os.getpid()}.cfg"
    )
    process: subprocess.Popen[str] | None = None
    configuration.write_text(
        haproxy_configuration(public_port), encoding="ascii"
    )
    try:
        validation = subprocess.run(
            ["/usr/bin/haproxy", "-c", "-f", str(configuration)],
            capture_output=True,
            text=True,
            check=False,
        )
        if validation.returncode:
            raise SMTP24Error(
                "HAProxy rejected SMTP-24 fixture: "
                f"{validation.stderr.strip() or validation.stdout.strip()}"
            )
        process = subprocess.Popen(
            ["/usr/bin/haproxy", "-db", "-f", str(configuration)],
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
        )
        deadline = time.monotonic() + 10
        last_error: BaseException | None = None
        while time.monotonic() < deadline:
            if process.poll() is not None:
                output = process.stdout.read() if process.stdout else ""
                raise SMTP24Error(
                    "HAProxy fixture exited before accepting SMTP: "
                    f"{output.strip()}"
                )
            try:
                smtp_dialog(None, host=HOST, port=HAPROXY_PORT)
                break
            except (OSError, SMTP24Error) as error:
                last_error = error
                time.sleep(0.1)
        else:
            raise SMTP24Error(
                f"HAProxy SMTP frontend did not become ready: {last_error}"
            )

        # The same loopback source is configured as a proxy. Bypassing
        # HAProxy must therefore fail before Bongo emits its public greeting.
        expect_silent_close(None, host=HOST, port=public_port)
    finally:
        if process is not None and process.poll() is None:
            process.terminate()
            try:
                process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                process.kill()
                process.wait(timeout=5)
        configuration.unlink(missing_ok=True)


def restore(original: dict, original_host: str, original_port: int) -> None:
    SMTP22.replace_configuration("smtp", original)
    SMTP22.SMTP07.INTERNAL_HOST = original_host
    SMTP22.SMTP07.INTERNAL_PORT = original_port
    SMTP22.restart_bongo()


def perform_test() -> None:
    original = SMTP22.read_configuration("smtp")
    original_global = SMTP22.read_configuration("global")
    original_host = str(
        original.get("internal_relay_bind_address", "")
    ).strip()
    if not original_host or original_host in {"0.0.0.0", "::"}:
        original_host = str(original_global.get("hostaddr", "")).strip()
    if not original_host or original_host in {"0.0.0.0", "::"}:
        original_host = "127.0.0.1"
    original_port = int(original.get("internal_relay_port", 26))
    if not SMTP22.SMTP07.service_active():
        raise SMTP24Error("bongo.service must be active before SMTP-24")

    changed = False
    try:
        SMTP22.replace_configuration(
            "smtp",
            proxy_configuration(
                original,
                proxy_networks=[TRUSTED_SOURCE],
                relay_networks=[
                    f"{V1_CLIENT}/32",
                    "203.0.113.0/24",
                ],
            ),
        )
        changed = True
        restart(proxy_v1())
        smtp_dialog(proxy_v1())
        smtp_dialog(proxy_v2())
        check_haproxy_public_listener(int(original.get("port", 25)))
        expect_silent_close(
            f"PROXY TCP4 {V1_CLIENT} {HOST} 70000 {PORT}\r\n".encode(
                "ascii"
            )
        )
        expect_silent_close(None)

        SMTP22.replace_configuration(
            "smtp",
            proxy_configuration(
                original,
                proxy_networks=[TRUSTED_SOURCE],
                relay_networks=[TRUSTED_SOURCE],
            ),
        )
        restart(b"PROXY UNKNOWN\r\n")
        smtp_dialog(b"PROXY UNKNOWN\r\n")

        SMTP22.replace_configuration(
            "smtp",
            proxy_configuration(
                original,
                proxy_networks=[UNTRUSTED_SOURCE],
                relay_networks=[TRUSTED_SOURCE],
            ),
        )
        restart(None)
        smtp_dialog(None)
        check_untrusted_forgery()
    finally:
        if changed:
            restore(original, original_host, original_port)

    if SMTP22.read_configuration("smtp") != original:
        raise SMTP24Error("SMTP configuration was not restored exactly")
    if not SMTP22.SMTP07.service_active():
        raise SMTP24Error("bongo.service is not active after restoration")
    print(
        "SMTP-24 PASS proxy=v1/v2 haproxy=v2/public-port25 "
        "client-address=applied-before-port26-acl "
        "trusted=exact/cidr unknown=passthrough malformed=closed "
        "missing=timeout/closed untrusted=native/forgery-rejected "
        "config-restored=yes service-restored=active"
    )


def main() -> int:
    try:
        perform_test()
        return 0
    except (
        json.JSONDecodeError,
        OSError,
        SMTP24Error,
        SMTP22.SMTP22Error,
        SMTP22.SMTP07.SMTP07Error,
        ValueError,
    ) as error:
        print(f"smtp-proxy-protocol-check: {error}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())