JsonPickler

Author: alexmojaki

Diff for: backend/main/views.py

@@ -121,8 +121,7 @@ def run_code(self, code, source):
 
                 for call in birdseye_objects["calls"]:
                     call["function_id"] = function_ids[call["function_id"]]
-                    if isinstance(call["start_time"], str):
-                        call["start_time"] = datetime.fromisoformat(call["start_time"])
+                    call["start_time"] = datetime.fromisoformat(call["start_time"])
                     call = eye.db.Call(**call)
                     session.add(call)
                     # TODO get correct call from top level

Diff for: backend/main/workers/master.py

@@ -1,9 +1,8 @@
 import atexit
-import multiprocessing
 import queue
 from collections import defaultdict
 from functools import lru_cache
-from multiprocessing.context import Process
+from multiprocessing import Queue, Process
 from threading import Thread
 
 from main import simple_settings
@@ -15,10 +14,10 @@
 
 
 class UserProcess:
-    def __init__(self, manager):
-        self.task_queue = manager.Queue()
-        self.input_queue = manager.Queue()
-        self.result_queue = manager.Queue()
+    def __init__(self):
+        self.task_queue = Queue()
+        self.input_queue = Queue()
+        self.result_queue = Queue()
         self.awaiting_input = False
         self.process = None
         self.start_process()
@@ -85,8 +84,7 @@ def _await_result(self):
 
 def master_consumer_loop(comms: AbstractCommunications):
     comms = comms.make_master_side_communications()
-    manager = multiprocessing.Manager()
-    user_processes = defaultdict(lambda: UserProcess(manager))
+    user_processes = defaultdict(UserProcess)
 
     while True:
         entry = comms.recv_entry()

Diff for: backend/main/workers/utils.py

@@ -1,6 +1,9 @@
 import sys
 import traceback
+import multiprocessing.queues
+import json
 
+from littleutils import DecentJSONEncoder
 from sentry_sdk import capture_exception
 
 
@@ -36,6 +39,8 @@ def string(self):
 
 output_buffer = OutputBuffer()
 
+json_encoder = DecentJSONEncoder()
+
 
 def make_result(
         passed=False,
@@ -52,7 +57,7 @@ def make_result(
     if output_parts is None:
         output_parts = output_buffer.pop()
 
-    return dict(
+    result = dict(
         passed=passed,
         message=message,
         awaiting_input=awaiting_input,
@@ -61,6 +66,10 @@ def make_result(
         birdseye_objects=birdseye_objects,
         error=error,
     )
+    # Check that JSON encoding works here
+    # because failures in the queue pickling are silent
+    json_pickler.dumps(result)
+    return result
 
 
 def internal_error_result(sentry_offline=False):
@@ -87,3 +96,16 @@ def internal_error_result(sentry_offline=False):
         output_parts=[dict(color="red", text=output)],
         error=dict(traceback=tb, sentry_event=sentry_event),
     )
+
+
+# Queues don't communicate in pickle so that the worker
+# can't put something malicious for the master to unpickle
+class JsonPickler:
+    def loads(self, b):
+        return json.loads(b.decode("utf8"))
+
+    def dumps(self, x):
+        return json_encoder.encode(x).encode("utf8")
+
+
+multiprocessing.queues._ForkingPickler = json_pickler = JsonPickler()