Merge pull request #81 - Limit STDOUT to prevent OOM

4 files changed, 76 insertions, 6 deletions

diff --git a/README.md b/README.md
index 6765260..187d029 100644
--- a/README.md
+++ b/README.md
@@ -24,6 +24,7 @@ result <- |             |<----------|           | <----------+
 
 The code is executed in a Python process that is launched through [NsJail], which is responsible for sandboxing the Python process. See [`snekbox.cfg`] for the NsJail configuration.
 
+The output returned by snekbox is truncated at around 1 MB.
 
 ## HTTP REST API
 
diff --git a/config/snekbox.cfg b/config/snekbox.cfg
index 5b47459..27caf27 100644
--- a/config/snekbox.cfg
+++ b/config/snekbox.cfg
@@ -18,7 +18,6 @@ envar: "NUMEXPR_NUM_THREADS=1"
 keep_caps: false
 
 rlimit_as: 700
-rlimit_fsize: 10
 
 clone_newnet: true
 clone_newuser: true
diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index cafde6d..cc1f682 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -27,6 +27,10 @@ NSJAIL_PATH = os.getenv("NSJAIL_PATH", "/usr/sbin/nsjail")
 NSJAIL_CFG = os.getenv("NSJAIL_CFG", "./config/snekbox.cfg")
 MEM_MAX = 52428800
 
+# Limit of stdout bytes we consume before terminating nsjail
+OUTPUT_MAX = 1_000_000  # 1 MB
+READ_CHUNK_SIZE = 10_000  # chars
+
 
 class NsJail:
     """
@@ -101,6 +105,39 @@ class NsJail:
                 # Treat fatal as error.
                 log.error(msg)
 
+    @staticmethod
+    def _consume_stdout(nsjail: subprocess.Popen) -> str:
+        """
+        Consume STDOUT, stopping when the output limit is reached or NsJail has exited.
+
+        The aim of this function is to limit the size of the output received from
+        NsJail to prevent container from claiming too much memory. If the output
+        received from STDOUT goes over the OUTPUT_MAX limit, the NsJail subprocess
+        is asked to terminate with a SIGKILL.
+
+        Once the subprocess has exited, either naturally or because it was terminated,
+        we return the output as a single string.
+        """
+        output_size = 0
+        output = []
+
+        # We'll consume STDOUT as long as the NsJail subprocess is running.
+        while nsjail.poll() is None:
+            chars = nsjail.stdout.read(READ_CHUNK_SIZE)
+            output_size += sys.getsizeof(chars)
+            output.append(chars)
+
+            if output_size > OUTPUT_MAX:
+                # Terminate the NsJail subprocess with SIGKILL.
+                log.info("Output exceeded the output limit, sending SIGKILL to NsJail.")
+                nsjail.kill()
+                break
+
+        # Ensure that we wait for the NsJail subprocess to terminate.
+        nsjail.wait()
+
+        return "".join(output)
+
     def python3(self, code: str) -> CompletedProcess:
         """Execute Python 3 code in an isolated environment and return the completed process."""
         with NamedTemporaryFile() as nsj_log:
@@ -124,7 +161,7 @@ class NsJail:
             log.info(msg)
 
             try:
-                result = subprocess.run(
+                nsjail = subprocess.Popen(
                     args,
                     stdout=subprocess.PIPE,
                     stderr=subprocess.STDOUT,
@@ -133,11 +170,19 @@ class NsJail:
             except ValueError:
                 return CompletedProcess(args, None, "ValueError: embedded null byte", None)
 
+            output = self._consume_stdout(nsjail)
+
+            # 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
+
             log_lines = nsj_log.read().decode("utf-8").splitlines()
-            if not log_lines and result.returncode == 255:
+            if not log_lines and returncode == 255:
                 # NsJail probably failed to parse arguments so log output will still be in stdout
-                log_lines = result.stdout.splitlines()
+                log_lines = output.splitlines()
 
             self._parse_log(log_lines)
 
-        return result
+        log.info(f"nsjail return code: {returncode}")
+        return CompletedProcess(args, returncode, output, None)
diff --git a/tests/test_nsjail.py b/tests/test_nsjail.py
index 0b755b2..40c27f9 100644
--- a/tests/test_nsjail.py
+++ b/tests/test_nsjail.py
@@ -1,8 +1,11 @@
+import io
 import logging
+import sys
 import unittest
+import unittest.mock
 from textwrap import dedent
 
-from snekbox.nsjail import MEM_MAX, NsJail
+from snekbox.nsjail import MEM_MAX, NsJail, OUTPUT_MAX, READ_CHUNK_SIZE
 
 
 class NsJailTests(unittest.TestCase):
@@ -174,3 +177,25 @@ class NsJailTests(unittest.TestCase):
             msg="stdout does not come before stderr"
         )
         self.assertEqual(result.stderr, None)
+
+    def test_stdout_flood_results_in_graceful_sigterm(self):
+        stdout_flood = dedent("""
+            while True:
+                print('abcdefghij')
+        """).strip()
+
+        result = self.nsjail.python3(stdout_flood)
+        self.assertEqual(result.returncode, 137)
+
+    def test_large_output_is_truncated(self):
+        chunk = "a" * READ_CHUNK_SIZE
+        expected_chunks = OUTPUT_MAX // sys.getsizeof(chunk) + 1
+
+        nsjail_subprocess = unittest.mock.Mock()
+
+        # Go 10 chunks over to make sure we exceed the limit
+        nsjail_subprocess.stdout = io.StringIO((expected_chunks + 10) * chunk)
+        nsjail_subprocess.poll.return_value = None
+
+        output = self.nsjail._consume_stdout(nsjail_subprocess)
+        self.assertEqual(output, chunk * expected_chunks)