diff options
| -rw-r--r-- | snekbox/nsjail.py | 3 | ||||
| -rw-r--r-- | tests/test_nsjail.py | 20 |
2 files changed, 21 insertions, 2 deletions
diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py index 8edf546..4dea3cb 100644 --- a/snekbox/nsjail.py +++ b/snekbox/nsjail.py @@ -29,6 +29,7 @@ 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: @@ -122,7 +123,7 @@ class NsJail: # We'll consume STDOUT as long as the NsJail subprocess is running. while nsjail.poll() is None: - chars = nsjail.stdout.read(10_000) + chars = nsjail.stdout.read(READ_CHUNK_SIZE) output_size += sys.getsizeof(chars) output.append(chars) diff --git a/tests/test_nsjail.py b/tests/test_nsjail.py index 691eb5b..f4e678c 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, READ_CHUNK_SIZE, OUTPUT_MAX class NsJailTests(unittest.TestCase): @@ -183,3 +186,18 @@ class NsJailTests(unittest.TestCase): 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.returncode = -9 + nsjail_subprocess.poll.return_value = None + + returncode, output = self.nsjail._consume_stdout(nsjail_subprocess) + self.assertEqual(returncode, 137) + self.assertEqual(output, chunk * expected_chunks) |