Validate IMU packet CRC

lib/ninedof_receiver.py

@@ -2,10 +2,12 @@
 
 import json
 import socket
+import struct
 import threading
 import time
 from typing import Any
 
+from lib.crc import crc32_ieee
 from lib.json_data_handler import JSONDataHandler
 from lib.runtime_paths import log_path, logs_dir
 
@@ -52,6 +54,7 @@ def __init__(self, host=UDP_IP, port=UDP_PORT, data_handler=None):
         # Stats
         self._lock = threading.Lock()
         self._packet_count = 0
+        self._crc_errors = 0
         self._last_data = {}
         self._last_recv_time = None
 
@@ -122,6 +125,7 @@ def get_stats(self) -> dict:
                 age_ms = round((time.monotonic() - self._last_recv_time) * 1000)
             return {
                 "packet_count": self._packet_count,
+                "crc_errors": self._crc_errors,
                 "last_data": self._last_data.copy(),
                 "age_ms": age_ms,
                 "tare_offset": self._tare_offset.copy(),
@@ -162,12 +166,10 @@ def _run(self):
 
     def _process_packet(self, data: bytes, addr: tuple):
         """Process incoming UDP packet with IMU data."""
-        try:
-            text = data.decode("utf-8", errors="strict")
-            msg = json.loads(text)
-        except Exception as e:
-            print(f"IMU: Bad JSON from {addr}: {e}")
+        decoded = self._decode_packet(data, addr)
+        if decoded is None:
             return
+        text, msg = decoded
 
         self._log_raw_packet(text)
 
@@ -243,6 +245,45 @@ def _val(key: str, default: float = float("nan")) -> float:
         except Exception as e:
             print(f"IMU: Error updating data: {e}")
 
+    def _decode_packet(self, data: bytes, addr: tuple) -> tuple[str, dict] | None:
+        """Decode legacy JSON or JSON followed by a big-endian CRC32."""
+        json_error = None
+        try:
+            text = data.decode("utf-8", errors="strict")
+            try:
+                return text, json.loads(text)
+            except Exception as exc:
+                json_error = exc
+        except UnicodeDecodeError as exc:
+            json_error = exc
+
+        if len(data) <= 4:
+            print(f"IMU: Bad JSON from {addr}: {json_error}")
+            return None
+
+        body = data[:-4]
+        try:
+            text = body.decode("utf-8", errors="strict")
+        except UnicodeDecodeError:
+            print(f"IMU: Bad JSON from {addr}: {json_error}")
+            return None
+
+        try:
+            msg = json.loads(text)
+        except Exception as exc:
+            print(f"IMU: Bad JSON from {addr}: {exc}")
+            return None
+
+        recv_crc = struct.unpack("!I", data[-4:])[0]
+        calc_crc = crc32_ieee(body)
+        if calc_crc != recv_crc:
+            with self._lock:
+                self._crc_errors += 1
+            print(f"IMU: CRC mismatch from {addr}: calc=0x{calc_crc:08X}, recv=0x{recv_crc:08X}")
+            return None
+
+        return text, msg
+
     def _log_raw_packet(self, text: str) -> None:
         try:
             with IMU_LOG.open("a", encoding="utf-8") as fp:

tests/test_protocols.py

@@ -1,8 +1,10 @@
+import json
 import struct
 
 import pytest
 
 import lib.control_telemetry as control_telem
+import lib.ninedof_receiver as imu_telem
 import lib.resource_receiver as resource_telem
 from lib import axis_config_sender, bitmask, crc, net_transport, pid_config_client, system_control_client
 from lib.json_data_handler import JSONDataHandler
@@ -114,6 +116,70 @@ def test_resource_telemetry_updates_json_handler(monkeypatch, tmp_path):
     assert receiver.get_udp_counters() == (84, 0)
 
 
+def test_imu_receiver_accepts_crc_trailer(monkeypatch, tmp_path, capsys):
+    monkeypatch.setattr(imu_telem, "IMU_LOG", tmp_path / "imu_raw.ndjson")
+    handler = DummyHandler()
+    receiver = imu_telem.IMUReceiver(data_handler=handler)
+    body = json.dumps(
+        {
+            "imu": {
+                "yaw": 1.25,
+                "pitch": -2.5,
+                "roll": 3.75,
+                "yr": 0.1,
+                "pr": 0.2,
+                "rr": 0.3,
+                "ax": 4.0,
+                "ay": 5.0,
+                "az": 6.0,
+            }
+        },
+        separators=(",", ":"),
+    ).encode()
+    packet = body + struct.pack("!I", crc.crc32_ieee(body))
+
+    receiver._process_packet(packet, ("10.77.0.2", 5002))  # pylint: disable=protected-access
+
+    assert capsys.readouterr().out == ""
+    assert handler.last_update["imu"]["yaw"] == 1.25
+    assert handler.last_update["imu"]["pitch"] == -2.5
+    assert handler.last_update["imu"]["roll"] == 3.75
+    assert receiver.get_stats()["packet_count"] == 1
+    assert receiver.get_stats()["crc_errors"] == 0
+
+
+def test_imu_receiver_keeps_legacy_json_without_crc(monkeypatch, tmp_path, capsys):
+    monkeypatch.setattr(imu_telem, "IMU_LOG", tmp_path / "imu_raw.ndjson")
+    handler = DummyHandler()
+    receiver = imu_telem.IMUReceiver(data_handler=handler)
+    packet = json.dumps({"imu": {"yaw": 4.0, "pitch": 5.0, "roll": 6.0}}, separators=(",", ":")).encode()
+
+    receiver._process_packet(packet, ("10.77.0.2", 5002))  # pylint: disable=protected-access
+
+    assert capsys.readouterr().out == ""
+    assert handler.last_update["imu"]["yaw"] == 4.0
+    assert handler.last_update["imu"]["pitch"] == 5.0
+    assert handler.last_update["imu"]["roll"] == 6.0
+    assert receiver.get_stats()["packet_count"] == 1
+    assert receiver.get_stats()["crc_errors"] == 0
+
+
+def test_imu_receiver_discards_bad_crc(monkeypatch, tmp_path, capsys):
+    monkeypatch.setattr(imu_telem, "IMU_LOG", tmp_path / "imu_raw.ndjson")
+    handler = DummyHandler()
+    receiver = imu_telem.IMUReceiver(data_handler=handler)
+    body = json.dumps({"imu": {"yaw": 10.0, "pitch": 20.0, "roll": 30.0}}, separators=(",", ":")).encode()
+    packet = body + struct.pack("!I", crc.crc32_ieee(body) ^ 0xFFFFFFFF)
+
+    receiver._process_packet(packet, ("10.77.0.2", 5002))  # pylint: disable=protected-access
+
+    assert handler.last_update is None
+    stats = receiver.get_stats()
+    assert stats["packet_count"] == 0
+    assert stats["crc_errors"] == 1
+    assert "IMU: CRC mismatch" in capsys.readouterr().out
+
+
 def test_json_data_handler_creates_parent_and_preserves_sections(tmp_path):
     data_file = tmp_path / "nested" / "data.json"
     handler = JSONDataHandler(file_path=data_file)