Keep output that took us over the output limit

Previously, the chunk of output that took us over the output limit was
dismissed. As we've already consumed it and it's not going to have a
very large size, we might as well include it in the final output we
return.

1 files changed, 6 insertions, 10 deletions

diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index 80d9fd5..9b39789 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -120,22 +120,18 @@ class NsJail:
         output_size = 0
         output = []
 
-        # We'll consume STDOUT as long as the NsJail subprocess is running
+        # We'll consume STDOUT as long as the NsJail subprocess is running.
         while nsjail.poll() is None:
-            # Read 100 characters from the STDOUT stream
             chars = nsjail.stdout.read(100)
-            chars_size = sys.getsizeof(chars)
+            output_size += sys.getsizeof(chars)
+            output.append(chars)
 
-            # Check if these characters take us over the output limit
-            if output_size + chars_size > OUTPUT_MAX:
-                # Ask nsjail to terminate itself using SIGTERM
+            if output_size > OUTPUT_MAX:
+                # Ask NsJail to terminate as we've gone over the output limit.
                 nsjail.terminate()
                 break
 
-            output_size += chars_size
-            output.append(chars)
-
-        # Ensure that we wait for the nsjail process to terminate
+        # Ensure that we wait for the NsJail subprocess to terminate.
         nsjail.wait()
 
         return "".join(output)