bongo/contrib/testing/queue-quota-partial-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 partial Queue delivery when a mailbox has limited free quota."""

from __future__ import annotations

import hashlib
import os
import re
import subprocess
import sys
import tempfile
from pathlib import Path

from bongo.store.StoreClient import StoreClient


USER = os.environ.get("BONGO_TEST_USER", "stqquota")
PASSWORD = os.environ.get("BONGO_TEST_PASSWORD", "")
ADMIN = os.environ.get("BONGO_TEST_ADMIN", "/usr/bin/bongo-admin")
QUEUE_TOOL = os.environ.get(
    "BONGO_TEST_QUEUE_TOOL", "/usr/bin/bongo-queuetool"
)
HOST = os.environ.get("BONGO_TEST_STORE_HOST", "127.0.0.1")
PORT = int(os.environ.get("BONGO_TEST_STORE_PORT", "689"))
MESSAGE_SIZE = int(os.environ.get("BONGO_TEST_QUOTA_MESSAGE_SIZE", "32768"))
ALLOW_LIVE = os.environ.get("BONGO_ALLOW_LIVE_USER_TEST") == "1"


class QuotaQueueTestError(RuntimeError):
    """Raised when the live Queue/Store quota contract is not met."""


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


def run_admin(*arguments: str) -> str:
    return run(["sudo", "-n", ADMIN, *arguments]).stdout


def run_queue(*arguments: str, check: bool = True, binary: bool = False):
    return run(
        ["sudo", "-n", "-u", "bongo", QUEUE_TOOL, *arguments],
        check=check,
        binary=binary,
    )


def read_admin_quota() -> tuple[int, int]:
    output = run_admin("user", "quota", USER)
    used_match = re.search(r"^Used: ([0-9]+) bytes$", output, re.MULTILINE)
    limit_match = re.search(
        r"^Quota: (?:(unlimited)|([0-9]+) bytes)$", output, re.MULTILINE
    )
    if not used_match or not limit_match:
        raise QuotaQueueTestError(
            f"cannot parse bongo-admin quota output: {output!r}"
        )
    limit = 0 if limit_match.group(1) else int(limit_match.group(2))
    return int(used_match.group(1)), limit


def set_admin_quota(limit: int) -> None:
    run_admin(
        "user",
        "quota",
        USER,
        "unlimited" if limit == 0 else str(limit),
    )


def create_message(marker: str) -> bytes:
    headers = (
        f"From: STQ-10 <stq10@bongo.test>\r\n"
        f"To: {USER}@bongo.test\r\n"
        f"Subject: STQ-10 partial quota delivery {marker}\r\n"
        f"Date: Thu, 23 Jul 2026 12:34:56 +0200\r\n"
        f"Message-ID: <stq10-{marker}@bongo.test>\r\n"
        "MIME-Version: 1.0\r\n"
        "Content-Type: text/plain; charset=UTF-8\r\n"
        "Content-Transfer-Encoding: 8bit\r\n"
        f"X-STQ-10-Marker: {marker}\r\n"
        "\r\n"
    ).encode()
    if len(headers) >= MESSAGE_SIZE:
        raise QuotaQueueTestError(
            "BONGO_TEST_QUOTA_MESSAGE_SIZE is too small"
        )
    return headers + marker.encode() + b":" + b"Q" * (
        MESSAGE_SIZE - len(headers) - len(marker) - 1
    )


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


def inbox_documents(store: StoreClient) -> dict[str, int]:
    return {
        document.uid: int(document.bodylen)
        for document in store.List("/mail/INBOX")
    }


def safe_queue_delete(queue_id: str | None) -> None:
    if queue_id:
        run_queue("delete", queue_id, check=False)


def safe_store_delete(store: StoreClient | None, document: str) -> None:
    if store is None:
        return
    try:
        store.Delete(document)
    except Exception as error:
        print(
            f"STQ-10 cleanup warning for Store document {document}: {error}",
            file=sys.stderr,
        )


def require_safe_environment() -> None:
    if not ALLOW_LIVE:
        raise QuotaQueueTestError(
            "set BONGO_ALLOW_LIVE_USER_TEST=1 for the disposable live account"
        )
    if not PASSWORD:
        raise QuotaQueueTestError("BONGO_TEST_PASSWORD must be set")
    if MESSAGE_SIZE < 4096:
        raise QuotaQueueTestError(
            "BONGO_TEST_QUOTA_MESSAGE_SIZE must be at least 4096"
        )


def main() -> int:
    require_safe_environment()
    baseline, original_limit = read_admin_quota()
    queue_ids: list[str] = []
    created_documents: list[str] = []
    store: StoreClient | None = None

    try:
        store = open_store()
        before = inbox_documents(store)
        queued_messages: list[bytes] = []

        with tempfile.TemporaryDirectory(prefix="bongo-stq10-") as work:
            work_path = Path(work)
            os.chmod(work, 0o755)
            for index in range(2):
                message_path = work_path / f"message-{index}.eml"
                message_path.write_bytes(
                    create_message(f"{os.getpid()}-{index}")
                )
                os.chmod(message_path, 0o644)
                queue_id = run_queue(
                    "hold-local", str(message_path)
                ).stdout.strip()
                if not queue_id.startswith("999-"):
                    raise QuotaQueueTestError(
                        f"unexpected held Queue ID: {queue_id!r}"
                    )
                queue_ids.append(queue_id)
                queued_messages.append(
                    run_queue(
                        "message", queue_id, binary=True
                    ).stdout
                )

        first_size = len(queued_messages[0])
        set_admin_quota(baseline + first_size)

        run_queue(
            "deliver-local",
            queue_ids[0],
            "stq10@bongo.test",
            USER,
            "INBOX",
        )
        after_first = inbox_documents(store)
        first_new = set(after_first) - set(before)
        if len(first_new) != 1:
            raise QuotaQueueTestError(
                "first delivery did not create exactly one Store document"
            )
        first_document = first_new.pop()
        created_documents.append(first_document)
        if store.Read(first_document) != queued_messages[0]:
            raise QuotaQueueTestError(
                "first Store document differs from its Queue message"
            )

        rejected = run_queue(
            "deliver-local",
            queue_ids[1],
            "stq10@bongo.test",
            USER,
            "INBOX",
            check=False,
        )
        rejected_output = rejected.stdout + rejected.stderr
        if rejected.returncode == 0:
            raise QuotaQueueTestError(
                "second delivery unexpectedly crossed the hard quota"
            )
        if "5220" not in rejected_output and "quota" not in rejected_output.lower():
            raise QuotaQueueTestError(
                "second delivery failed without a quota result: "
                + rejected_output.strip()
            )

        retained = run_queue(
            "message", queue_ids[1], binary=True
        ).stdout
        if retained != queued_messages[1]:
            raise QuotaQueueTestError(
                "quota rejection changed or removed the second Queue message"
            )
        retained_hash = hashlib.sha256(retained).hexdigest()

        store.Delete(first_document)
        created_documents.remove(first_document)
        used_after_cleanup, limit_after_cleanup = read_admin_quota()
        if (used_after_cleanup, limit_after_cleanup) != (
            baseline,
            baseline + first_size,
        ):
            raise QuotaQueueTestError(
                "deleting the first message did not release its quota"
            )

        run_queue(
            "deliver-local",
            queue_ids[1],
            "stq10@bongo.test",
            USER,
            "INBOX",
        )
        after_retry = inbox_documents(store)
        retry_new = set(after_retry) - set(before)
        if len(retry_new) != 1:
            raise QuotaQueueTestError(
                "retry did not create exactly one Store document"
            )
        retry_document = retry_new.pop()
        created_documents.append(retry_document)
        if store.Read(retry_document) != queued_messages[1]:
            raise QuotaQueueTestError(
                "retried Store document differs from retained Queue message"
            )

        store.Delete(retry_document)
        created_documents.remove(retry_document)
        restored_used, temporary_limit = read_admin_quota()
        if restored_used != baseline or temporary_limit != baseline + first_size:
            raise QuotaQueueTestError(
                "test Store cleanup did not restore baseline usage"
            )
    finally:
        for document in reversed(created_documents):
            safe_store_delete(store, document)
        if store is not None:
            try:
                store.Quit()
            except Exception:
                pass
        for queue_id in queue_ids:
            safe_queue_delete(queue_id)
        set_admin_quota(original_limit)

    restored_used, restored_limit = read_admin_quota()
    if (restored_used, restored_limit) != (baseline, original_limit):
        raise QuotaQueueTestError(
            "final cleanup did not restore the original quota state"
        )

    print(
        "STQ-10 PARTIAL PASS "
        f"user={USER} first-bytes={first_size} "
        f"second-bytes={len(queued_messages[1])} "
        f"retained-sha256={retained_hash} "
        "first-delivered=yes second-retained=yes retry-exact=yes "
        "quota-restored=yes"
    )
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (OSError, QuotaQueueTestError) as error:
        print(f"STQ-10 PARTIAL FAIL: {error}", file=sys.stderr)
        raise SystemExit(1)