fix(mqtt): harden kobrax_client.py against malformed data and concurrent requests

Found during a targeted code review, not from a user report:

- _drain() could get permanently stuck: valid JSON that isn't an object
  (e.g. a bare number or list, which json.loads() happily accepts) crashed
  _dispatch()'s dict-oriented logic, and the exception escaped before
  self._buf was advanced past the bad packet - so the same malformed bytes
  sat at the front of the buffer and re-crashed every subsequent _drain()
  call, forcing a reconnect each time. Now the buffer always advances (via
  try/finally) and _dispatch() rejects non-dict payloads with a warning log
  instead of crashing.

- _pending_msgid/_pending_report were mutated without synchronization
  while the reader thread reads/resolves them in _dispatch() - a
  check-then-set race on the shared report_key slot meant two concurrent
  publish() calls for the same msg_type could have one silently miss its
  reply. Added a dedicated lock around all registration/cleanup.

- The report-topic resolution path never verified a reply's msgid matched
  the waiter it was about to resolve, so a late reply for an
  already-timed-out request could be delivered to an unrelated, newer
  caller waiting on the same report_key. Entries now carry their own msgid
  for this comparison; replies without a msgid still resolve normally
  (most printer push-reports don't carry one).

- upload_gcode() silently sent a request with an empty session token when
  the upload URL was missing "?s=", instead of raising a clear error at
  the actual point of failure. It also leaked the upload socket's file
  descriptor on any send/recv failure other than a timeout, since there
  was no try/finally around its lifetime.

New tests in tests/test_client_robustness.py cover all of the above.

kobrax_client.py

@@ -178,6 +178,14 @@ class KobraXClient:
         self._pending_msgid: dict[str, dict] = {}
         # Pending requests by msg_type/report topic suffix
         self._pending_report: dict[str, dict] = {}
+        # Guards _pending_msgid/_pending_report against concurrent mutation:
+        # the reader thread resolves entries in _dispatch() while publish()
+        # (called from the poll loop and, via run_in_executor, HTTP handler
+        # threads) registers/cleans them up - without this, two concurrent
+        # publish() calls for the same msg_type can race on the
+        # check-then-set for a report_key slot, and _dispatch() could observe
+        # a dict mid-mutation.
+        self._pending_lock = threading.Lock()
 
         # Optional callbacks: topic_suffix → callable(payload_dict)
         self.callbacks: dict[str, callable] = {}
@@ -446,34 +454,46 @@ class KobraXClient:
     def _drain(self):
         buf = self._buf
         idx = 0
-        while idx < len(buf):
-            ptype = buf[idx] & 0xF0
-            i = idx + 1
-            mul = 1
-            rem = 0
-            while i < len(buf):
-                b = buf[i]
-                rem += (b & 0x7F) * mul
-                mul *= 128
-                i += 1
-                if not (b & 0x80):
+        try:
+            while idx < len(buf):
+                ptype = buf[idx] & 0xF0
+                i = idx + 1
+                mul = 1
+                rem = 0
+                while i < len(buf):
+                    b = buf[i]
+                    rem += (b & 0x7F) * mul
+                    mul *= 128
+                    i += 1
+                    if not (b & 0x80):
+                        break
+                if i + rem > len(buf):
                     break
-            if i + rem > len(buf):
-                break
-            pkt = buf[i:i + rem]
-            idx = i + rem
+                pkt = buf[i:i + rem]
+                idx = i + rem
 
-            if ptype == 0x30:
-                topic, raw_payload = _parse_publish(pkt)
-                if topic is None:
-                    continue
-                try:
-                    payload = json.loads(raw_payload)
-                except Exception:
-                    payload = {"_raw": raw_payload.decode("utf-8", errors="replace")}
-                self._dispatch(topic, payload)
-
-        self._buf = buf[idx:]
+                if ptype == 0x30:
+                    topic, raw_payload = _parse_publish(pkt)
+                    if topic is None:
+                        continue
+                    try:
+                        payload = json.loads(raw_payload)
+                    except Exception:
+                        payload = {"_raw": raw_payload.decode("utf-8", errors="replace")}
+                    try:
+                        self._dispatch(topic, payload)
+                    except Exception as e:
+                        # A single malformed/unexpected message (e.g. valid JSON
+                        # that isn't an object, like a bare number or list) must
+                        # not be reprocessed forever: without this, an exception
+                        # here would skip the buffer-advance below, leaving the
+                        # same bad packet at the front of self._buf so every
+                        # future _drain() call crashes on it again - each one
+                        # forcing a reconnect via the reader loop's exception
+                        # handler, an endless self-inflicted reconnect loop.
+                        log.warning("dispatch error for %s: %s", topic, e)
+        finally:
+            self._buf = buf[idx:]
 
     def _dedup_hash(self, suffix: str, payload: dict) -> str:
         """Hash payload ignoring volatile per-tick fields for dedup check."""
@@ -484,6 +504,9 @@ class KobraXClient:
         return hashlib.md5(json.dumps(stable, sort_keys=True).encode(), usedforsecurity=False).hexdigest()
 
     def _dispatch(self, topic: str, payload: dict):
+        if not isinstance(payload, dict):
+            log.warning("dispatch: non-dict payload on %s: %r", topic, payload)
+            return
         suffix = "/".join(topic.split("/")[-2:])
 
         if self._raw_log:
@@ -511,18 +534,29 @@ class KobraXClient:
                 log.info("RX %-25s  state=%-12s  data=%s",
                          suffix, state, json.dumps(payload.get("data"), ensure_ascii=False))
 
-        # Resolve by report topic suffix (e.g. "info/report")
-        if suffix in self._pending_report:
-            entry = self._pending_report[suffix]
-            entry["result"] = payload
-            entry["event"].set()
+        msgid = payload.get("msgid")
+        with self._pending_lock:
+            report_entry = self._pending_report.get(suffix)
+            msgid_entry  = self._pending_msgid.get(msgid) if msgid else None
+
+        # Resolve by report topic suffix (e.g. "info/report"). If the payload
+        # carries a msgid that doesn't match what this waiter is actually
+        # expecting, it's a stale/late reply for a different, already-timed-out
+        # request that happens to share the same report_key - don't deliver it
+        # to the wrong caller.
+        if report_entry is not None:
+            entry_msgid = report_entry.get("msgid")
+            if not entry_msgid or not msgid or entry_msgid == msgid:
+                report_entry["result"] = payload
+                report_entry["event"].set()
+            else:
+                log.debug("dispatch: msgid mismatch for %s report (waiting=%s, got=%s) - ignoring stale reply",
+                           suffix, entry_msgid, msgid)
 
         # Resolve by msgid (for generic response ACK)
-        msgid = payload.get("msgid")
-        if msgid and msgid in self._pending_msgid:
-            entry = self._pending_msgid[msgid]
-            entry["result"] = payload
-            entry["event"].set()
+        if msgid_entry is not None:
+            msgid_entry["result"] = payload
+            msgid_entry["event"].set()
 
         # User callbacks by topic suffix (last two path components)
         if suffix in self.callbacks:
@@ -558,13 +592,18 @@ class KobraXClient:
         # Also register by report topic as fallback for responses without msgid.
         report_key = f"{msg_type}/report"
         event  = threading.Event()
-        entry  = {"event": event, "result": None}
-        self._pending_msgid[msgid] = entry
-        # Only register report-key waiter if nobody else is waiting on it
+        # entry carries its own msgid so _dispatch()'s report-suffix path can
+        # confirm a reply actually belongs to THIS request before delivering
+        # it - without that, a late reply for an already-timed-out request A
+        # could be handed to a newer request B waiting on the same report_key.
+        entry  = {"event": event, "result": None, "msgid": msgid}
         report_registered = False
-        if report_key not in self._pending_report:
-            self._pending_report[report_key] = entry
-            report_registered = True
+        with self._pending_lock:
+            self._pending_msgid[msgid] = entry
+            # Only register report-key waiter if nobody else is waiting on it
+            if report_key not in self._pending_report:
+                self._pending_report[report_key] = entry
+                report_registered = True
 
         topic = self._pub_topic(msg_type)
         # Status poll TX (query/getInfo) is pure noise (every few seconds) ->
@@ -578,9 +617,10 @@ class KobraXClient:
                 self._sock.sendall(_build_publish(topic, payload))
         except Exception as e:
             log.error("send error: %s, reconnecting…", e)
-            self._pending_msgid.pop(msgid, None)
-            if report_registered:
-                self._pending_report.pop(report_key, None)
+            with self._pending_lock:
+                self._pending_msgid.pop(msgid, None)
+                if report_registered:
+                    self._pending_report.pop(report_key, None)
             # Non-blocking: never hang the poll-loop thread inside publish()
             # while a reconnect is running / during backoff (see _reconnect
             # docstring) - it must return so kobra_state can flip to "offline".
@@ -590,22 +630,25 @@ class KobraXClient:
             try:
                 with self._lock:
                     self._sock.sendall(_build_publish(topic, payload))
-                self._pending_msgid[msgid] = entry
-                if report_registered:
-                    self._pending_report[report_key] = entry
+                with self._pending_lock:
+                    self._pending_msgid[msgid] = entry
+                    if report_registered:
+                        self._pending_report[report_key] = entry
             except Exception:
                 return None
 
         if timeout <= 0:
-            self._pending_msgid.pop(msgid, None)
-            if report_registered:
-                self._pending_report.pop(report_key, None)
+            with self._pending_lock:
+                self._pending_msgid.pop(msgid, None)
+                if report_registered:
+                    self._pending_report.pop(report_key, None)
             return None
 
         received = event.wait(timeout)
-        self._pending_msgid.pop(msgid, None)
-        if report_registered:
-            self._pending_report.pop(report_key, None)
+        with self._pending_lock:
+            self._pending_msgid.pop(msgid, None)
+            if report_registered:
+                self._pending_report.pop(report_key, None)
         if not received:
             return None
         return entry["result"]
@@ -710,7 +753,9 @@ class KobraXClient:
                 raise RuntimeError("Could not get info/report for upload URL")
             upload_url = info["data"]["urls"]["fileUploadurl"]
         # parse token from URL query string
-        token = upload_url.split("?s=")[1] if "?s=" in upload_url else ""
+        if "?s=" not in upload_url:
+            raise RuntimeError(f"Upload: no session token ('?s=') in upload URL: {upload_url!r}")
+        token = upload_url.split("?s=")[1]
 
         with open(filepath, "rb") as f:
             file_data = f.read()
@@ -760,19 +805,25 @@ class KobraXClient:
         # (the printer processes the file before replying).
         _ai  = socket.getaddrinfo(self.host, 18910, socket.AF_INET, socket.SOCK_STREAM)
         sock = socket.create_connection(_ai[0][4], timeout=10)
-        sock.settimeout(None)  # blocking during send
-        sock.sendall(headers + body)
-        sock.settimeout(180)
-        response = b""
         try:
-            while True:
-                chunk = sock.recv(65536)
-                if not chunk:
-                    break
-                response += chunk
-        except socket.timeout:
-            pass
-        sock.close()
+            sock.settimeout(None)  # blocking during send
+            sock.sendall(headers + body)
+            sock.settimeout(180)
+            response = b""
+            try:
+                while True:
+                    chunk = sock.recv(65536)
+                    if not chunk:
+                        break
+                    response += chunk
+            except socket.timeout:
+                pass
+        finally:
+            # Without this, a sendall()/recv() failure other than
+            # socket.timeout (e.g. ConnectionResetError/BrokenPipeError if
+            # the printer drops the connection mid-upload) skipped
+            # sock.close() entirely, leaking the fd on every failed attempt.
+            sock.close()
 
         # parse HTTP response body
         if b"\r\n\r\n" in response:

tests/test_client_robustness.py

@@ -0,0 +1,189 @@
+"""Robustness fixes found during a targeted code review of kobrax_client.py:
+
+1. _drain()/_dispatch() must not get stuck reprocessing the same malformed
+   packet forever when the printer sends valid JSON that isn't an object
+   (e.g. a bare number or list) - previously an exception from _dispatch()
+   escaped before self._buf was advanced, so the same bytes sat at the
+   front of the buffer and crashed every subsequent _drain() call, each one
+   forcing a reconnect via the reader loop's exception handler.
+2. publish()'s pending-dict registration/cleanup must be safe against
+   concurrent callers for the same msg_type, and a stale/late reply must
+   not be delivered to a newer, unrelated caller waiting on the same
+   report-topic suffix.
+3. upload_gcode() must raise a clear error for a malformed upload URL
+   (missing "?s=") instead of silently sending an unauthenticated request,
+   and must not leak the upload socket on a send/recv failure.
+"""
+import argparse
+import json
+import socket
+import tempfile
+import threading
+import time
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from kobrax_client import KobraXClient, _build_publish
+
+
+def _client(**overrides):
+    kwargs = dict(
+        host="192.168.1.100", username="u", password="p",
+        mode_id="20030", device_id="abc123", port=9883,
+        client_id="test",
+    )
+    kwargs.update(overrides)
+    return KobraXClient(**kwargs)
+
+
+def _frame(topic: str, payload) -> bytes:
+    """Builds a raw MQTT PUBLISH frame carrying `payload` as JSON, the same
+    wire format _drain() parses."""
+    return _build_publish(topic, json.dumps(payload))
+
+
+# --- Fix 1: malformed (non-dict) JSON payload must not wedge the buffer ---
+
+def test_drain_advances_buffer_past_non_dict_payload():
+    c = _client()
+    # A bare JSON number is valid JSON but not a dict - json.loads succeeds,
+    # but _dispatch()'s dict-oriented logic would previously crash on it.
+    bad = _frame("anycubic/anycubicCloud/v1/printer/public/20030/abc123/info/report", 42)
+    good = _frame("anycubic/anycubicCloud/v1/printer/public/20030/abc123/info/report", {"state": "done"})
+    c._buf = bad + good
+
+    c._drain()  # must not raise
+
+    # Both packets must have been consumed - the malformed one logged/skipped,
+    # the buffer advanced past it so the following valid packet is also
+    # processed, not left stuck behind it.
+    assert c._buf == b""
+
+
+def test_drain_does_not_reprocess_bad_packet_on_repeated_calls():
+    """Regression guard for the original bug: before the fix, a crash in
+    _dispatch() left self._buf untouched, so the bad packet stayed at the
+    front and every _drain() call re-crashed on it."""
+    c = _client()
+    bad = _frame("anycubic/anycubicCloud/v1/printer/public/20030/abc123/info/report", [1, 2, 3])
+    c._buf = bad
+
+    c._drain()
+    assert c._buf == b""  # consumed, not stuck
+
+    # A second call on the now-empty buffer must be a no-op, not a re-crash.
+    c._drain()
+    assert c._buf == b""
+
+
+def test_dispatch_rejects_non_dict_payload_directly():
+    c = _client()
+    c._dispatch("anycubic/anycubicCloud/v1/printer/public/20030/abc123/info/report", "not a dict")
+    # No exception, no registered pending entries get corrupted.
+    assert c._pending_report == {}
+
+
+# --- Fix 2: pending-dict thread-safety + msgid correlation ---
+
+def test_publish_concurrent_same_msg_type_both_get_delivered():
+    """Two overlapping publish() calls for the same msg_type must each
+    receive their OWN reply, not have one silently miss out because the
+    other already claimed the shared report_key slot."""
+    c = _client()
+    c._running = True
+    c._sock = MagicMock()
+    c._ensure_reader = lambda: None  # no real reader thread needed for this test
+
+    results = {}
+
+    def call(label):
+        results[label] = c.publish("info", "query", timeout=2.0)
+
+    t1 = threading.Thread(target=call, args=("a",))
+    t2 = threading.Thread(target=call, args=("b",))
+    t1.start()
+    time.sleep(0.02)
+    t2.start()
+    time.sleep(0.05)
+
+    # Simulate the printer replying to whichever msgid-bearing requests are
+    # currently pending, by msgid (the reliable path both requests can use).
+    with c._pending_lock:
+        pending = dict(c._pending_msgid)
+    for msgid, entry in pending.items():
+        c._dispatch("anycubic/anycubicCloud/v1/printer/public/20030/abc123/info/report",
+                    {"msgid": msgid, "state": "done", "data": {"ok": True}})
+
+    t1.join(timeout=3)
+    t2.join(timeout=3)
+
+    assert results["a"] is not None
+    assert results["b"] is not None
+
+
+def test_dispatch_ignores_stale_reply_with_mismatched_msgid():
+    """A late reply carrying a msgid that doesn't match what the current
+    report_key waiter is expecting must not be delivered to it - it belongs
+    to an earlier, already-resolved/timed-out request."""
+    c = _client()
+    c._running = True
+    c._sock = MagicMock()
+
+    event = threading.Event()
+    entry = {"event": event, "result": None, "msgid": "the-real-one"}
+    with c._pending_lock:
+        c._pending_report["info/report"] = entry
+
+    # A stale reply for a different msgid must not resolve this waiter.
+    c._dispatch("anycubic/anycubicCloud/v1/printer/public/20030/abc123/info/report",
+                {"msgid": "some-other-stale-id", "state": "done"})
+    assert not event.is_set()
+    assert entry["result"] is None
+
+    # The matching reply must resolve it.
+    c._dispatch("anycubic/anycubicCloud/v1/printer/public/20030/abc123/info/report",
+                {"msgid": "the-real-one", "state": "done"})
+    assert event.is_set()
+    assert entry["result"]["msgid"] == "the-real-one"
+
+
+def test_dispatch_delivers_reply_without_msgid_as_before():
+    """Regression guard: plenty of printer reports carry no msgid at all
+    (e.g. spontaneous status pushes) - those must still resolve a waiter
+    registered without one (entry["msgid"] is falsy)."""
+    c = _client()
+    event = threading.Event()
+    entry = {"event": event, "result": None, "msgid": None}
+    with c._pending_lock:
+        c._pending_report["info/report"] = entry
+
+    c._dispatch("anycubic/anycubicCloud/v1/printer/public/20030/abc123/info/report",
+                {"state": "done"})
+    assert event.is_set()
+
+
+# --- Fix 3: upload_gcode() URL validation + socket cleanup ---
+
+def test_upload_gcode_raises_clear_error_on_missing_session_token(tmp_path):
+    c = _client()
+    f = tmp_path / "test.gcode"
+    f.write_text("; gcode")
+
+    with pytest.raises(RuntimeError, match="session token"):
+        c.upload_gcode(str(f), upload_url="http://192.168.1.100:18910/gcode_upload")
+
+
+def test_upload_gcode_closes_socket_on_send_failure(tmp_path):
+    c = _client()
+    f = tmp_path / "test.gcode"
+    f.write_text("; gcode")
+
+    fake_sock = MagicMock()
+    fake_sock.sendall.side_effect = ConnectionResetError("printer went away")
+
+    with patch("socket.create_connection", return_value=fake_sock):
+        with pytest.raises(ConnectionResetError):
+            c.upload_gcode(str(f), upload_url="http://192.168.1.100:18910/gcode_upload?s=tok123")
+
+    fake_sock.close.assert_called_once()