Merge #145 - make the output size limit customisable

9 files changed, 82 insertions, 46 deletions

diff --git a/README.md b/README.md
index 2e2ee9c..7540e21 100644
--- a/README.md
+++ b/README.md
@@ -28,7 +28,7 @@ Snekbox -->>- Client: JSON response
 
 The code is executed in a Python process that is launched through [NsJail], which is responsible for sandboxing the Python process.
 
-The output returned by snekbox is truncated at around 1 MB.
+The output returned by snekbox is truncated at around 1 MB by default, but this can be [configured](#gunicorn).
 
 ## HTTP REST API
 
@@ -66,7 +66,9 @@ NsJail is configured through [`snekbox.cfg`]. It contains the exact values for t
 
 ### Gunicorn
 
-[Gunicorn settings] can be found in [`gunicorn.conf.py`]. In the default configuration, the worker count and the bind address are likely the only things of any interest. Since it uses the default synchronous workers, the [worker count] effectively determines how many concurrent code evaluations can be performed.
+[Gunicorn settings] can be found in [`gunicorn.conf.py`]. In the default configuration, the worker count, the bind address, and the WSGI app URI are likely the only things of any interest. Since it uses the default synchronous workers, the [worker count] effectively determines how many concurrent code evaluations can be performed.
+
+`wsgi_app` can be given arguments which are forwarded to the [`NsJail`] object. For example, `wsgi_app = "snekbox:SnekAPI(max_output_size=2_000_000, read_chunk_size=20_000)"`.
 
 ### Environment Variables
 
@@ -74,13 +76,9 @@ All environment variables have defaults and are therefore not required to be set
 
 Name | Description
 ---- | -----------
-`DEBUG` | Enable debug logging if set to a non-empty value.
-`NSJAIL_CFG` | Path to the NsJail configuration file.
-`NSJAIL_PATH` | Path to the NsJail binary.
+`SNEKBOX_DEBUG` | Enable debug logging if set to a non-empty value.
 `SNEKBOX_SENTRY_DSN` | [Data Source Name] for Sentry. Sentry is disabled if left unset.
 
-Note: relative paths are relative to the root of the repository.
-
 ## Third-party Packages
 
 By default, the Python interpreter has no access to any packages besides the
@@ -125,3 +123,4 @@ See [CONTRIBUTING.md](.github/CONTRIBUTING.md).
 [sentry release]: https://docs.sentry.io/platforms/python/configuration/releases/
 [data source name]: https://docs.sentry.io/product/sentry-basics/dsn-explainer/
 [GitHub Container Registry]: https://github.com/orgs/python-discord/packages/container/package/snekbox
+[`NsJail`]: snekbox/nsjail.py
diff --git a/docker-compose.yml b/docker-compose.yml
index 07548cd..9d3ae71 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -13,7 +13,7 @@ services:
     ipc: none
     tty: true
     environment:
-      DEBUG: 1
+      SNEKBOX_DEBUG: 1
       PYTHONDONTWRITEBYTECODE: 1
     build:
       context: .
diff --git a/pyproject.toml b/pyproject.toml
index 7a13370..e0a3d26 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -65,7 +65,7 @@ relative_files = true
 [tool.black]
 line-length = 100
 target-version = ["py310"]
-force-exclude = ["snekbox/config_pb2.py"]
+force-exclude = "snekbox/config_pb2.py"
 
 [tool.isort]
 line_length = 100
diff --git a/snekbox/__init__.py b/snekbox/__init__.py
index 657c032..bed3692 100644
--- a/snekbox/__init__.py
+++ b/snekbox/__init__.py
@@ -1,7 +1,7 @@
 import os
 from importlib import metadata
 
-DEBUG = os.environ.get("DEBUG", False)
+DEBUG = os.environ.get("SNEKBOX_DEBUG", False)
 
 try:
     __version__ = metadata.version("snekbox")
diff --git a/snekbox/api/resources/eval.py b/snekbox/api/resources/eval.py
index 7c23a52..1df6c1b 100644
--- a/snekbox/api/resources/eval.py
+++ b/snekbox/api/resources/eval.py
@@ -29,8 +29,8 @@ class EvalResource:
         "required": ["input"],
     }
 
-    def __init__(self):
-        self.nsjail = NsJail()
+    def __init__(self, nsjail: NsJail):
+        self.nsjail = nsjail
 
     @validate(REQ_SCHEMA)
     def on_post(self, req: falcon.Request, resp: falcon.Response) -> None:
diff --git a/snekbox/api/snekapi.py b/snekbox/api/snekapi.py
index a1804e6..5a8d390 100644
--- a/snekbox/api/snekapi.py
+++ b/snekbox/api/snekapi.py
@@ -1,5 +1,7 @@
 import falcon
 
+from snekbox.nsjail import NsJail
+
 from .resources import EvalResource
 
 
@@ -7,6 +9,8 @@ class SnekAPI(falcon.App):
     """
     The main entry point to the snekbox JSON API.
 
+    Forward arguments to a new `NsJail` object.
+
     Routes:
 
     - /eval
@@ -21,6 +25,7 @@ class SnekAPI(falcon.App):
     """
 
     def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+        super().__init__()
 
-        self.add_route("/eval", EvalResource())
+        nsjail = NsJail(*args, **kwargs)
+        self.add_route("/eval", EvalResource(nsjail))
diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index a6b202e..63afdef 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -1,5 +1,4 @@
 import logging
-import os
 import re
 import subprocess
 import sys
@@ -21,14 +20,6 @@ log = logging.getLogger(__name__)
 LOG_PATTERN = re.compile(
     r"\[(?P<level>(I)|[DWEF])\]\[.+?\](?(2)|(?P<func>\[\d+\] .+?:\d+ )) ?(?P<msg>.+)"
 )
-LOG_BLACKLIST = ("Process will be ",)
-
-NSJAIL_PATH = os.getenv("NSJAIL_PATH", "/usr/sbin/nsjail")
-NSJAIL_CFG = os.getenv("NSJAIL_CFG", "./config/snekbox.cfg")
-
-# Limit of stdout bytes we consume before terminating nsjail
-OUTPUT_MAX = 1_000_000  # 1 MB
-READ_CHUNK_SIZE = 10_000  # chars
 
 
 class NsJail:
@@ -38,33 +29,43 @@ class NsJail:
     See config/snekbox.cfg for the default NsJail configuration.
     """
 
-    def __init__(self, nsjail_binary: str = NSJAIL_PATH):
-        self.nsjail_binary = nsjail_binary
-        self.config = self._read_config()
+    def __init__(
+        self,
+        nsjail_path: str = "/usr/sbin/nsjail",
+        config_path: str = "./config/snekbox.cfg",
+        max_output_size: int = 1_000_000,
+        read_chunk_size: int = 10_000,
+    ):
+        self.nsjail_path = nsjail_path
+        self.config_path = config_path
+        self.max_output_size = max_output_size
+        self.read_chunk_size = read_chunk_size
+
+        self.config = self._read_config(config_path)
         self.cgroup_version = utils.cgroup.init(self.config)
         self.ignore_swap_limits = utils.swap.should_ignore_limit(self.config, self.cgroup_version)
 
         log.info(f"Assuming cgroup version {self.cgroup_version}.")
 
     @staticmethod
-    def _read_config() -> NsJailConfig:
-        """Read the NsJail config at `NSJAIL_CFG` and return a protobuf Message object."""
+    def _read_config(config_path: str) -> NsJailConfig:
+        """Read the NsJail config at `config_path` and return a protobuf Message object."""
         config = NsJailConfig()
 
         try:
-            with open(NSJAIL_CFG, encoding="utf-8") as f:
+            with open(config_path, encoding="utf-8") as f:
                 config_text = f.read()
         except FileNotFoundError:
-            log.fatal(f"The NsJail config at {NSJAIL_CFG!r} could not be found.")
+            log.fatal(f"The NsJail config at {config_path!r} could not be found.")
             sys.exit(1)
         except OSError as e:
-            log.fatal(f"The NsJail config at {NSJAIL_CFG!r} could not be read.", exc_info=e)
+            log.fatal(f"The NsJail config at {config_path!r} could not be read.", exc_info=e)
             sys.exit(1)
 
         try:
             text_format.Parse(config_text, config)
         except text_format.ParseError as e:
-            log.fatal(f"The NsJail config at {NSJAIL_CFG!r} could not be parsed.", exc_info=e)
+            log.fatal(f"The NsJail config at {config_path!r} could not be parsed.", exc_info=e)
             sys.exit(1)
 
         return config
@@ -79,10 +80,6 @@ class NsJail:
                 continue
 
             msg = match["msg"]
-            if not DEBUG and any(msg.startswith(s) for s in LOG_BLACKLIST):
-                # Skip blacklisted messages if not debugging.
-                continue
-
             if DEBUG and match["func"]:
                 # Prepend PID, function signature, and line number if debugging.
                 msg = f"{match['func']}{msg}"
@@ -99,8 +96,7 @@ class NsJail:
                 # Treat fatal as error.
                 log.error(msg)
 
-    @staticmethod
-    def _consume_stdout(nsjail: subprocess.Popen) -> str:
+    def _consume_stdout(self, nsjail: subprocess.Popen) -> str:
         """
         Consume STDOUT, stopping when the output limit is reached or NsJail has exited.
 
@@ -119,11 +115,11 @@ class NsJail:
         with nsjail:
             # We'll consume STDOUT as long as the NsJail subprocess is running.
             while nsjail.poll() is None:
-                chars = nsjail.stdout.read(READ_CHUNK_SIZE)
+                chars = nsjail.stdout.read(self.read_chunk_size)
                 output_size += sys.getsizeof(chars)
                 output.append(chars)
 
-                if output_size > OUTPUT_MAX:
+                if output_size > self.max_output_size:
                     # Terminate the NsJail subprocess with SIGTERM.
                     # This in turn reaps and kills children with SIGKILL.
                     log.info("Output exceeded the output limit, sending SIGTERM to NsJail.")
@@ -158,9 +154,9 @@ class NsJail:
 
         with NamedTemporaryFile() as nsj_log:
             args = (
-                self.nsjail_binary,
+                self.nsjail_path,
                 "--config",
-                NSJAIL_CFG,
+                self.config_path,
                 "--log",
                 nsj_log.name,
                 *nsjail_args,
diff --git a/tests/api/__init__.py b/tests/api/__init__.py
index 3f7d250..0e6e422 100644
--- a/tests/api/__init__.py
+++ b/tests/api/__init__.py
@@ -11,7 +11,7 @@ class SnekAPITestCase(testing.TestCase):
     def setUp(self):
         super().setUp()
 
-        self.patcher = mock.patch("snekbox.api.resources.eval.NsJail", autospec=True)
+        self.patcher = mock.patch("snekbox.api.snekapi.NsJail", autospec=True)
         self.mock_nsjail = self.patcher.start()
         self.mock_nsjail.return_value.python3.return_value = CompletedProcess(
             args=[], returncode=0, stdout="output", stderr="error"
diff --git a/tests/test_nsjail.py b/tests/test_nsjail.py
index a3632e6..492c2f9 100644
--- a/tests/test_nsjail.py
+++ b/tests/test_nsjail.py
@@ -1,11 +1,13 @@
 import io
 import logging
+import shutil
 import sys
+import tempfile
 import unittest
 import unittest.mock
 from textwrap import dedent
 
-from snekbox.nsjail import OUTPUT_MAX, READ_CHUNK_SIZE, NsJail
+from snekbox.nsjail import NsJail
 
 
 class NsJailTests(unittest.TestCase):
@@ -255,8 +257,8 @@ class NsJailTests(unittest.TestCase):
         self.assertEqual(result.returncode, 143)
 
     def test_large_output_is_truncated(self):
-        chunk = "a" * READ_CHUNK_SIZE
-        expected_chunks = OUTPUT_MAX // sys.getsizeof(chunk) + 1
+        chunk = "a" * self.nsjail.read_chunk_size
+        expected_chunks = self.nsjail.max_output_size // sys.getsizeof(chunk) + 1
 
         nsjail_subprocess = unittest.mock.MagicMock()
 
@@ -280,3 +282,37 @@ class NsJailTests(unittest.TestCase):
 
         self.assertEqual(result.returncode, 0)
         self.assertEqual(result.args[-3:-1], args)
+
+
+class NsJailArgsTests(unittest.TestCase):
+    def setUp(self):
+        self._temp_dir = tempfile.TemporaryDirectory()
+        self.addClassCleanup(self._temp_dir.cleanup)
+
+        self.nsjail_path = shutil.copy2("/usr/sbin/nsjail", self._temp_dir.name)
+        self.config_path = shutil.copy2("./config/snekbox.cfg", self._temp_dir.name)
+        self.max_output_size = 1_234_567
+        self.read_chunk_size = 12_345
+
+        self.nsjail = NsJail(
+            self.nsjail_path, self.config_path, self.max_output_size, self.read_chunk_size
+        )
+
+        logging.getLogger("snekbox.nsjail").setLevel(logging.WARNING)
+
+    def test_nsjail_path(self):
+        result = self.nsjail.python3("")
+
+        self.assertEqual(result.args[0], self.nsjail_path)
+
+    def test_config_path(self):
+        result = self.nsjail.python3("")
+
+        i = result.args.index("--config") + 1
+        self.assertEqual(result.args[i], self.config_path)
+
+    def test_init_args(self):
+        self.assertEqual(self.nsjail.nsjail_path, self.nsjail_path)
+        self.assertEqual(self.nsjail.config_path, self.config_path)
+        self.assertEqual(self.nsjail.max_output_size, self.max_output_size)
+        self.assertEqual(self.nsjail.read_chunk_size, self.read_chunk_size)