aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
Diffstat
-rw-r--r--config/snekbox.cfg2
-rw-r--r--docker-compose.yml2
-rw-r--r--snekbox/__main__.py2
-rw-r--r--snekbox/api/resources/eval.py6
-rw-r--r--snekbox/memfs.py140
-rw-r--r--snekbox/nsjail.py52
-rw-r--r--snekbox/process.py32
-rw-r--r--snekbox/snekio.py70
-rw-r--r--tests/test_main.py4
9 files changed, 291 insertions, 19 deletions
diff --git a/config/snekbox.cfg b/config/snekbox.cfg
index 87c216e..aaa5d89 100644
--- a/config/snekbox.cfg
+++ b/config/snekbox.cfg
@@ -104,6 +104,8 @@ mount {
rw: false
}
+rlimit_fsize: 134217728
+
cgroup_mem_max: 52428800
cgroup_mem_swap_max: 0
cgroup_mem_mount: "/sys/fs/cgroup/memory"
diff --git a/docker-compose.yml b/docker-compose.yml
index 9d3ae71..0613abc 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -8,7 +8,7 @@ services:
image: ghcr.io/python-discord/snekbox${IMAGE_SUFFIX:--venv:dev}
pull_policy: never
ports:
- - 8060:8060
+ - "8060:8060"
init: true
ipc: none
tty: true
diff --git a/snekbox/__main__.py b/snekbox/__main__.py
index 2382c4c..6cbd1ea 100644
--- a/snekbox/__main__.py
+++ b/snekbox/__main__.py
@@ -16,7 +16,7 @@ def parse_args() -> argparse.Namespace:
"nsjail_args", nargs="?", default=[], help="override configured NsJail options"
)
parser.add_argument(
- "py_args", nargs="?", default=["-c"], help="arguments to pass to the Python process"
+ "py_args", nargs="?", default=[], help="arguments to pass to the Python process"
)
# nsjail_args and py_args are just dummies for documentation purposes.
diff --git a/snekbox/api/resources/eval.py b/snekbox/api/resources/eval.py
index 1df6c1b..1e1cc8b 100644
--- a/snekbox/api/resources/eval.py
+++ b/snekbox/api/resources/eval.py
@@ -82,4 +82,8 @@ class EvalResource:
log.exception("An exception occurred while trying to process the request")
raise falcon.HTTPInternalServerError
- resp.media = {"stdout": result.stdout, "returncode": result.returncode}
+ resp.media = {
+ "stdout": result.stdout,
+ "returncode": result.returncode,
+ "attachments": [atc.to_dict() for atc in result.attachments],
+ }
diff --git a/snekbox/memfs.py b/snekbox/memfs.py
new file mode 100644
index 0000000..b7295a6
--- /dev/null
+++ b/snekbox/memfs.py
@@ -0,0 +1,140 @@
+"""Memory filesystem for snekbox."""
+from __future__ import annotations
+
+import logging
+import subprocess
+from collections.abc import Generator
+from contextlib import contextmanager
+from pathlib import Path
+from shutil import rmtree
+from threading import BoundedSemaphore
+from types import TracebackType
+from typing import Type
+from uuid import uuid4
+
+from snekbox.snekio import FileAttachment
+
+log = logging.getLogger(__name__)
+
+
+def mount_tmpfs(name: str) -> Path:
+ """Create and mount a tmpfs directory."""
+ namespace = Path("/snekbox/memfs")
+ tmp = namespace / name
+ if not tmp.exists() or not tmp.is_dir():
+ # Create the directory
+ tmp.mkdir(parents=True, exist_ok=True)
+ tmp.chmod(0o777)
+ # Mount the tmpfs
+ subprocess.check_call(
+ [
+ "mount",
+ "-t",
+ "tmpfs",
+ "-o",
+ f"size={MemFSOptions.MEMFS_SIZE}",
+ "tmpfs",
+ str(tmp),
+ ]
+ )
+ # Execute only access for other users
+ tmp.chmod(0o711)
+ namespace.chmod(0o711)
+ return tmp
+
+
+def unmount_tmpfs(name: str) -> None:
+ """Unmount and remove a tmpfs directory."""
+ tmp = Path("/snekbox/memfs", name)
+ if tmp.exists() and tmp.is_dir():
+ subprocess.check_call(["umount", str(tmp)])
+ rmtree(tmp, ignore_errors=True)
+
+
+class MemFSOptions:
+ """Options for memory file system."""
+
+ # Size of the memory filesystem (per instance)
+ MEMFS_SIZE = "32M"
+ # Maximum number of files attachments will be scanned for
+ MAX_FILES = 2
+ # Maximum size of a file attachment (32 MB)
+ MAX_FILE_SIZE = 32 * 1024 * 1024
+
+
+class MemoryTempDir:
+ """A temporary directory using tmpfs."""
+
+ assignment_lock = BoundedSemaphore(1)
+ assigned_names: set[str] = set() # Pool of tempdir names in use
+
+ def __init__(self) -> None:
+ self.path: Path | None = None
+
+ @property
+ def name(self) -> str | None:
+ """Name of the temp dir."""
+ return self.path.name if self.path else None
+
+ @property
+ def home(self) -> Path | None:
+ """Path to home directory."""
+ return Path(self.path, "home") if self.path else None
+
+ def __enter__(self) -> MemoryTempDir:
+ # Generates a uuid tempdir
+ with self.assignment_lock:
+ for _ in range(10):
+ name = str(uuid4())
+ if name not in self.assigned_names:
+ self.path = mount_tmpfs(name)
+ self.path.chmod(0o555)
+ # Create a home folder
+ home = self.path / "home"
+ home.mkdir()
+ home.chmod(0o777)
+ self.assigned_names.add(name)
+ return self
+ else:
+ raise RuntimeError("Failed to generate a unique tempdir name in 10 attempts")
+
+ def __exit__(
+ self,
+ exc_type: Type[BaseException] | None,
+ exc_val: BaseException | None,
+ exc_tb: TracebackType | None,
+ ) -> None:
+ self.cleanup()
+
+ @contextmanager
+ def allow_write(self) -> None:
+ """Temporarily allow writes to the root tempdir."""
+ self.path.chmod(0o777)
+ yield
+ self.path.chmod(0o555)
+
+ def attachments(self) -> Generator[FileAttachment, None, None]:
+ """Return a list of attachments in the tempdir."""
+ # Look for any file starting with `output`
+ for file in self.home.glob("output*"):
+ if file.is_file():
+ yield FileAttachment.from_path(file, MemFSOptions.MAX_FILE_SIZE)
+
+ def cleanup(self) -> None:
+ """Remove files in temp dir, releases name."""
+ if self.path is None:
+ return
+ # Remove the path folder
+ unmount_tmpfs(self.name)
+
+ if not self.path.exists():
+ with self.assignment_lock:
+ self.assigned_names.remove(self.name)
+ else:
+ # Don't remove name from pool if failed to delete folder
+ logging.warning(f"Failed to remove {self.path} in cleanup")
+
+ self.path = None
+
+ def __repr__(self):
+ return f"<MemoryTempDir {self.name or '(Uninitialized)'}>"
diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index 63afdef..0344c3c 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -3,17 +3,21 @@ import re
import subprocess
import sys
import textwrap
-from subprocess import CompletedProcess
from tempfile import NamedTemporaryFile
from typing import Iterable
from google.protobuf import text_format
+# noinspection PyProtectedMember
from snekbox import DEBUG, utils
from snekbox.config_pb2 import NsJailConfig
+from snekbox.memfs import MemoryTempDir
__all__ = ("NsJail",)
+from snekbox.process import EvalResult
+from snekbox.snekio import AttachmentError
+
log = logging.getLogger(__name__)
# [level][timestamp][PID]? function_signature:line_no? message
@@ -129,8 +133,8 @@ class NsJail:
return "".join(output)
def python3(
- self, code: str, *, nsjail_args: Iterable[str] = (), py_args: Iterable[str] = ("-c",)
- ) -> CompletedProcess:
+ self, code: str, *, nsjail_args: Iterable[str] = (), py_args: Iterable[str] = ("",)
+ ) -> EvalResult:
"""
Execute Python 3 code in an isolated environment and return the completed process.
@@ -152,7 +156,24 @@ class NsJail:
*nsjail_args,
)
- with NamedTemporaryFile() as nsj_log:
+ with NamedTemporaryFile() as nsj_log, MemoryTempDir() as temp_dir:
+ # Write the code to a python file in the temp directory.
+ with temp_dir.allow_write():
+ code_path = temp_dir.path / "main.py"
+ code_path.write_text(code)
+ log.info(f"Created code file at [{code_path!r}].")
+
+ # Add the temp dir to be mounted as cwd
+ nsjail_args = (
+ "--bindmount", # Mount temp dir in R/W mode
+ f"{temp_dir.home}:home",
+ "--cwd", # Set cwd to temp dir
+ "home",
+ "--env", # Set $HOME to temp dir
+ "HOME=home",
+ *nsjail_args,
+ )
+
args = (
self.nsjail_path,
"--config",
@@ -163,8 +184,8 @@ class NsJail:
"--",
self.config.exec_bin.path,
*self.config.exec_bin.arg,
- *py_args,
- code,
+ *[arg for arg in py_args if arg != "-c"],
+ code_path,
)
msg = "Executing code..."
@@ -177,23 +198,26 @@ class NsJail:
args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
)
except ValueError:
- return CompletedProcess(args, None, "ValueError: embedded null byte", None)
+ return EvalResult(args, None, "ValueError: embedded null byte")
try:
output = self._consume_stdout(nsjail)
except UnicodeDecodeError:
- return CompletedProcess(
- args,
- None,
- "UnicodeDecodeError: invalid Unicode in output pipe",
- None,
- )
+ return EvalResult(args, None, "UnicodeDecodeError: invalid Unicode in output pipe")
# When you send signal `N` to a subprocess to terminate it using Popen, it
# will return `-N` as its exit code. As we normally get `N + 128` back, we
# convert negative exit codes to the `N + 128` form.
returncode = -nsjail.returncode + 128 if nsjail.returncode < 0 else nsjail.returncode
+ # Parse attachments
+ try:
+ attachments = list(temp_dir.attachments())
+ log.info(f"Found {len(attachments)} attachments.")
+ except AttachmentError as err:
+ log.error(f"Failed to parse attachments: {err}")
+ return EvalResult(args, returncode, f"AttachmentError: {err}")
+
log_lines = nsj_log.read().decode("utf-8").splitlines()
if not log_lines and returncode == 255:
# NsJail probably failed to parse arguments so log output will still be in stdout
@@ -203,4 +227,4 @@ class NsJail:
log.info(f"nsjail return code: {returncode}")
- return CompletedProcess(args, returncode, output, None)
+ return EvalResult(args, returncode, output, attachments=attachments)
diff --git a/snekbox/process.py b/snekbox/process.py
new file mode 100644
index 0000000..d6d25c0
--- /dev/null
+++ b/snekbox/process.py
@@ -0,0 +1,32 @@
+"""Utilities for process management."""
+from collections.abc import Sequence
+from os import PathLike
+from subprocess import CompletedProcess
+from typing import TypeVar
+
+from snekbox.snekio import FileAttachment
+
+_T = TypeVar("_T")
+ArgType = (
+ str
+ | bytes
+ | PathLike[str]
+ | PathLike[bytes]
+ | Sequence[str | bytes | PathLike[str] | PathLike[bytes]]
+)
+
+
+class EvalResult(CompletedProcess[_T]):
+ """An evaluation job that has finished running."""
+
+ def __init__(
+ self,
+ args: ArgType,
+ returncode: int | None,
+ stdout: _T | None = None,
+ stderr: _T | None = None,
+ attachments: list[FileAttachment] | None = None,
+ ) -> None:
+ """Create an evaluation result."""
+ super().__init__(args, returncode, stdout, stderr)
+ self.attachments: list[FileAttachment] = attachments or []
diff --git a/snekbox/snekio.py b/snekbox/snekio.py
new file mode 100644
index 0000000..12645cb
--- /dev/null
+++ b/snekbox/snekio.py
@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+import mimetypes
+import zlib
+from base64 import b64encode
+from dataclasses import dataclass
+from pathlib import Path
+
+SUPPORTED_MIME_TYPES = {
+ "image/png",
+ "image/jpeg",
+}
+
+
+def sizeof_fmt(num: int, suffix: str = "B") -> str:
+ """Return a human-readable file size."""
+ for unit in ("", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi"):
+ if abs(num) < 1024:
+ return f"{num:3.1f}{unit}{suffix}"
+ num /= 1024
+ return f"{num:.1f}Yi{suffix}"
+
+
+class AttachmentError(ValueError):
+ """Raised when an attachment is invalid."""
+
+
+@dataclass
+class FileAttachment:
+ """A file attachment."""
+
+ name: str
+ content: bytes
+
+ @classmethod
+ def from_path(cls, file: Path, max_size: int) -> FileAttachment:
+ """Create an attachment from a path."""
+ size = file.stat().st_size
+ if size > max_size:
+ raise AttachmentError(
+ f"File {file.name} too large: {sizeof_fmt(size)} "
+ f"exceeds the limit of {sizeof_fmt(max_size)}"
+ )
+
+ with file.open("rb") as f:
+ content = f.read(max_size + 1)
+ size = len(content)
+ if len(content) > max_size:
+ raise AttachmentError(
+ f"File {file.name} too large: {sizeof_fmt(len(content))} "
+ f"exceeds the limit of {sizeof_fmt(max_size)}"
+ )
+ return cls(file.name, content)
+
+ @property
+ def mime(self) -> str:
+ """MIME type of the attachment."""
+ return mimetypes.guess_type(self.name)[0]
+
+ def is_supported(self) -> bool:
+ """Return whether the attachment is supported."""
+ if self.mime.startswith("text/"):
+ return True
+ return self.mime in SUPPORTED_MIME_TYPES
+
+ def to_dict(self) -> dict[str, str]:
+ """Convert the attachment to a dict."""
+ cmp = zlib.compress(self.content)
+ content = b64encode(cmp).decode("ascii")
+ return {"name": self.name, "mime": self.mime, "content": content}
diff --git a/tests/test_main.py b/tests/test_main.py
index 77b3130..1e6cbc5 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -12,10 +12,10 @@ import snekbox.__main__ as snekbox_main
class ArgParseTests(unittest.TestCase):
def test_parse_args(self):
subtests = (
- (["", "code"], Namespace(code="code", nsjail_args=[], py_args=["-c"])),
+ (["", "code"], Namespace(code="code", nsjail_args=[], py_args=[])),
(
["", "code", "--time_limit", "0"],
- Namespace(code="code", nsjail_args=["--time_limit", "0"], py_args=["-c"]),
+ Namespace(code="code", nsjail_args=["--time_limit", "0"], py_args=[]),
),
(
["", "code", "---", "-m", "timeit"],
generated by cgit v1.2.3 (git 2.39.1) at 2025-10-22 18:20:02 +0000