fix(server): guard WS on_message against non-dict JSON + unknown types (#805)

buckaroo/server/websocket_handler.py

@@ -42,11 +42,27 @@ def on_message(self, message):
             self.write_message(json.dumps({"type": "error", "error_code": "invalid_json", "message": "Invalid JSON"}))
             return
 
+        # Reject non-object JSON (null, arrays, scalars) explicitly.
+        # Without this guard, ``msg.get(...)`` below raises ``AttributeError``
+        # on ``None`` etc., Tornado swallows it, and the WS closes
+        # permanently. See #805.
+        if not isinstance(msg, dict):
+            self.write_message(json.dumps({"type": "error",
+                "error_code": "invalid_message_shape",
+                "message": f"WS messages must be JSON objects; got {type(msg).__name__}"}))
+            return
+
         msg_type = msg.get("type")
         if msg_type == "infinite_request":
             self._handle_infinite_request(msg.get("payload_args", {}))
         elif msg_type == "buckaroo_state_change":
             self._handle_buckaroo_state_change(msg.get("new_state") or {})
+        else:
+            # Pre-#805 unknown / missing ``type`` was silently dropped —
+            # clients had no way to debug. Tell them what they sent.
+            self.write_message(json.dumps({"type": "error",
+                "error_code": "unknown_message_type",
+                "message": f"Unknown WS message type: {msg_type!r}"}))
 
     def _handle_buckaroo_state_change(self, new_state):
         sessions = self.application.settings["sessions"]

tests/unit/server/test_load_expr.py

@@ -1,5 +1,6 @@
 """End-to-end tests for POST /load_expr — server load path for
 XorqBuckarooInfiniteWidget over a xorq/ibis expression."""
+import asyncio
 import io
 import json
 import os
@@ -141,6 +142,67 @@ async def test_ws_search_pushdown(self):
         finally:
             shutil.rmtree(builds_root, ignore_errors=True)
 
+    @tornado.testing.gen_test
+    async def test_ws_message_robustness(self):
+        """Regression for #805: ``on_message`` did ``msg.get(...)`` on
+        whatever ``json.loads`` returned, which is unsafe when the JSON
+        is not an object (``null``, bare arrays, scalars). ``null`` in
+        particular killed the WS — ``None.get`` raises ``AttributeError``,
+        Tornado swallows it, the stream closes.
+
+        Adjacent: unknown message types (``{"type": 42}``, missing
+        ``type``, empty ``{}``) were silently dropped. Clients couldn't
+        debug because no response came.
+
+        This test sends each malformed shape and asserts the server
+        returns a structured error frame, NOT a silent drop or a
+        crashed WS. After the barrage, a well-formed infinite_request
+        confirms the WS is still alive — the headline pre-fix bug was
+        that ``null`` permanently killed the stream.
+
+        No ``/load`` setup: the guards run before any session lookup,
+        so the test works against an empty session and is portable
+        across machines.
+        """
+        ws = await tornado.websocket.websocket_connect(
+            f"ws://localhost:{self.get_http_port()}/ws/ws-guard")
+
+        # Each entry: (label, raw_ws_message). Server must respond
+        # to each with a structured error frame within 3s.
+        cases = [
+            ("bare_null", "null"),
+            ("bare_array", "[1,2,3]"),
+            ("bare_scalar", "42"),
+            ("empty_object", "{}"),
+            ("missing_type", json.dumps({"payload": "x"})),
+            ("type_as_int", json.dumps({"type": 42, "payload": "x"})),
+            ("unknown_type",
+                json.dumps({"type": "buckaroo_invented_command"})),
+        ]
+        for label, raw in cases:
+            ws.write_message(raw)
+            frame = await asyncio.wait_for(ws.read_message(), timeout=3.0)
+            self.assertIsNotNone(frame, f"{label}: no response (silent drop)")
+            d = json.loads(frame)
+            self.assertEqual(d.get("type"), "error",
+                f"{label}: expected error frame, got {d.get('type')!r}")
+            self.assertIn("error_code", d,
+                f"{label}: error frame missing error_code")
+
+        # Liveness: WS survived the malformed barrage. Send a well-formed
+        # request and confirm we still get a response. With no session
+        # data loaded, _handle_infinite_request returns an infinite_resp
+        # with empty data — what matters is that *any* frame arrives.
+        ws.write_message(json.dumps({
+            "type": "infinite_request",
+            "payload_args": {"start": 0, "end": 10,
+                "sourceName": "default", "origEnd": 10}}))
+        frame = await asyncio.wait_for(ws.read_message(), timeout=3.0)
+        self.assertIsNotNone(frame, "WS died after malformed barrage")
+        d = json.loads(frame)
+        self.assertEqual(d.get("type"), "infinite_resp")
+        ws.close()
+
     @tornado.testing.gen_test
     async def test_session_reuse_xorq_then_pandas(self):
         """A client that POSTs /load_expr and then POSTs /load with the