Refactor to centralize AST security checks in security.py

- Moves shared AST blocklist logic to `security.py` for reuse across `preview_function` and `evaluation_function`.
- Updates `Dockerfile` to enable `nsjail`-based sandboxing for untrusted submissions.
- Adds unit tests for `check_code_safety` and integration tests for security gate in `evaluation_function`.
- Updates documentation to reflect the centralized security check, sandbox requirements, and testing changes.

Author: m-messer

CLAUDE.md

@@ -125,16 +125,19 @@ docker build -t evaluatepython .
 # Cross-platform (CI uses linux/x86_64):
 docker build --platform=linux/x86_64 .
 
-# Run the server locally (port 8080)
-docker run -it --rm -p 8080:8080 evaluatepython
+# Run the server locally (port 8080).
+# --privileged is required: the image enables shimmy's nsjail sandbox
+# (SANDBOX_ENABLED=true), and nsjail needs it to create namespaces.
+docker run -it --rm --privileged -p 8080:8080 evaluatepython
 ```
 
 ## Tests
 
-Two test files, run with `pytest`:
+Three test files, run with `pytest`:
 
-- `evaluation_function/evaluation_test.py` — integration tests covering: all modes (demo, io_test, unit_test), all pass, partial fail, hidden test failure, runtime error, matplotlib plot capture, Hypothesis support
+- `evaluation_function/evaluation_test.py` — integration tests covering: all modes (demo, io_test, unit_test), all pass, partial fail, hidden test failure, runtime error, matplotlib plot capture, Hypothesis support, the pre-execution security gate
 - `evaluation_function/preview_test.py` — unit tests covering: valid Python, syntax errors, dangerous imports, dangerous builtins, dunder access
+- `evaluation_function/security_test.py` — unit tests for `check_code_safety` (the shared AST blocklist used by both `preview_function` and `evaluation_function`)
 
 CI runs on Python 3.12 and uploads JUnit XML results (`.github/workflows/test-lint.yml`).
 
@@ -146,14 +149,18 @@ CI runs on Python 3.12 and uploads JUnit XML results (`.github/workflows/test-li
 | `MPLBACKEND` | `Agg` | Set at subprocess runtime to suppress GUI |
 | `FUNCTION_COMMAND` | `python` | lf_toolkit runner |
 | `FUNCTION_ARGS` | `-m,evaluation_function.main` | lf_toolkit runner |
-| `FUNCTION_RPC_TRANSPORT` | `ipc` | lf_toolkit transport |
+| `FUNCTION_RPC_TRANSPORT` | `stdio` | shimmy↔worker transport (stdio so it survives the sandbox mount namespace) |
 | `LOG_LEVEL` | `debug` | Logging verbosity |
 | `AWS_*` / boto3 credentials | Runtime env | Required for S3 plot uploads |
+| `SANDBOX_ENABLED` | `true` | Wrap the worker in shimmy's nsjail sandbox (needs `--privileged` at run time) |
+| `SANDBOX_SECCOMP` | `true` | nsjail seccomp syscall filter |
+| `SANDBOX_RO_BINDS` | `/usr:/lib:/lib64:/bin:/sbin:/etc:/app` | Read-only bind mounts visible inside the jail |
+| `SANDBOX_TMPFS` | `/tmp` | Writable tmpfs inside the jail (student scripts, plot dirs, MPLCONFIGDIR) |
 
 Dependencies managed via Poetry; `.venv` is created in-project (`poetry.toml`).
 
 ## Deployment
 
 - Push to `main` triggers GitHub Actions (`.github/workflows/`) which builds and deploys to Lambda Feedback automatically
 - The function name is declared in `config.json` as `EvaluationFunctionName: "evaluatePython"` (lowerCamelCase)
-- The base Docker image is `ghcr.io/lambda-feedback/evaluation-function-base/python:test-sandbox-3.12`
+- The base Docker image is `ghcr.io/lambda-feedback/evaluation-function-base/python:3.12` (bundles shimmy + nsjail; sandboxing is enabled via the `SANDBOX_*` env vars in the Dockerfile, not by the base tag)

Dockerfile

@@ -1,4 +1,4 @@
-FROM ghcr.io/lambda-feedback/evaluation-function-base/python:test-sandbox-3.12 AS builder
+FROM ghcr.io/lambda-feedback/evaluation-function-base/python:3.12 AS builder
 
 RUN pip install poetry==1.8.3
 
@@ -12,7 +12,7 @@ COPY pyproject.toml poetry.lock ./
 RUN --mount=type=cache,target=$POETRY_CACHE_DIR \
     poetry install --without dev --no-root
 
-FROM ghcr.io/lambda-feedback/evaluation-function-base/python:test-sandbox-3.12
+FROM ghcr.io/lambda-feedback/evaluation-function-base/python:3.12
 
 ENV VIRTUAL_ENV=/app/.venv \
     PATH="/app/.venv/bin:$PATH"
@@ -31,7 +31,38 @@ ENV FUNCTION_COMMAND="python"
 # Args to start the evaluation function with
 ENV FUNCTION_ARGS="-m,evaluation_function.main"
 
-# The transport to use for the RPC server
-ENV FUNCTION_RPC_TRANSPORT="ipc"
+# The transport to use for the RPC server.
+# stdio (not ipc): the worker runs inside an nsjail mount namespace with a
+# private tmpfs /tmp (see the sandbox settings below), so a host unix-socket
+# rendezvous at /tmp/eval.sock would be unreachable. shimmy's sandbox is
+# designed around stdio; lf_toolkit routes its logs to stderr, keeping stdout
+# clean for the RPC framing.
+ENV FUNCTION_RPC_TRANSPORT="stdio"
+
+# --- Sandboxed execution of untrusted student code (shimmy + nsjail) ---
+# shimmy wraps the worker process -- and every `python` subprocess it spawns to
+# run a submission -- in an nsjail sandbox: unprivileged uid (nobody:nogroup),
+# a minimal bind-mounted view of the filesystem, and seccomp syscall filtering.
+#
+# Run-time requirements (cannot be expressed in the image):
+#   * the container must run with --privileged (or --cap-add SYS_ADMIN) so
+#     nsjail can create its namespaces -- see the shimmy README, "Sandboxed
+#     Execution". Without it the worker will not boot.
+#   * nsjail must exist at /usr/sbin/nsjail (provided by the base image's
+#     shimmy stage).
+# If the worker fails to start, drop SANDBOX_SECCOMP first, then widen
+# SANDBOX_RO_BINDS (the list is linux/x86_64 + Debian-specific).
+#
+# Network is deliberately left enabled -- the function uploads matplotlib
+# plots to S3 via boto3. Untrusted network/filesystem use is blocked one layer
+# up, at the AST gate in evaluation_function/security.py (check_code_safety).
+#
+# No CPU/memory rlimits here: the RPC worker is long-lived and shared across
+# requests, so a cumulative RLIMIT_CPU / RLIMIT_AS would eventually kill it.
+# Per-execution wall-clock limits are enforced in evaluation.py (_TIMEOUT).
+ENV SANDBOX_ENABLED="true" \
+    SANDBOX_SECCOMP="true" \
+    SANDBOX_RO_BINDS="/usr:/lib:/lib64:/bin:/sbin:/etc:/app" \
+    SANDBOX_TMPFS="/tmp"
 
 ENV LOG_LEVEL="debug"

README.md

@@ -12,11 +12,13 @@ Push to `main` triggers GitHub Actions which automatically builds and deploys to
 ### Run the Docker Image
 
 ```bash
-docker run -it --rm -p 8080:8080 ghcr.io/lambda-feedback/evaluatepython:latest
+docker run -it --rm --privileged -p 8080:8080 ghcr.io/lambda-feedback/evaluatepython:latest
 ```
 
 The image includes [Shimmy](https://github.com/lambda-feedback/shimmy), which listens for HTTP requests on port 8080 and forwards them to the evaluation function.
 
+`--privileged` (or `--cap-add SYS_ADMIN`) is required: the image runs student code inside Shimmy's [nsjail](https://github.com/google/nsjail) sandbox (`SANDBOX_ENABLED=true`), and nsjail needs those privileges to create its namespaces. The sandbox runs the worker as `nobody` with a minimal bind-mounted filesystem and seccomp filtering; untrusted imports/builtins are additionally rejected before execution by the AST check in `evaluation_function/security.py`.
+
 ### Evaluation Modes
 
 The function supports three modes, set via `params.mode`.
@@ -80,10 +82,12 @@ Add `"pep8_feedback": true` to any mode to append a style check to the feedback.
 ```
 evaluation_function/main.py             # IPC server entry point
 evaluation_function/evaluation.py       # core evaluation pipeline (all three modes)
-evaluation_function/preview.py          # AST-based security validator
+evaluation_function/security.py         # shared AST blocklist (check_code_safety)
+evaluation_function/preview.py          # pre-submission validator (syntax + security)
 evaluation_function/dev.py              # CLI wrapper for local testing
 evaluation_function/evaluation_test.py  # integration tests
 evaluation_function/preview_test.py     # preview/security tests
+evaluation_function/security_test.py    # check_code_safety unit tests
 config.json                             # deployment configuration
 ```
 
@@ -144,9 +148,11 @@ docker build --platform=linux/x86_64 -t evaluatepython .
 ### Running the Docker Image
 
 ```bash
-docker run -it --rm -p 8080:8080 evaluatepython
+docker run -it --rm --privileged -p 8080:8080 evaluatepython
 ```
 
+`--privileged` is required for the nsjail sandbox (see [Run the Docker Image](#run-the-docker-image)). To run without it for local debugging, disable the sandbox: add `-e SANDBOX_ENABLED=false`.
+
 ## Deployment to Lambda Feedback
 
 The function name is declared in [`config.json`](config.json) as `"evaluatePython"` (lowerCamelCase). Pushing to `main` triggers automated deployment via GitHub Actions.

docs/dev.md

@@ -7,6 +7,12 @@
 - **`evaluation_function`** — run and grade student code
 - **`preview_function`** — static AST-based security check (called before execution)
 
+Both share the AST blocklist in `security.py` (`check_code_safety`). `preview_function`
+runs it pre-submission; `evaluation_function` also runs it as a hard gate — if the
+submission trips the blocklist it is **not executed** and an `error` item is returned.
+At the container level, shimmy additionally wraps execution in an nsjail sandbox
+(`SANDBOX_*` env in the `Dockerfile`).
+
 ---
 
 ## `evaluation_function`
@@ -173,7 +179,9 @@ The check runs in-process via `pycodestyle.Checker` with `max_line_length=200` (
 
 ## `preview_function`
 
-Called before evaluation. Parses the student code as an AST and checks for security violations.
+Called before evaluation. Parses the student code as an AST and checks for security
+violations via `security.check_code_safety`. The same check gates `evaluation_function`,
+so the blocklist below applies to both entry points.
 
 ### Blocked constructs
 

evaluation_function/evaluation.py

@@ -11,6 +11,8 @@
 from lf_toolkit.evaluation import Result, Params
 from lf_toolkit.evaluation.image_upload import upload_image, ImageUploadError
 
+from .security import check_code_safety
+
 _TIMEOUT = 25
 _UPLOAD_FOLDER = "evaluatePython"
 
@@ -303,6 +305,14 @@ def evaluation_function(response: Any, answer: Any, params: Params) -> Result:
         result.add_feedback("error", f"Unknown or missing mode: {mode!r}. Expected 'demo', 'io_test', or 'unit_test'.")
         return result
 
+    violations = check_code_safety(str(response))
+    if violations:
+        result.add_feedback(
+            "error",
+            "Unsafe code detected -- not executed:\n" + "\n".join(f"- {v}" for v in violations),
+        )
+        return result
+
     if mode == "demo":
         result = _evaluate_demo(str(response), result)
     elif mode == "io_test":

evaluation_function/evaluation_test.py

@@ -330,4 +330,40 @@ def test_pep8_appended_to_io_test_result(self):
         params = {**_params(_test("5\n", "25\n")), "pep8_feedback": True}
         result = evaluation_function(_SQUARE_CODE, None, params).to_dict()
         self.assertTrue(result["is_correct"])
-        self.assertIn("No style issues found", result["feedback"])
+        self.assertIn("No style issues found", result["feedback"])
+
+
+class TestSecurityGate(unittest.TestCase):
+
+    def test_blocked_import_not_executed(self):
+        result = evaluation_function(
+            "import os\nprint(os.getcwd())", None, {"mode": "demo"}
+        ).to_dict()
+        self.assertFalse(result["is_correct"])
+        self.assertIn("Unsafe code detected", result["feedback"])
+        self.assertIn("os", result["feedback"])
+
+    def test_blocked_builtin_not_executed(self):
+        result = evaluation_function(
+            "eval('1 + 1')", None, _params(_test("", ""))
+        ).to_dict()
+        self.assertIn("Unsafe code detected", result["feedback"])
+        self.assertNotIn("tests passed", result["feedback"])
+
+    def test_dunder_access_blocked_in_unit_test(self):
+        result = evaluation_function(
+            "().__class__.__bases__", None, _unit_params(_SQUARE_TESTS)
+        ).to_dict()
+        self.assertIn("Unsafe code detected", result["feedback"])
+
+    def test_safe_stdlib_import_still_runs(self):
+        result = evaluation_function(
+            "import math\nprint(math.sqrt(16))", None, {"mode": "demo"}
+        ).to_dict()
+        self.assertIn("4.0", result["feedback"])
+        self.assertNotIn("Unsafe", result["feedback"])
+
+    def test_syntax_error_passes_gate_and_surfaces_downstream(self):
+        result = evaluation_function("def f(:\n", None, {"mode": "demo"}).to_dict()
+        self.assertNotIn("Unsafe code detected", result["feedback"])
+        self.assertIn("SyntaxError", result["feedback"])

evaluation_function/preview.py

@@ -2,55 +2,19 @@
 from typing import Any
 from lf_toolkit.preview import Result, Params, Preview
 
-_BLOCKED_MODULES = {
-    "os", "sys", "subprocess", "socket", "urllib", "http",
-    "requests", "shutil", "pathlib", "ftplib", "smtplib",
-    "ctypes", "multiprocessing", "threading", "importlib",
-    "pickle", "builtins",
-}
-
-_BLOCKED_BUILTINS = {"exec", "eval", "compile", "open", "__import__"}
-
-
-class _SecurityVisitor(ast.NodeVisitor):
-    def __init__(self):
-        self.violations: list[str] = []
-
-    def visit_Import(self, node):
-        for alias in node.names:
-            root = alias.name.split(".")[0]
-            if root in _BLOCKED_MODULES:
-                self.violations.append(f"import of '{root}' is not allowed")
-        self.generic_visit(node)
-
-    def visit_ImportFrom(self, node):
-        if node.module:
-            root = node.module.split(".")[0]
-            if root in _BLOCKED_MODULES:
-                self.violations.append(f"import of '{root}' is not allowed")
-        self.generic_visit(node)
-
-    def visit_Call(self, node):
-        if isinstance(node.func, ast.Name) and node.func.id in _BLOCKED_BUILTINS:
-            self.violations.append(f"use of '{node.func.id}()' is not allowed")
-        self.generic_visit(node)
-
-    def visit_Attribute(self, node):
-        if node.attr.startswith("__") and node.attr.endswith("__"):
-            self.violations.append(f"access to '{node.attr}' is not allowed")
-        self.generic_visit(node)
+from .security import check_code_safety
 
 
 def preview_function(response: Any, params: Params) -> Result:
+    code = str(response)
     try:
-        tree = ast.parse(str(response))
+        ast.parse(code)
     except SyntaxError as e:
         return Result(preview=Preview(feedback=f"SyntaxError: {e.msg} (line {e.lineno})"))
 
-    visitor = _SecurityVisitor()
-    visitor.visit(tree)
-    if visitor.violations:
-        lines = "\n".join(f"- {v}" for v in visitor.violations)
+    violations = check_code_safety(code)
+    if violations:
+        lines = "\n".join(f"- {v}" for v in violations)
         return Result(preview=Preview(feedback=f"Unsafe code detected:\n{lines}"))
 
     return Result(preview=Preview(feedback="Valid Python syntax."))

evaluation_function/security.py

@@ -0,0 +1,55 @@
+import ast
+
+_BLOCKED_MODULES = {
+    "os", "sys", "subprocess", "socket", "urllib", "http",
+    "requests", "shutil", "pathlib", "ftplib", "smtplib",
+    "ctypes", "multiprocessing", "threading", "importlib",
+    "pickle", "builtins",
+}
+
+_BLOCKED_BUILTINS = {"exec", "eval", "compile", "open", "__import__"}
+
+
+class _SecurityVisitor(ast.NodeVisitor):
+    def __init__(self):
+        self.violations: list[str] = []
+
+    def visit_Import(self, node):
+        for alias in node.names:
+            root = alias.name.split(".")[0]
+            if root in _BLOCKED_MODULES:
+                self.violations.append(f"import of '{root}' is not allowed")
+        self.generic_visit(node)
+
+    def visit_ImportFrom(self, node):
+        if node.module:
+            root = node.module.split(".")[0]
+            if root in _BLOCKED_MODULES:
+                self.violations.append(f"import of '{root}' is not allowed")
+        self.generic_visit(node)
+
+    def visit_Call(self, node):
+        if isinstance(node.func, ast.Name) and node.func.id in _BLOCKED_BUILTINS:
+            self.violations.append(f"use of '{node.func.id}()' is not allowed")
+        self.generic_visit(node)
+
+    def visit_Attribute(self, node):
+        if node.attr.startswith("__") and node.attr.endswith("__"):
+            self.violations.append(f"access to '{node.attr}' is not allowed")
+        self.generic_visit(node)
+
+
+def check_code_safety(code: str) -> list[str]:
+    """Return the list of security violations in ``code`` (empty means safe).
+
+    Shared by ``preview_function`` (pre-submission check) and
+    ``evaluation_function`` (pre-execution gate). A syntax error is not a
+    safety violation -- it is surfaced elsewhere -- so it yields an empty list.
+    """
+    try:
+        tree = ast.parse(code)
+    except SyntaxError:
+        return []
+    visitor = _SecurityVisitor()
+    visitor.visit(tree)
+    return visitor.violations