geometry/kernel.py

"""Sandboxed executor for model-written geometry generators (SPEC_OPTION_A.md §4.1).

`run(code, entry, params)` executes a generator `build(p) -> (verts, faces)` in a
**separate interpreter process** with a wall-clock timeout, an address-space /
CPU rlimit, and an import allowlist (numpy, math, itertools).

SECURITY NOTE: the restricted builtins + import allowlist are a *correctness*
guard (keep generators small and honest), NOT a security boundary — restricted
Python is leaky. The real isolation is the separate process + rlimits + timeout,
and, in production, the Modal function's network isolation. Do not rely on the
allowlist to contain hostile code on a networked host.
"""
import json
import os
import subprocess
import sys

import numpy as np


class KernelError(Exception):
    """Raised when a generator fails to execute or returns invalid geometry."""


# The runner executes inside a fresh `python -c` process. It reads a JSON request
# from stdin and prints exactly one JSON result line to stdout.
_RUNNER = r'''
import sys, json, math, resource
import builtins as _b
try:
    import numpy as np
except Exception:
    print(json.dumps({"tag": "err", "msg": "numpy import failed"})); sys.exit(0)

req = json.loads(sys.stdin.read())
code = req["code"]; entry = req.get("entry", "build"); params = req.get("params") or {}
mem_mb = int(req.get("mem_mb", 512)); cpu = int(req.get("cpu", 4))
for lim, val in ((resource.RLIMIT_CPU, (cpu, cpu)),
                 (resource.RLIMIT_AS, (mem_mb * 1024 * 1024, mem_mb * 1024 * 1024))):
    try:
        resource.setrlimit(lim, val)
    except Exception:
        pass  # best-effort; macOS does not enforce RLIMIT_AS reliably

ALLOWED = {"numpy", "math", "itertools"}
def _imp(name, *a, **k):
    if name.split(".")[0] not in ALLOWED:
        raise ImportError("disallowed import: " + name)
    return _b.__import__(name, *a, **k)

SAFE = {n: getattr(_b, n) for n in (
    "range", "len", "min", "max", "abs", "sum", "enumerate", "zip", "map",
    "filter", "list", "dict", "tuple", "set", "frozenset", "float", "int",
    "bool", "str", "round", "sorted", "reversed", "any", "all", "pow",
    "divmod", "print", "isinstance", "slice", "complex", "repr", "ValueError",
    "TypeError", "IndexError", "Exception")}
SAFE["__import__"] = _imp
g = {"__builtins__": SAFE}

# Inject the assembly helpers as `asm` (SPEC_ASSEMBLY.md §6) WITHOUT widening the
# generator's import allowlist: the runner imports the trusted module here (real
# builtins) and hands the generator only the safe callables. Best-effort.
try:
    sys.path.insert(0, req.get("geomdir", ""))
    import assembly as _asmmod
    g["asm"] = _asmmod.namespace()
except Exception:
    pass

try:
    exec(code, g)
    fn = g.get(entry)
    if not callable(fn):
        raise ValueError("no callable %r" % entry)
    verts, faces = fn(params)
    V = np.asarray(verts, dtype=float).reshape(-1, 3)
    if not np.isfinite(V).all():
        raise ValueError("non-finite vertices")
    n = len(V)
    F = []
    for f in faces:
        i, j, k = int(f[0]), int(f[1]), int(f[2]); cid = str(f[3])
        if not (0 <= i < n and 0 <= j < n and 0 <= k < n):
            raise ValueError("face index out of range")
        F.append([i, j, k, cid])
    if not F:
        raise ValueError("no faces returned")
    print(json.dumps({"tag": "ok", "verts": V.tolist(), "faces": F}))
except Exception as e:
    print(json.dumps({"tag": "err", "msg": repr(e)}))
'''


def run(code, entry="build", params=None, timeout=None, mem_mb=512):
    """Execute `code`'s `entry(params)` in a sandbox. Return (verts ndarray, faces list).

    Raises KernelError on timeout, disallowed import, exception, or invalid output.
    """
    if timeout is None:
        timeout = float(os.environ.get("TRIKONA_KERNEL_TIMEOUT", "3"))
    payload = json.dumps({
        "code": code, "entry": entry, "params": params or {},
        "geomdir": os.path.dirname(os.path.abspath(__file__)),  # for `import assembly`
        "mem_mb": mem_mb, "cpu": int(timeout) + 1,
    })
    try:
        proc = subprocess.run(
            [sys.executable, "-c", _RUNNER],
            input=payload, capture_output=True, text=True, timeout=timeout + 1.0,
        )
    except subprocess.TimeoutExpired:
        raise KernelError("timeout")
    out = (proc.stdout or "").strip()
    if not out:
        tail = (proc.stderr or "").strip()[-200:]
        raise KernelError("no output: " + (tail or "process crashed"))
    try:
        res = json.loads(out.splitlines()[-1])
    except Exception:
        raise KernelError("unparseable output")
    if res.get("tag") == "err":
        raise KernelError(res.get("msg", "unknown error"))
    return np.asarray(res["verts"], dtype=float), res["faces"]