diff options
| -rw-r--r-- | snekbox/api/resources/eval.py | 4 | ||||
| -rw-r--r-- | snekbox/memfs.py | 12 | ||||
| -rw-r--r-- | snekbox/nsjail.py | 29 | ||||
| -rw-r--r-- | snekbox/process.py | 32 |
4 files changed, 53 insertions, 24 deletions
diff --git a/snekbox/api/resources/eval.py b/snekbox/api/resources/eval.py index 80c9ec4..1e1cc8b 100644 --- a/snekbox/api/resources/eval.py +++ b/snekbox/api/resources/eval.py @@ -77,7 +77,7 @@ class EvalResource: args = req.media.get("args", ("-c",)) try: - result, attachments = self.nsjail.python3(code, py_args=args) + result = self.nsjail.python3(code, py_args=args) except Exception: log.exception("An exception occurred while trying to process the request") raise falcon.HTTPInternalServerError @@ -85,5 +85,5 @@ class EvalResource: resp.media = { "stdout": result.stdout, "returncode": result.returncode, - "attachments": [atc.to_dict() for atc in attachments], + "attachments": [atc.to_dict() for atc in result.attachments], } diff --git a/snekbox/memfs.py b/snekbox/memfs.py index 520954b..2eb83d2 100644 --- a/snekbox/memfs.py +++ b/snekbox/memfs.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging import subprocess +from collections.abc import Generator from contextlib import contextmanager from functools import cache from pathlib import Path @@ -93,13 +94,12 @@ class MemoryTempDir: yield self.path.chmod(0o555) - def attachments(self) -> list[FileAttachment] | None: + def attachments(self) -> Generator[FileAttachment, None, None]: """Return a list of attachments in the tempdir.""" - # First look for any file named `output` (any extension) - output = next((f for f in self.home.glob("output*") if f.is_file()), None) - if output: - return [FileAttachment.from_path(output, MAX_FILE_SIZE)] - return None + # Look for any file starting with `output` + for file in self.path.glob("output*"): + if file.is_file(): + yield FileAttachment.from_path(file, MAX_FILE_SIZE) def cleanup(self) -> None: """Remove files in temp dir, releases name.""" diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py index d8b642a..0344c3c 100644 --- a/snekbox/nsjail.py +++ b/snekbox/nsjail.py @@ -3,19 +3,20 @@ 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.snekio import FileAttachment +from snekbox.process import EvalResult +from snekbox.snekio import AttachmentError log = logging.getLogger(__name__) @@ -133,7 +134,7 @@ class NsJail: def python3( self, code: str, *, nsjail_args: Iterable[str] = (), py_args: Iterable[str] = ("",) - ) -> tuple[CompletedProcess, list[FileAttachment]]: + ) -> EvalResult: """ Execute Python 3 code in an isolated environment and return the completed process. @@ -197,20 +198,12 @@ 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 @@ -218,8 +211,12 @@ class NsJail: returncode = -nsjail.returncode + 128 if nsjail.returncode < 0 else nsjail.returncode # Parse attachments - attachments = temp_dir.attachments() or [] - log.info(f"Found {len(attachments)} 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: @@ -230,4 +227,4 @@ class NsJail: log.info(f"nsjail return code: {returncode}") - return CompletedProcess(args, returncode, output, None), attachments + 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 [] |