Upload sandbox.py with huggingface_hub

sandbox.py

@@ -0,0 +1,420 @@
+"""Isolated execution of untrusted Python and Rust.
+
+The app compiles and runs code that arbitrary visitors typed into a web page.
+Running that in-process would hand every visitor the host's environment - which
+on a deployed instance includes the API keys. Everything here exists to keep
+those two worlds apart.
+
+Guarantees, strongest first:
+
+  1. Untrusted code runs in a *separate process* with an explicitly rebuilt
+     environment. No API key, token, or secret is inherited - the child gets a
+     hand-written allowlist of variables (PATH and the few the OS needs to
+     start a process at all), never a copy of os.environ.
+  2. Every run is wall-clock bounded and the whole process tree is killed on
+     timeout, so an infinite loop cannot pin a worker forever.
+  3. Work happens in a throwaway temp directory that is deleted afterwards, so
+     nothing can accumulate on disk or read the repository.
+  4. On Linux (how this deploys) POSIX rlimits additionally cap CPU seconds,
+     address space, file size, and forked processes.
+
+Item 4 is unavailable on Windows, so local dev is protected by 1-3 only. That
+is fine for running your own code; the deployed container is where hostile
+input actually arrives, and there the full set applies.
+
+What this is NOT: a security boundary strong enough to run genuinely malicious
+code with no further thought. A determined attacker inside the container is
+still inside the container. The container is therefore the real boundary - the
+process-level work here reduces blast radius, it does not replace isolation at
+the platform layer.
+"""
+
+from __future__ import annotations
+
+import os
+import platform
+import shutil
+import subprocess
+import sys
+import tempfile
+import time
+from dataclasses import dataclass
+
+IS_WINDOWS = platform.system() == "Windows"
+
+# Wall-clock ceilings, seconds.
+PYTHON_TIMEOUT = int(os.getenv("PYTHON_TIMEOUT", "60"))
+COMPILE_TIMEOUT = int(os.getenv("COMPILE_TIMEOUT", "90"))
+RUN_TIMEOUT = int(os.getenv("RUN_TIMEOUT", "60"))
+
+# POSIX-only resource ceilings.
+MAX_CPU_SECONDS = int(os.getenv("MAX_CPU_SECONDS", "60"))
+MAX_MEMORY_MB = int(os.getenv("MAX_MEMORY_MB", "2048"))
+MAX_OUTPUT_CHARS = 20_000
+
+
+@dataclass
+class Result:
+    """Outcome of one sandboxed run."""
+
+    ok: bool
+    stdout: str
+    stderr: str
+    seconds: float
+    stage: str = "run"  # "compile" or "run" - which step produced the error
+
+    @property
+    def display(self) -> str:
+        """Human-readable text for the UI output box."""
+        if self.ok:
+            return self.stdout or "(the program produced no output)"
+        label = "Compilation failed" if self.stage == "compile" else "Execution failed"
+        detail = (self.stderr or self.stdout or "").strip() or "no diagnostics available"
+        return f"{label}\n\n{detail}"
+
+
+def _safe_env() -> dict[str, str]:
+    """Build a minimal environment from scratch.
+
+    Deliberately constructed by allowlist rather than by copying os.environ and
+    deleting keys - a denylist silently fails the moment a new provider adds an
+    env var, and that failure mode leaks a credential.
+    """
+    allowed: dict[str, str] = {}
+
+    # The bare minimum for a process to start and find a toolchain.
+    passthrough = ["PATH"]
+    if IS_WINDOWS:
+        # Windows cannot create a process without these.
+        passthrough += ["SYSTEMROOT", "SYSTEMDRIVE", "COMSPEC", "PATHEXT",
+                        "NUMBER_OF_PROCESSORS", "PROCESSOR_ARCHITECTURE",
+                        "USERPROFILE", "WINDIR"]
+    else:
+        passthrough += ["LANG", "LC_ALL", "HOME"]
+
+    for name in passthrough:
+        value = os.environ.get(name)
+        if value:
+            allowed[name] = value
+
+    # Keep the child's Python quiet and deterministic.
+    allowed["PYTHONUNBUFFERED"] = "1"
+    allowed["PYTHONDONTWRITEBYTECODE"] = "1"
+
+    return allowed
+
+
+def assert_safe_to_deploy() -> None:
+    """Refuse to serve the public if secrets are sitting on disk.
+
+    Scrubbing the child's environment stops submitted code from *inheriting* a
+    key, but it cannot stop it from calling open('.env'). Verified by test: a
+    sandboxed program read the dotenv file in plaintext.
+
+    On a public host that path has to be closed by not having the file at all -
+    secrets belong in the platform's secret store, injected as environment
+    variables into the parent only. This turns that requirement into a
+    startup check instead of a thing to remember.
+    """
+    if os.getenv("PUBLIC_DEPLOYMENT", "0") != "1":
+        return
+
+    here = os.path.dirname(os.path.abspath(__file__))
+    offenders = [
+        name for name in (".env", ".env.local", ".env.production")
+        if os.path.exists(os.path.join(here, name))
+    ]
+    if offenders:
+        raise RuntimeError(
+            "Refusing to start in public mode with secret files on disk: "
+            f"{', '.join(offenders)}.\n"
+            "Submitted code runs on this filesystem and can read them. "
+            "Delete the file and supply secrets as environment variables "
+            "(on Hugging Face Spaces: Settings -> Variables and secrets)."
+        )
+
+
+def _preexec():
+    """Apply POSIX resource limits in the child before exec. Linux/macOS only."""
+    import resource
+
+    resource.setrlimit(resource.RLIMIT_CPU, (MAX_CPU_SECONDS, MAX_CPU_SECONDS))
+    mem = MAX_MEMORY_MB * 1024 * 1024
+    resource.setrlimit(resource.RLIMIT_AS, (mem, mem))
+    resource.setrlimit(resource.RLIMIT_FSIZE, (64 * 1024 * 1024, 64 * 1024 * 1024))
+    # No forking - blocks fork bombs and most attempts to outlive the timeout.
+    resource.setrlimit(resource.RLIMIT_NPROC, (64, 64))
+    # New process group so a timeout kill takes any children with it.
+    os.setsid()
+
+
+def _popen_kwargs() -> dict:
+    kwargs: dict = {}
+    if IS_WINDOWS:
+        # Lets us terminate the whole tree on timeout.
+        kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
+    else:
+        kwargs["preexec_fn"] = _preexec
+    return kwargs
+
+
+def _truncate(text: str) -> str:
+    if text is None:
+        return ""
+    if len(text) <= MAX_OUTPUT_CHARS:
+        return text
+    return text[:MAX_OUTPUT_CHARS] + f"\n\n... output truncated at {MAX_OUTPUT_CHARS:,} characters ..."
+
+
+def _kill_tree(proc: subprocess.Popen) -> None:
+    """Kill a process and everything it spawned."""
+    try:
+        if IS_WINDOWS:
+            subprocess.run(["taskkill", "/F", "/T", "/PID", str(proc.pid)],
+                           capture_output=True, timeout=10)
+        else:
+            import signal
+            os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
+    except Exception:
+        pass
+    try:
+        proc.wait(timeout=5)
+    except Exception:
+        pass
+
+
+def _run(cmd: list[str], cwd: str, timeout: int, stage: str) -> Result:
+    """Run one command under the full set of restrictions."""
+    started = time.perf_counter()
+    try:
+        proc = subprocess.Popen(
+            cmd,
+            cwd=cwd,
+            env=_safe_env(),
+            stdin=subprocess.DEVNULL,   # never block waiting on input
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            **_popen_kwargs(),
+        )
+    except FileNotFoundError as e:
+        return Result(False, "", f"Toolchain not found: {e}", 0.0, stage)
+
+    try:
+        stdout, stderr = proc.communicate(timeout=timeout)
+    except subprocess.TimeoutExpired:
+        _kill_tree(proc)
+        elapsed = time.perf_counter() - started
+        return Result(
+            False, "",
+            f"Timed out after {timeout}s and was terminated.\n"
+            f"The sandbox caps how long submitted code may run.",
+            elapsed, stage,
+        )
+
+    elapsed = time.perf_counter() - started
+    return Result(proc.returncode == 0, _truncate(stdout), _truncate(stderr), elapsed, stage)
+
+
+# --------------------------------------------------------------------------
+# public API
+# --------------------------------------------------------------------------
+
+
+def run_python(code: str) -> Result:
+    """Execute untrusted Python in a subprocess.
+
+    The notebook original used exec() in the web server's own process with full
+    builtins, which let any visitor read os.environ - and therefore every API
+    key on the host. This runs it out-of-process against a scrubbed environment
+    instead.
+    """
+    if not code.strip():
+        return Result(False, "", "No Python code to run.", 0.0)
+
+    with tempfile.TemporaryDirectory(prefix="pyrun_") as tmp:
+        path = os.path.join(tmp, "program.py")
+        with open(path, "w", encoding="utf-8") as f:
+            f.write(code)
+        # -I: isolated mode. Ignores PYTHON* env vars and keeps the script's
+        # directory off sys.path, so it cannot import anything from this repo.
+        return _run([sys.executable, "-I", "-S", path], tmp, PYTHON_TIMEOUT, "run")
+
+
+@dataclass(frozen=True)
+class Language:
+    """A compiled target the app can translate Python into."""
+
+    key: str            # internal id
+    display: str        # shown in the UI
+    extension: str      # source file suffix
+    highlight: str      # gradio Code component language
+    compilers: tuple[str, ...]   # executables to look for, in order of preference
+    install_hint: str
+
+    def find_compiler(self) -> str | None:
+        for name in self.compilers:
+            found = shutil.which(name)
+            if found:
+                return found
+        return None
+
+
+LANGUAGES: dict[str, Language] = {
+    "rust": Language(
+        key="rust",
+        display="Rust",
+        extension="rs",
+        highlight="cpp",     # gradio has no rust mode; cpp reads acceptably
+        compilers=("rustc",),
+        install_hint="https://rustup.rs",
+    ),
+    "cpp": Language(
+        key="cpp",
+        display="C++",
+        extension="cpp",
+        highlight="cpp",
+        compilers=("g++", "clang++", "c++"),
+        install_hint="https://gcc.gnu.org or https://winlibs.com on Windows",
+    ),
+}
+
+DEFAULT_LANGUAGE = "rust"
+
+
+def build_compile_command(language: Language, compiler: str, source: str,
+                          binary: str) -> list[str]:
+    """Optimisation flags per toolchain.
+
+    Both sets deliberately trade compile time for runtime - whole-program
+    optimisation, native CPU targeting, no unwinding machinery where it can be
+    dropped. Compile time is excluded from the measurement, so this is free.
+    """
+    if language.key == "rust":
+        return [
+            compiler, source,
+            "-C", "opt-level=3",
+            "-C", "target-cpu=native",
+            "-C", "codegen-units=1",
+            "-C", "lto=fat",
+            "-C", "panic=abort",
+            "-C", "strip=symbols",
+            "-o", binary,
+        ]
+    return [
+        compiler, source,
+        "-std=c++17",
+        "-O3",
+        "-march=native",
+        "-funroll-loops",
+        "-fno-plt",
+        "-flto",
+        "-o", binary,
+    ]
+
+
+def run_compiled(language_key: str, code: str) -> tuple[Result, str]:
+    """Compile and execute untrusted code. Returns (result, compile command)."""
+    language = LANGUAGES.get(language_key)
+    if language is None:
+        return Result(False, "", f"Unknown language: {language_key}", 0.0, "compile"), ""
+
+    if not code.strip():
+        return Result(False, "", f"No {language.display} code to compile.",
+                      0.0, "compile"), ""
+
+    compiler = language.find_compiler()
+    if not compiler:
+        return Result(
+            False, "",
+            f"No {language.display} compiler found on PATH "
+            f"(looked for: {', '.join(language.compilers)}).\n"
+            f"Install one from {language.install_hint} to run generated code.",
+            0.0, "compile",
+        ), ""
+
+    binary = "program.exe" if IS_WINDOWS else "program"
+    source_name = f"main.{language.extension}"
+
+    with tempfile.TemporaryDirectory(prefix=f"{language.key}run_") as tmp:
+        with open(os.path.join(tmp, source_name), "w", encoding="utf-8") as f:
+            f.write(code)
+
+        cmd = build_compile_command(language, compiler, source_name, binary)
+        pretty = " ".join([os.path.basename(cmd[0])] + cmd[1:])
+
+        compiled = _run(cmd, tmp, COMPILE_TIMEOUT, "compile")
+        if not compiled.ok:
+            return compiled, pretty
+
+        # Time execution only - trading compile time for runtime is the entire
+        # premise of these flags, so charging the result for it measures wrong.
+        return _run([os.path.join(tmp, binary)], tmp, RUN_TIMEOUT, "run"), pretty
+
+
+def available_languages() -> list[str]:
+    """Language keys with a working compiler on this machine."""
+    return [key for key, lang in LANGUAGES.items() if lang.find_compiler()]
+
+
+_probe_cache: dict[str, bool] = {}
+
+
+def cpp_has_int128() -> bool:
+    """Whether this C++ toolchain supports the __int128 extension.
+
+    Worth probing rather than assuming. It is a GCC/Clang extension present on
+    64-bit Linux (so the deployed container has it) but missing from MinGW on
+    Windows. Models reach for it unprompted when a benchmark needs wide
+    accumulators, and telling them "avoid non-portable extensions" is too weak
+    to stop it - stating the capability as a fact is not.
+    """
+    if "int128" in _probe_cache:
+        return _probe_cache["int128"]
+
+    lang = LANGUAGES["cpp"]
+    compiler = lang.find_compiler()
+    if not compiler:
+        _probe_cache["int128"] = False
+        return False
+
+    snippet = "unsigned __int128 x = 1; int main(){ return (int)(x & 1) - 1; }"
+    with tempfile.TemporaryDirectory(prefix="probe_") as tmp:
+        path = os.path.join(tmp, "probe.cpp")
+        with open(path, "w", encoding="utf-8") as f:
+            f.write(snippet)
+        out = os.path.join(tmp, "probe.exe" if IS_WINDOWS else "probe")
+        try:
+            proc = subprocess.run(
+                [compiler, "-std=c++17", "-fsyntax-only", path],
+                capture_output=True, text=True, timeout=25, env=_safe_env(),
+            )
+            ok = proc.returncode == 0
+        except Exception:
+            ok = False
+
+    _probe_cache["int128"] = ok
+    return ok
+
+
+def _version_of(path: str) -> str:
+    try:
+        out = subprocess.run([path, "--version"], capture_output=True,
+                             text=True, timeout=5)
+        return (out.stdout or out.stderr).strip().splitlines()[0]
+    except Exception:
+        return os.path.basename(path)
+
+
+def toolchain_status() -> dict[str, str]:
+    """Report what the sandbox can actually run, for the UI to display."""
+    status: dict[str, str] = {}
+    for key, lang in LANGUAGES.items():
+        found = lang.find_compiler()
+        status[key] = _version_of(found) if found else ""
+    status["python"] = f"Python {platform.python_version()}"
+    status["isolation"] = (
+        "full (rlimits + scrubbed env)" if not IS_WINDOWS
+        else "partial (scrubbed env, no rlimits on Windows)"
+    )
+    return status