#!/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)