diff --git a/kobrax_client.py b/kobrax_client.py index 40230f6..a234e73 100644 --- a/kobrax_client.py +++ b/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: diff --git a/tests/test_client_robustness.py b/tests/test_client_robustness.py new file mode 100644 index 0000000..bcf7d3a --- /dev/null +++ b/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()