aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
authorGravatar MarkKoz <[email protected]>2021-02-04 19:04:29 -0800
committerGravatar MarkKoz <[email protected]>2021-02-04 19:30:42 -0800
commita8b9c331fd7a69c9a89083886b7b47b3bba6f78a (patch)
tree85e80936f1dfcf3e927d15c8fd8eaf1244446ea5
parentSuppress NsJail info logs during tests (diff)
Close file descriptors when subprocess ends
-rw-r--r--snekbox/nsjail.py27
-rw-r--r--tests/test_nsjail.py2
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)