Skip to content

Commit 8034dc0

Browse files
committed
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.
1 parent 83ed482 commit 8034dc0

8 files changed

Lines changed: 174 additions & 57 deletions

File tree

CLAUDE.md

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -125,16 +125,19 @@ docker build -t evaluatepython .
125125
# Cross-platform (CI uses linux/x86_64):
126126
docker build --platform=linux/x86_64 .
127127

128-
# Run the server locally (port 8080)
129-
docker run -it --rm -p 8080:8080 evaluatepython
128+
# Run the server locally (port 8080).
129+
# --privileged is required: the image enables shimmy's nsjail sandbox
130+
# (SANDBOX_ENABLED=true), and nsjail needs it to create namespaces.
131+
docker run -it --rm --privileged -p 8080:8080 evaluatepython
130132
```
131133

132134
## Tests
133135

134-
Two test files, run with `pytest`:
136+
Three test files, run with `pytest`:
135137

136-
- `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
138+
- `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
137139
- `evaluation_function/preview_test.py` — unit tests covering: valid Python, syntax errors, dangerous imports, dangerous builtins, dunder access
140+
- `evaluation_function/security_test.py` — unit tests for `check_code_safety` (the shared AST blocklist used by both `preview_function` and `evaluation_function`)
138141

139142
CI runs on Python 3.12 and uploads JUnit XML results (`.github/workflows/test-lint.yml`).
140143

@@ -146,14 +149,18 @@ CI runs on Python 3.12 and uploads JUnit XML results (`.github/workflows/test-li
146149
| `MPLBACKEND` | `Agg` | Set at subprocess runtime to suppress GUI |
147150
| `FUNCTION_COMMAND` | `python` | lf_toolkit runner |
148151
| `FUNCTION_ARGS` | `-m,evaluation_function.main` | lf_toolkit runner |
149-
| `FUNCTION_RPC_TRANSPORT` | `ipc` | lf_toolkit transport |
152+
| `FUNCTION_RPC_TRANSPORT` | `stdio` | shimmy↔worker transport (stdio so it survives the sandbox mount namespace) |
150153
| `LOG_LEVEL` | `debug` | Logging verbosity |
151154
| `AWS_*` / boto3 credentials | Runtime env | Required for S3 plot uploads |
155+
| `SANDBOX_ENABLED` | `true` | Wrap the worker in shimmy's nsjail sandbox (needs `--privileged` at run time) |
156+
| `SANDBOX_SECCOMP` | `true` | nsjail seccomp syscall filter |
157+
| `SANDBOX_RO_BINDS` | `/usr:/lib:/lib64:/bin:/sbin:/etc:/app` | Read-only bind mounts visible inside the jail |
158+
| `SANDBOX_TMPFS` | `/tmp` | Writable tmpfs inside the jail (student scripts, plot dirs, MPLCONFIGDIR) |
152159

153160
Dependencies managed via Poetry; `.venv` is created in-project (`poetry.toml`).
154161

155162
## Deployment
156163

157164
- Push to `main` triggers GitHub Actions (`.github/workflows/`) which builds and deploys to Lambda Feedback automatically
158165
- The function name is declared in `config.json` as `EvaluationFunctionName: "evaluatePython"` (lowerCamelCase)
159-
- The base Docker image is `ghcr.io/lambda-feedback/evaluation-function-base/python:test-sandbox-3.12`
166+
- 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

Lines changed: 35 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
FROM ghcr.io/lambda-feedback/evaluation-function-base/python:test-sandbox-3.12 AS builder
1+
FROM ghcr.io/lambda-feedback/evaluation-function-base/python:3.12 AS builder
22

33
RUN pip install poetry==1.8.3
44

@@ -12,7 +12,7 @@ COPY pyproject.toml poetry.lock ./
1212
RUN --mount=type=cache,target=$POETRY_CACHE_DIR \
1313
poetry install --without dev --no-root
1414

15-
FROM ghcr.io/lambda-feedback/evaluation-function-base/python:test-sandbox-3.12
15+
FROM ghcr.io/lambda-feedback/evaluation-function-base/python:3.12
1616

1717
ENV VIRTUAL_ENV=/app/.venv \
1818
PATH="/app/.venv/bin:$PATH"
@@ -31,7 +31,38 @@ ENV FUNCTION_COMMAND="python"
3131
# Args to start the evaluation function with
3232
ENV FUNCTION_ARGS="-m,evaluation_function.main"
3333

34-
# The transport to use for the RPC server
35-
ENV FUNCTION_RPC_TRANSPORT="ipc"
34+
# The transport to use for the RPC server.
35+
# stdio (not ipc): the worker runs inside an nsjail mount namespace with a
36+
# private tmpfs /tmp (see the sandbox settings below), so a host unix-socket
37+
# rendezvous at /tmp/eval.sock would be unreachable. shimmy's sandbox is
38+
# designed around stdio; lf_toolkit routes its logs to stderr, keeping stdout
39+
# clean for the RPC framing.
40+
ENV FUNCTION_RPC_TRANSPORT="stdio"
41+
42+
# --- Sandboxed execution of untrusted student code (shimmy + nsjail) ---
43+
# shimmy wraps the worker process -- and every `python` subprocess it spawns to
44+
# run a submission -- in an nsjail sandbox: unprivileged uid (nobody:nogroup),
45+
# a minimal bind-mounted view of the filesystem, and seccomp syscall filtering.
46+
#
47+
# Run-time requirements (cannot be expressed in the image):
48+
# * the container must run with --privileged (or --cap-add SYS_ADMIN) so
49+
# nsjail can create its namespaces -- see the shimmy README, "Sandboxed
50+
# Execution". Without it the worker will not boot.
51+
# * nsjail must exist at /usr/sbin/nsjail (provided by the base image's
52+
# shimmy stage).
53+
# If the worker fails to start, drop SANDBOX_SECCOMP first, then widen
54+
# SANDBOX_RO_BINDS (the list is linux/x86_64 + Debian-specific).
55+
#
56+
# Network is deliberately left enabled -- the function uploads matplotlib
57+
# plots to S3 via boto3. Untrusted network/filesystem use is blocked one layer
58+
# up, at the AST gate in evaluation_function/security.py (check_code_safety).
59+
#
60+
# No CPU/memory rlimits here: the RPC worker is long-lived and shared across
61+
# requests, so a cumulative RLIMIT_CPU / RLIMIT_AS would eventually kill it.
62+
# Per-execution wall-clock limits are enforced in evaluation.py (_TIMEOUT).
63+
ENV SANDBOX_ENABLED="true" \
64+
SANDBOX_SECCOMP="true" \
65+
SANDBOX_RO_BINDS="/usr:/lib:/lib64:/bin:/sbin:/etc:/app" \
66+
SANDBOX_TMPFS="/tmp"
3667

3768
ENV LOG_LEVEL="debug"

README.md

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,13 @@ Push to `main` triggers GitHub Actions which automatically builds and deploys to
1212
### Run the Docker Image
1313

1414
```bash
15-
docker run -it --rm -p 8080:8080 ghcr.io/lambda-feedback/evaluatepython:latest
15+
docker run -it --rm --privileged -p 8080:8080 ghcr.io/lambda-feedback/evaluatepython:latest
1616
```
1717

1818
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.
1919

20+
`--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`.
21+
2022
### Evaluation Modes
2123

2224
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.
8082
```
8183
evaluation_function/main.py # IPC server entry point
8284
evaluation_function/evaluation.py # core evaluation pipeline (all three modes)
83-
evaluation_function/preview.py # AST-based security validator
85+
evaluation_function/security.py # shared AST blocklist (check_code_safety)
86+
evaluation_function/preview.py # pre-submission validator (syntax + security)
8487
evaluation_function/dev.py # CLI wrapper for local testing
8588
evaluation_function/evaluation_test.py # integration tests
8689
evaluation_function/preview_test.py # preview/security tests
90+
evaluation_function/security_test.py # check_code_safety unit tests
8791
config.json # deployment configuration
8892
```
8993

@@ -144,9 +148,11 @@ docker build --platform=linux/x86_64 -t evaluatepython .
144148
### Running the Docker Image
145149

146150
```bash
147-
docker run -it --rm -p 8080:8080 evaluatepython
151+
docker run -it --rm --privileged -p 8080:8080 evaluatepython
148152
```
149153

154+
`--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`.
155+
150156
## Deployment to Lambda Feedback
151157

152158
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

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,12 @@
77
- **`evaluation_function`** — run and grade student code
88
- **`preview_function`** — static AST-based security check (called before execution)
99

10+
Both share the AST blocklist in `security.py` (`check_code_safety`). `preview_function`
11+
runs it pre-submission; `evaluation_function` also runs it as a hard gate — if the
12+
submission trips the blocklist it is **not executed** and an `error` item is returned.
13+
At the container level, shimmy additionally wraps execution in an nsjail sandbox
14+
(`SANDBOX_*` env in the `Dockerfile`).
15+
1016
---
1117

1218
## `evaluation_function`
@@ -173,7 +179,9 @@ The check runs in-process via `pycodestyle.Checker` with `max_line_length=200` (
173179

174180
## `preview_function`
175181

176-
Called before evaluation. Parses the student code as an AST and checks for security violations.
182+
Called before evaluation. Parses the student code as an AST and checks for security
183+
violations via `security.check_code_safety`. The same check gates `evaluation_function`,
184+
so the blocklist below applies to both entry points.
177185

178186
### Blocked constructs
179187

evaluation_function/evaluation.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
from lf_toolkit.evaluation import Result, Params
1212
from lf_toolkit.evaluation.image_upload import upload_image, ImageUploadError
1313

14+
from .security import check_code_safety
15+
1416
_TIMEOUT = 25
1517
_UPLOAD_FOLDER = "evaluatePython"
1618

@@ -303,6 +305,14 @@ def evaluation_function(response: Any, answer: Any, params: Params) -> Result:
303305
result.add_feedback("error", f"Unknown or missing mode: {mode!r}. Expected 'demo', 'io_test', or 'unit_test'.")
304306
return result
305307

308+
violations = check_code_safety(str(response))
309+
if violations:
310+
result.add_feedback(
311+
"error",
312+
"Unsafe code detected -- not executed:\n" + "\n".join(f"- {v}" for v in violations),
313+
)
314+
return result
315+
306316
if mode == "demo":
307317
result = _evaluate_demo(str(response), result)
308318
elif mode == "io_test":

evaluation_function/evaluation_test.py

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -330,4 +330,40 @@ def test_pep8_appended_to_io_test_result(self):
330330
params = {**_params(_test("5\n", "25\n")), "pep8_feedback": True}
331331
result = evaluation_function(_SQUARE_CODE, None, params).to_dict()
332332
self.assertTrue(result["is_correct"])
333-
self.assertIn("No style issues found", result["feedback"])
333+
self.assertIn("No style issues found", result["feedback"])
334+
335+
336+
class TestSecurityGate(unittest.TestCase):
337+
338+
def test_blocked_import_not_executed(self):
339+
result = evaluation_function(
340+
"import os\nprint(os.getcwd())", None, {"mode": "demo"}
341+
).to_dict()
342+
self.assertFalse(result["is_correct"])
343+
self.assertIn("Unsafe code detected", result["feedback"])
344+
self.assertIn("os", result["feedback"])
345+
346+
def test_blocked_builtin_not_executed(self):
347+
result = evaluation_function(
348+
"eval('1 + 1')", None, _params(_test("", ""))
349+
).to_dict()
350+
self.assertIn("Unsafe code detected", result["feedback"])
351+
self.assertNotIn("tests passed", result["feedback"])
352+
353+
def test_dunder_access_blocked_in_unit_test(self):
354+
result = evaluation_function(
355+
"().__class__.__bases__", None, _unit_params(_SQUARE_TESTS)
356+
).to_dict()
357+
self.assertIn("Unsafe code detected", result["feedback"])
358+
359+
def test_safe_stdlib_import_still_runs(self):
360+
result = evaluation_function(
361+
"import math\nprint(math.sqrt(16))", None, {"mode": "demo"}
362+
).to_dict()
363+
self.assertIn("4.0", result["feedback"])
364+
self.assertNotIn("Unsafe", result["feedback"])
365+
366+
def test_syntax_error_passes_gate_and_surfaces_downstream(self):
367+
result = evaluation_function("def f(:\n", None, {"mode": "demo"}).to_dict()
368+
self.assertNotIn("Unsafe code detected", result["feedback"])
369+
self.assertIn("SyntaxError", result["feedback"])

evaluation_function/preview.py

Lines changed: 6 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -2,55 +2,19 @@
22
from typing import Any
33
from lf_toolkit.preview import Result, Params, Preview
44

5-
_BLOCKED_MODULES = {
6-
"os", "sys", "subprocess", "socket", "urllib", "http",
7-
"requests", "shutil", "pathlib", "ftplib", "smtplib",
8-
"ctypes", "multiprocessing", "threading", "importlib",
9-
"pickle", "builtins",
10-
}
11-
12-
_BLOCKED_BUILTINS = {"exec", "eval", "compile", "open", "__import__"}
13-
14-
15-
class _SecurityVisitor(ast.NodeVisitor):
16-
def __init__(self):
17-
self.violations: list[str] = []
18-
19-
def visit_Import(self, node):
20-
for alias in node.names:
21-
root = alias.name.split(".")[0]
22-
if root in _BLOCKED_MODULES:
23-
self.violations.append(f"import of '{root}' is not allowed")
24-
self.generic_visit(node)
25-
26-
def visit_ImportFrom(self, node):
27-
if node.module:
28-
root = node.module.split(".")[0]
29-
if root in _BLOCKED_MODULES:
30-
self.violations.append(f"import of '{root}' is not allowed")
31-
self.generic_visit(node)
32-
33-
def visit_Call(self, node):
34-
if isinstance(node.func, ast.Name) and node.func.id in _BLOCKED_BUILTINS:
35-
self.violations.append(f"use of '{node.func.id}()' is not allowed")
36-
self.generic_visit(node)
37-
38-
def visit_Attribute(self, node):
39-
if node.attr.startswith("__") and node.attr.endswith("__"):
40-
self.violations.append(f"access to '{node.attr}' is not allowed")
41-
self.generic_visit(node)
5+
from .security import check_code_safety
426

437

448
def preview_function(response: Any, params: Params) -> Result:
9+
code = str(response)
4510
try:
46-
tree = ast.parse(str(response))
11+
ast.parse(code)
4712
except SyntaxError as e:
4813
return Result(preview=Preview(feedback=f"SyntaxError: {e.msg} (line {e.lineno})"))
4914

50-
visitor = _SecurityVisitor()
51-
visitor.visit(tree)
52-
if visitor.violations:
53-
lines = "\n".join(f"- {v}" for v in visitor.violations)
15+
violations = check_code_safety(code)
16+
if violations:
17+
lines = "\n".join(f"- {v}" for v in violations)
5418
return Result(preview=Preview(feedback=f"Unsafe code detected:\n{lines}"))
5519

5620
return Result(preview=Preview(feedback="Valid Python syntax."))

evaluation_function/security.py

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
import ast
2+
3+
_BLOCKED_MODULES = {
4+
"os", "sys", "subprocess", "socket", "urllib", "http",
5+
"requests", "shutil", "pathlib", "ftplib", "smtplib",
6+
"ctypes", "multiprocessing", "threading", "importlib",
7+
"pickle", "builtins",
8+
}
9+
10+
_BLOCKED_BUILTINS = {"exec", "eval", "compile", "open", "__import__"}
11+
12+
13+
class _SecurityVisitor(ast.NodeVisitor):
14+
def __init__(self):
15+
self.violations: list[str] = []
16+
17+
def visit_Import(self, node):
18+
for alias in node.names:
19+
root = alias.name.split(".")[0]
20+
if root in _BLOCKED_MODULES:
21+
self.violations.append(f"import of '{root}' is not allowed")
22+
self.generic_visit(node)
23+
24+
def visit_ImportFrom(self, node):
25+
if node.module:
26+
root = node.module.split(".")[0]
27+
if root in _BLOCKED_MODULES:
28+
self.violations.append(f"import of '{root}' is not allowed")
29+
self.generic_visit(node)
30+
31+
def visit_Call(self, node):
32+
if isinstance(node.func, ast.Name) and node.func.id in _BLOCKED_BUILTINS:
33+
self.violations.append(f"use of '{node.func.id}()' is not allowed")
34+
self.generic_visit(node)
35+
36+
def visit_Attribute(self, node):
37+
if node.attr.startswith("__") and node.attr.endswith("__"):
38+
self.violations.append(f"access to '{node.attr}' is not allowed")
39+
self.generic_visit(node)
40+
41+
42+
def check_code_safety(code: str) -> list[str]:
43+
"""Return the list of security violations in ``code`` (empty means safe).
44+
45+
Shared by ``preview_function`` (pre-submission check) and
46+
``evaluation_function`` (pre-execution gate). A syntax error is not a
47+
safety violation -- it is surfaced elsewhere -- so it yields an empty list.
48+
"""
49+
try:
50+
tree = ast.parse(code)
51+
except SyntaxError:
52+
return []
53+
visitor = _SecurityVisitor()
54+
visitor.visit(tree)
55+
return visitor.violations

0 commit comments

Comments
 (0)