diff options
author | 2021-02-04 19:04:29 -0800 | |
---|---|---|
committer | 2021-02-04 19:30:42 -0800 | |
commit | a8b9c331fd7a69c9a89083886b7b47b3bba6f78a (patch) | |
tree | 85e80936f1dfcf3e927d15c8fd8eaf1244446ea5 | |
parent | Suppress NsJail info logs during tests (diff) |
Close file descriptors when subprocess ends
-rw-r--r-- | snekbox/nsjail.py | 27 | ||||
-rw-r--r-- | tests/test_nsjail.py | 2 |
2 files changed, 14 insertions, 15 deletions
diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py index 55ece65..d48f72b 100644 --- a/snekbox/nsjail.py +++ b/snekbox/nsjail.py @@ -148,20 +148,19 @@ class NsJail: 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() + # Context manager will wait for process to terminate and close file descriptors. + 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) + 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 return "".join(output) diff --git a/tests/test_nsjail.py b/tests/test_nsjail.py index 7e7df9b..820f57e 100644 --- a/tests/test_nsjail.py +++ b/tests/test_nsjail.py @@ -192,7 +192,7 @@ class NsJailTests(unittest.TestCase): chunk = "a" * READ_CHUNK_SIZE expected_chunks = OUTPUT_MAX // sys.getsizeof(chunk) + 1 - nsjail_subprocess = unittest.mock.Mock() + nsjail_subprocess = unittest.mock.MagicMock() # Go 10 chunks over to make sure we exceed the limit nsjail_subprocess.stdout = io.StringIO((expected_chunks + 10) * chunk) |