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

@@ -141,6 +141,54 @@ 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.
+        """
+        import asyncio
+        await _post(self.get_http_port(), "/load",
+            {"session": "ws-guard",
+             "path": "/tmp/restaurant-complaints-pandas.parquet",
+             "mode": "buckaroo", "no_browser": True})
+        ws = await tornado.websocket.websocket_connect(
+            f"ws://localhost:{self.get_http_port()}/ws/ws-guard")
+        await ws.read_message()  # discard initial_state
+
+        # Each entry: (label, raw_ws_message). Server must respond
+        # to each with a structured error frame within 2s.
+        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")
+        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