timmy-talking-turd/tests/reencode-image.test.py

#!/usr/bin/env python3
"""Resource-bound contract for the image re-encoder subprocess.

Rejection of oversized or bomb-shaped images must happen from the container
header, before any full pixel decode, so a hostile upload cannot allocate
hundreds of megabytes inside a 512 MiB service. Decompression-bomb warnings are
treated as errors: a warning that still returns pixels is not a rejection.

Peak RSS is measured per child with os.wait4, so each measurement belongs to
exactly one subprocess rather than a running maximum.
"""
from __future__ import annotations

import json
import os
from pathlib import Path
import resource
import subprocess
import sys
import unittest

ROOT = Path(__file__).resolve().parents[1]
SCRIPT = ROOT / "scripts" / "reencode_image.py"
FIXTURES = ROOT / "tests" / "fixtures"
TMP = Path(os.environ.get("TMPDIR", "/tmp"))

# Interpreter + Pillow import costs ~21 MiB and a header-only open adds ~1 MiB.
# 64 MiB leaves generous headroom for that while remaining far below any full
# decode of the committed bomb fixtures (144 MP would need >140 MiB).
REJECTION_RSS_CEILING_KIB = 64 * 1024
# A hard address-space cap well below what a full decode of the committed bomb
# fixtures needs (~163 MiB observed). Header-only rejection must still complete
# cleanly under it, which proves the rejection never allocates the pixel buffer.
ADDRESS_SPACE_CEILING_BYTES = 112 * 1024 * 1024

EXIT_OK = 0
EXIT_REJECTED = 2
EXIT_UNAVAILABLE = 3


def run_measured(source: Path, *extra: str, out_name: str = "out.jpg",
                 env: dict | None = None, address_space: int | None = None):
    """Spawn the re-encoder in a fork we reap ourselves, for exact per-child rusage."""
    target = TMP / f"reencode-test-{out_name}"
    argv = [sys.executable, str(SCRIPT), "--in", str(source), "--out", str(target), *extra]
    stdout_r, stdout_w = os.pipe()
    stderr_r, stderr_w = os.pipe()
    pid = os.fork()
    if pid == 0:  # pragma: no cover - child process
        try:
            os.close(stdout_r)
            os.close(stderr_r)
            os.dup2(stdout_w, 1)
            os.dup2(stderr_w, 2)
            os.close(stdout_w)
            os.close(stderr_w)
            if address_space is not None:
                resource.setrlimit(resource.RLIMIT_AS, (address_space, address_space))
            os.execve(argv[0], argv, env or os.environ)
        except BaseException:
            os._exit(127)
    os.close(stdout_w)
    os.close(stderr_w)
    with os.fdopen(stdout_r, "r") as out_handle, os.fdopen(stderr_r, "r") as err_handle:
        stdout = out_handle.read()
        stderr = err_handle.read()
    _, status, usage = os.wait4(pid, 0)
    code = os.waitstatus_to_exitcode(status)
    return subprocess.CompletedProcess(argv, code, stdout, stderr), usage.ru_maxrss


def verdict(process: subprocess.CompletedProcess) -> dict:
    try:
        return json.loads(process.stdout.strip().splitlines()[-1])
    except (ValueError, IndexError):
        return {}


class ReencoderResourceBounds(unittest.TestCase):
    def test_bomb_dimensions_rejected_without_full_decode(self) -> None:
        process, peak_kib = run_measured(FIXTURES / "ingress-bomb.png", out_name="bomb.jpg")
        self.assertEqual(process.returncode, EXIT_REJECTED, process.stderr[:300])
        self.assertEqual(verdict(process).get("error"), "dimensions")
        self.assertLess(
            peak_kib, REJECTION_RSS_CEILING_KIB,
            f"12000x12000 rejection allocated {peak_kib} KiB; it must reject from the header",
        )

    def test_oversized_jpeg_rejected_without_full_decode(self) -> None:
        process, peak_kib = run_measured(FIXTURES / "ingress-oversized.jpg", out_name="oversized.jpg")
        self.assertEqual(process.returncode, EXIT_REJECTED, process.stderr[:300])
        self.assertEqual(verdict(process).get("error"), "dimensions")
        self.assertLess(
            peak_kib, REJECTION_RSS_CEILING_KIB,
            f"6000x6000 rejection allocated {peak_kib} KiB; it must reject from the header",
        )

    def test_bomb_fixtures_reject_under_a_hard_address_space_cap(self) -> None:
        # Proof of boundedness that does not depend on RSS sampling: under a
        # hard 192 MiB address-space limit the rejection must still complete
        # cleanly rather than dying from allocation failure.
        for name in ("ingress-bomb.png", "ingress-oversized.jpg"):
            with self.subTest(fixture=name):
                process, _ = run_measured(
                    FIXTURES / name, out_name=f"capped-{name}.jpg",
                    address_space=ADDRESS_SPACE_CEILING_BYTES,
                )
                self.assertEqual(process.returncode, EXIT_REJECTED, process.stderr[:300])
                self.assertEqual(verdict(process).get("error"), "dimensions")
                self.assertNotIn("MemoryError", process.stderr)

    def test_decompression_bomb_warnings_are_errors_not_warnings(self) -> None:
        process, _ = run_measured(FIXTURES / "ingress-bomb.png", out_name="warn.jpg")
        self.assertEqual(process.returncode, EXIT_REJECTED)
        self.assertNotIn(
            "DecompressionBombWarning", process.stderr,
            "a decompression bomb must raise an error, not emit a warning and continue",
        )

    def test_total_pixel_ceiling_is_enforced_independently(self) -> None:
        # Both dimensions are well inside the per-dimension cap; only the total
        # pixel budget rejects this, proving the ceiling exists on its own.
        process, _ = run_measured(
            FIXTURES / "ingress-clean.jpg", "--max-pixels", "1024", out_name="pixels.jpg",
        )
        self.assertEqual(process.returncode, EXIT_REJECTED)
        self.assertEqual(verdict(process).get("error"), "dimensions")

    def test_rejection_output_never_leaks_image_data_or_traces(self) -> None:
        for name in ("ingress-bomb.png", "ingress-oversized.jpg", "ingress-garbage.jpg"):
            with self.subTest(fixture=name):
                process, _ = run_measured(FIXTURES / name, out_name=f"leak-{name}.jpg")
                combined = process.stdout + process.stderr
                self.assertNotIn("Traceback", combined)
                self.assertLess(len(combined), 400, "subprocess output must stay short and fixed")

    def test_clean_image_still_reencodes_successfully(self) -> None:
        process, _ = run_measured(FIXTURES / "ingress-clean.jpg", out_name="clean.jpg")
        self.assertEqual(process.returncode, EXIT_OK, process.stderr[:300])
        result = verdict(process)
        self.assertTrue(result.get("ok"))
        self.assertEqual(result.get("format"), "jpeg")
        self.assertTrue(result.get("metadataStripped"))

    def test_missing_pillow_reports_unavailable_exit_code(self) -> None:
        # A production interpreter without Pillow must be distinguishable from
        # hostile client input, so it can map to processing-unavailable.
        stub = TMP / "reencode-no-pil"
        stub.mkdir(exist_ok=True)
        (stub / "PIL.py").write_text("raise ImportError('no pillow here')\n", encoding="utf-8")
        env = dict(os.environ, PYTHONPATH=str(stub))
        process, _ = run_measured(
            FIXTURES / "ingress-clean.jpg", out_name="nopil.jpg", env=env,
        )
        self.assertEqual(process.returncode, EXIT_UNAVAILABLE)
        self.assertEqual(verdict(process).get("error"), "unavailable")


class PinContract(unittest.TestCase):
    def test_verify_pin_reports_the_pinned_immutable_runtime(self) -> None:
        process, _ = run_measured(FIXTURES / "ingress-clean.jpg", "--verify-pin", out_name="pin.json")
        self.assertEqual(process.returncode, EXIT_OK, process.stderr[:300])
        result = verdict(process)
        self.assertTrue(result.get("ok"))
        self.assertEqual(result.get("python"), "3.11")
        self.assertEqual(result.get("pillow"), "12.3.0")
        self.assertEqual(result.get("pinned"), {"python": "3.11", "pillow": "12.3.0"})

    def test_verify_pin_fails_when_pillow_is_not_the_pinned_version(self) -> None:
        # Force a mismatched Pillow version to prove the gate refuses a drifted
        # runtime (server-side provisioning fault, not client input).
        stub = TMP / "reencode-fake-pil"
        stub.mkdir(exist_ok=True)
        (stub / "PIL.py").write_text(
            "class Image:\n    __version__ = '99.0.0'\n", encoding="utf-8")
        env = dict(os.environ, PYTHONPATH=str(stub))
        process, _ = run_measured(
            FIXTURES / "ingress-clean.jpg", "--verify-pin", out_name="pinbad.json", env=env,
        )
        self.assertEqual(process.returncode, EXIT_UNAVAILABLE)
        self.assertEqual(verdict(process).get("ok"), False)


if __name__ == "__main__":
    unittest.main(verbosity=2)