Decoding in a box¶

"""The worker: set the limits on itself, then open, decode and write. Nothing else.

Run only by boxed.py. The limits are set here, in the child, rather than
with `preexec_fn` in the parent, which is unsafe in a program with threads.
On Linux `RLIMIT_AS` bounds what the process may address, so an allocation
beyond it fails instead of growing; macOS does not enforce it, and there
the wall-clock limit is the one that holds.
"""

from __future__ import annotations

import argparse
import resource
import sys


def limit(kind: int, value: int) -> None:
    soft, hard = resource.getrlimit(kind)
    if hard != resource.RLIM_INFINITY:
        value = min(value, hard)
    resource.setrlimit(kind, (value, value))


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("source")
    parser.add_argument("destination")
    parser.add_argument("--memory", type=int, required=True)
    parser.add_argument("--cpu-seconds", type=int, required=True)
    parser.add_argument("--max-output", type=int, required=True)
    parser.add_argument("--formats", required=True)
    args = parser.parse_args()

    limit(resource.RLIMIT_AS, args.memory)
    limit(resource.RLIMIT_CPU, args.cpu_seconds)
    limit(resource.RLIMIT_FSIZE, args.max_output)
    limit(resource.RLIMIT_CORE, 0)

    from PIL import Image

    Image.MAX_IMAGE_PIXELS = None  # the check already judged the size; the box is what holds it
    with Image.open(args.source, formats=args.formats.split(",")) as image:
        image.load()  # the decode: this is the line the box exists for
        image.save(args.destination, format="PNG")
    return 0


if __name__ == "__main__":
    sys.exit(main())

"""Decode an image in a boxed subprocess, after the check, so a flaw in the decoder costs one file.

The parent checks the file with sillybos and then runs the worker below in
a fresh process with the memory, CPU and output limits set inside it,
a wall-clock limit from outside, no environment, a scratch directory as
its working directory and nothing on standard input. The worker opens the
file with only the allowed readers, decodes it, and writes the pixels out
as a PNG of its own making. Whatever happens in the worker, the parent
gets either that PNG or a sentence saying why not.
"""

from __future__ import annotations

import os
import signal
import subprocess
import sys
import tempfile
from pathlib import Path

from sillybos import DEFAULT, Policy, Refused, require
from sillybos.pillow import pillow_formats

WORKER = Path(__file__).with_name("boxed_worker.py")
BASELINE_BYTES = 256 * 2**20  # the interpreter, Pillow and its codecs, with room
BYTES_PER_PIXEL = 8  # the decoded image and one working copy, four channels each
SECONDS_PER_MEGAPIXEL = 1.0


class DecodeRefused(Exception):
    """The box refused the file: the limit it hit and, where there is one, the decoder's own words."""


def memory_limit(policy: Policy) -> int:
    """What the worker may address, from the policy's pixel ceiling. Measure yours and adjust."""
    return BASELINE_BYTES + BYTES_PER_PIXEL * policy.max_pixels


def decode_in_box(
    path: Path, out: Path, policy: Policy = DEFAULT, *, memory: int | None = None
) -> Path:
    """Check `path`, then decode it in a limited subprocess, writing a PNG to `out`.

    `memory` is the worker's address-space limit in bytes; by default it is
    sized from the policy with `memory_limit`.
    """
    verdict = require(path, policy)  # Refused, with the reasons, before anything decodes
    assert verdict.facts is not None
    wall_seconds = (
        5 + SECONDS_PER_MEGAPIXEL * verdict.facts.max_pixels / 1_000_000
        if verdict.facts.max_pixels
        else 30
    )
    with tempfile.TemporaryDirectory() as scratch:
        try:
            done = subprocess.run(
                [
                    sys.executable,
                    "-I",
                    str(WORKER),
                    str(path.resolve()),
                    str(out.resolve()),
                    "--memory",
                    str(memory if memory is not None else memory_limit(policy)),
                    "--cpu-seconds",
                    str(int(wall_seconds) + 1),
                    "--max-output",
                    str(64 * 2**20),
                    "--formats",
                    ",".join(pillow_formats(policy)),
                ],
                cwd=scratch,
                env={"PATH": os.defpath},
                stdin=subprocess.DEVNULL,
                capture_output=True,
                text=True,
                errors="replace",
                timeout=wall_seconds,
                check=False,
            )
        except subprocess.TimeoutExpired:
            raise DecodeRefused(
                f"the decoder did not finish within {wall_seconds:.0f} seconds"
            ) from None
    if done.returncode == 0:
        return out
    if done.returncode < 0:
        name = signal.Signals(-done.returncode).name
        why = {
            "SIGKILL": "was killed, which is what running out of memory looks like",
            "SIGXCPU": "used more processor time than it was allowed",
            "SIGXFSZ": "tried to write more output than it was allowed",
            "SIGSEGV": "crashed",
        }.get(name, f"ended with {name}")
        raise DecodeRefused(f"the decoder {why}")
    last = (done.stderr.strip().splitlines() or ["no message"])[-1]
    raise DecodeRefused(f"the decoder failed: {last}")

1. a normal scan: /tmp/box/page.png (1200, 1600)
2. the bomb under the default policy: refused before decoding:
   ('a page of 144,000,000 pixels; the limit is 89,478,485',)
3. the bomb under a policy that allows it, so only the box stands in the
   way (the policy's own allowance would be 1782 MiB; the box is given
   512 MiB): refused by the box: the decoder failed: MemoryError after 0.1s
4. a normal scan under the same 512 MiB box: page2.png