342 lines
12 KiB
Python
342 lines
12 KiB
Python
import hashlib
|
|
import json
|
|
import stat
|
|
from collections.abc import Callable
|
|
from pathlib import Path
|
|
from typing import Any, cast
|
|
|
|
import paho.mqtt.client as mqtt
|
|
import pytest
|
|
from paho.mqtt.packettypes import PacketTypes
|
|
from paho.mqtt.reasoncodes import ReasonCode
|
|
|
|
import k1link.device_plugins.xgrids_k1.mqtt.capture as capture_module
|
|
from k1link.device_plugins.xgrids_k1.mqtt.capture import (
|
|
FRAME_HEADER,
|
|
RAW_MAGIC,
|
|
REPORT_TOPICS,
|
|
CaptureError,
|
|
CaptureFormatError,
|
|
capture_mqtt,
|
|
iter_capture_frames,
|
|
validate_private_ipv4,
|
|
)
|
|
|
|
|
|
class FakeClient:
|
|
def __init__(self, *, topic: str = "RealtimePath", payload: bytes = b"pose-data") -> None:
|
|
self.on_connect: Callable[..., None] | None = None
|
|
self.on_subscribe: Callable[..., None] | None = None
|
|
self.on_message: Callable[..., None] | None = None
|
|
self.on_disconnect: Callable[..., None] | None = None
|
|
self.topic = topic
|
|
self.payload = payload
|
|
self.connect_calls: list[tuple[str, int, int]] = []
|
|
self.subscribe_calls: list[Any] = []
|
|
self.disconnect_count = 0
|
|
self._step = 0
|
|
|
|
def connect(self, host: str, port: int, keepalive: int) -> mqtt.MQTTErrorCode:
|
|
self.connect_calls.append((host, port, keepalive))
|
|
return mqtt.MQTT_ERR_SUCCESS
|
|
|
|
def subscribe(self, topics: Any) -> tuple[mqtt.MQTTErrorCode, int]:
|
|
self.subscribe_calls.append(topics)
|
|
return mqtt.MQTT_ERR_SUCCESS, 7
|
|
|
|
def loop(self, timeout: float) -> mqtt.MQTTErrorCode:
|
|
assert timeout > 0
|
|
self._step += 1
|
|
if self._step == 1:
|
|
assert self.on_connect is not None
|
|
self.on_connect(
|
|
self,
|
|
None,
|
|
mqtt.ConnectFlags(session_present=False),
|
|
ReasonCode(PacketTypes.CONNACK, "Success"),
|
|
None,
|
|
)
|
|
elif self._step == 2:
|
|
assert self.on_subscribe is not None
|
|
self.on_subscribe(
|
|
self,
|
|
None,
|
|
7,
|
|
[ReasonCode(PacketTypes.SUBACK, identifier=0) for _ in REPORT_TOPICS],
|
|
None,
|
|
)
|
|
elif self._step == 3:
|
|
assert self.on_message is not None
|
|
message = mqtt.MQTTMessage(topic=self.topic.encode())
|
|
message.payload = self.payload
|
|
message.qos = 0
|
|
self.on_message(self, None, message)
|
|
else:
|
|
raise KeyboardInterrupt
|
|
return mqtt.MQTT_ERR_SUCCESS
|
|
|
|
def disconnect(self) -> mqtt.MQTTErrorCode:
|
|
self.disconnect_count += 1
|
|
return mqtt.MQTT_ERR_SUCCESS
|
|
|
|
|
|
@pytest.mark.parametrize("address", ["10.0.0.1", "172.16.0.1", "172.31.255.254", "192.168.4.2"])
|
|
def test_validate_private_ipv4_accepts_only_rfc1918(address: str) -> None:
|
|
assert validate_private_ipv4(address) == address
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"address",
|
|
["8.8.8.8", "127.0.0.1", "169.254.1.2", "::1", "k1.local", " 192.168.1.2"],
|
|
)
|
|
def test_validate_private_ipv4_rejects_other_targets(address: str) -> None:
|
|
with pytest.raises(ValueError, match="private IPv4|RFC1918"):
|
|
validate_private_ipv4(address)
|
|
|
|
|
|
def test_capture_writes_verifiable_frames_metadata_and_summary(tmp_path: Path) -> None:
|
|
fake = FakeClient()
|
|
observed = []
|
|
|
|
summary = capture_mqtt(
|
|
"192.168.1.50",
|
|
tmp_path / "capture",
|
|
duration_seconds=30,
|
|
on_message_recorded=observed.append,
|
|
_client_factory=lambda: cast(mqtt.Client, fake),
|
|
)
|
|
|
|
assert fake.connect_calls == [("192.168.1.50", 1883, 30)]
|
|
assert fake.subscribe_calls == [[(topic, 0) for topic in REPORT_TOPICS]]
|
|
assert fake.disconnect_count == 1
|
|
assert summary["stop_reason"] == "keyboard_interrupt"
|
|
assert summary["message_count"] == 1
|
|
assert summary["payload_bytes"] == len(fake.payload)
|
|
assert summary["subscriptions"] == list(REPORT_TOPICS)
|
|
assert len(observed) == 1
|
|
assert observed[0].topic == fake.topic
|
|
assert observed[0].payload == fake.payload
|
|
assert observed[0].received_at_epoch_ns > 0
|
|
|
|
capture_dir = tmp_path / "capture"
|
|
raw = (capture_dir / "mqtt.raw.k1mqtt").read_bytes()
|
|
assert raw.startswith(RAW_MAGIC)
|
|
topic_length, payload_length = FRAME_HEADER.unpack_from(raw, len(RAW_MAGIC))
|
|
topic_start = len(RAW_MAGIC) + FRAME_HEADER.size
|
|
payload_start = topic_start + topic_length
|
|
assert raw[topic_start:payload_start].decode() == fake.topic
|
|
assert payload_length == len(fake.payload)
|
|
assert raw[payload_start:] == fake.payload
|
|
|
|
records: list[dict[str, object]] = [
|
|
json.loads(line) for line in (capture_dir / "mqtt.metadata.jsonl").read_text().splitlines()
|
|
]
|
|
assert len(records) == 1
|
|
record = records[0]
|
|
assert record["record_type"] == "message"
|
|
assert record["topic"] == fake.topic
|
|
assert record["payload_bytes"] == len(fake.payload)
|
|
assert isinstance(record["received_at_epoch_ns"], int)
|
|
assert record["payload_sha256"] == hashlib.sha256(fake.payload).hexdigest()
|
|
assert record["raw_frame_offset"] == len(RAW_MAGIC)
|
|
assert record["raw_payload_offset"] == payload_start
|
|
|
|
frames = list(iter_capture_frames(capture_dir / "mqtt.raw.k1mqtt"))
|
|
assert len(frames) == 1
|
|
assert frames[0].sequence == 1
|
|
assert frames[0].topic == fake.topic
|
|
assert frames[0].payload == fake.payload
|
|
assert frames[0].raw_frame_offset == len(RAW_MAGIC)
|
|
assert frames[0].raw_payload_offset == payload_start
|
|
|
|
saved_summary = json.loads((capture_dir / "mqtt.summary.json").read_text())
|
|
assert saved_summary == summary
|
|
assert saved_summary["artifact_hashes"]["raw_sha256"] == hashlib.sha256(raw).hexdigest()
|
|
for artifact_name in ("mqtt.raw.k1mqtt", "mqtt.metadata.jsonl", "mqtt.summary.json"):
|
|
assert stat.S_IMODE((capture_dir / artifact_name).stat().st_mode) == 0o600
|
|
|
|
|
|
def test_oversize_payload_is_not_written_to_raw_capture(tmp_path: Path) -> None:
|
|
fake = FakeClient(payload=b"oversize")
|
|
capture_dir = tmp_path / "capture"
|
|
|
|
with pytest.raises(CaptureError, match="limit is 4 bytes") as error:
|
|
capture_mqtt(
|
|
"10.1.2.3",
|
|
capture_dir,
|
|
max_message_bytes=4,
|
|
_client_factory=lambda: cast(mqtt.Client, fake),
|
|
)
|
|
|
|
assert (capture_dir / "mqtt.raw.k1mqtt").read_bytes() == RAW_MAGIC
|
|
record = json.loads((capture_dir / "mqtt.metadata.jsonl").read_text())
|
|
assert record["record_type"] == "rejected_message"
|
|
assert record["payload_bytes"] == len(fake.payload)
|
|
assert error.value.summary is not None
|
|
assert error.value.summary["stop_reason"] == "message_too_large"
|
|
assert error.value.summary["message_count"] == 0
|
|
assert error.value.summary["rejected_message_count"] == 1
|
|
|
|
|
|
def test_capture_refuses_to_overwrite_existing_artifacts(tmp_path: Path) -> None:
|
|
capture_dir = tmp_path / "capture"
|
|
capture_dir.mkdir()
|
|
raw = capture_dir / "mqtt.raw.k1mqtt"
|
|
raw.write_bytes(b"existing evidence")
|
|
|
|
with pytest.raises(FileExistsError, match="refusing to overwrite"):
|
|
capture_mqtt(
|
|
"192.168.1.2",
|
|
capture_dir,
|
|
_client_factory=lambda: cast(mqtt.Client, FakeClient()),
|
|
)
|
|
|
|
assert raw.read_bytes() == b"existing evidence"
|
|
|
|
|
|
def test_capture_can_be_stopped_by_owner_without_losing_artifacts(tmp_path: Path) -> None:
|
|
fake = FakeClient()
|
|
stop_checks = 0
|
|
|
|
def should_stop() -> bool:
|
|
nonlocal stop_checks
|
|
stop_checks += 1
|
|
return stop_checks >= 4
|
|
|
|
summary = capture_mqtt(
|
|
"192.168.1.50",
|
|
tmp_path / "capture",
|
|
duration_seconds=30,
|
|
should_stop=should_stop,
|
|
_client_factory=lambda: cast(mqtt.Client, fake),
|
|
)
|
|
|
|
assert summary["stop_reason"] == "external_stop"
|
|
assert summary["message_count"] == 1
|
|
assert list(iter_capture_frames(tmp_path / "capture" / "mqtt.raw.k1mqtt"))[0].payload
|
|
|
|
|
|
def test_preview_failure_happens_after_raw_message_is_preserved(tmp_path: Path) -> None:
|
|
fake = FakeClient()
|
|
capture_dir = tmp_path / "capture"
|
|
|
|
def fail_preview(_message: object) -> None:
|
|
raise ValueError("synthetic preview failure")
|
|
|
|
with pytest.raises(CaptureError, match="preview callback failed") as error:
|
|
capture_mqtt(
|
|
"192.168.1.50",
|
|
capture_dir,
|
|
on_message_recorded=fail_preview,
|
|
_client_factory=lambda: cast(mqtt.Client, fake),
|
|
)
|
|
|
|
assert error.value.summary is not None
|
|
assert error.value.summary["message_count"] == 1
|
|
frames = list(iter_capture_frames(capture_dir / "mqtt.raw.k1mqtt"))
|
|
assert frames[0].payload == fake.payload
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("raw", "message"),
|
|
[
|
|
(b"not-mqtt", "magic/version"),
|
|
(RAW_MAGIC + b"\x00", "truncated length header"),
|
|
(RAW_MAGIC + FRAME_HEADER.pack(4, 1) + b"ab", "topic .* is truncated"),
|
|
(RAW_MAGIC + FRAME_HEADER.pack(1, 4) + b"t" + b"ab", "payload .* is truncated"),
|
|
(RAW_MAGIC + FRAME_HEADER.pack(1, 5) + b"t" + b"abcde", "exceeds 4"),
|
|
],
|
|
)
|
|
def test_capture_reader_rejects_invalid_or_unbounded_frames(
|
|
tmp_path: Path,
|
|
raw: bytes,
|
|
message: str,
|
|
) -> None:
|
|
path = tmp_path / "capture.raw"
|
|
path.write_bytes(raw)
|
|
|
|
with pytest.raises(CaptureFormatError, match=message):
|
|
list(iter_capture_frames(path, max_payload_bytes=4))
|
|
|
|
|
|
def _mqtt_message(topic: str, payload: bytes) -> mqtt.MQTTMessage:
|
|
message = mqtt.MQTTMessage(topic=topic.encode("utf-8"))
|
|
message.payload = payload
|
|
message.qos = 0
|
|
message.retain = False
|
|
message.dup = False
|
|
return message
|
|
|
|
|
|
def test_group_commit_fsyncs_raw_before_publishing_metadata(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
writer = capture_module._CaptureWriter(tmp_path / "capture", 1024) # noqa: SLF001
|
|
writer.open()
|
|
assert writer._raw is not None # noqa: SLF001
|
|
assert writer._metadata is not None # noqa: SLF001
|
|
raw_fd = writer._raw.fileno() # noqa: SLF001
|
|
metadata_fd = writer._metadata.fileno() # noqa: SLF001
|
|
fsync_calls: list[int] = []
|
|
real_fsync = capture_module.os.fsync
|
|
|
|
def observe_fsync(descriptor: int) -> None:
|
|
fsync_calls.append(descriptor)
|
|
real_fsync(descriptor)
|
|
|
|
monkeypatch.setattr(capture_module.os, "fsync", observe_fsync)
|
|
monkeypatch.setattr(capture_module, "GROUP_COMMIT_MAX_MESSAGES", 2)
|
|
|
|
writer.record(_mqtt_message("RealtimePath", b"one"))
|
|
assert writer.metadata_path.read_bytes() == b""
|
|
writer.record(_mqtt_message("RealtimePath", b"two"))
|
|
|
|
assert fsync_calls[:2] == [raw_fd, metadata_fd]
|
|
assert len(writer.metadata_path.read_text(encoding="utf-8").splitlines()) == 2
|
|
writer.close()
|
|
|
|
|
|
def test_group_commit_timer_bounds_quiet_stream_rpo(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
writer = capture_module._CaptureWriter(tmp_path / "capture", 1024) # noqa: SLF001
|
|
writer.open()
|
|
writer.record(_mqtt_message("RealtimePath", b"one"))
|
|
assert writer.metadata_path.read_bytes() == b""
|
|
|
|
writer.maybe_commit(
|
|
writer._last_commit_monotonic # noqa: SLF001
|
|
+ capture_module.GROUP_COMMIT_INTERVAL_SECONDS
|
|
)
|
|
|
|
assert len(writer.metadata_path.read_text(encoding="utf-8").splitlines()) == 1
|
|
writer.close()
|
|
|
|
|
|
def test_raw_fsync_failure_never_publishes_metadata_ahead_of_raw(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
writer = capture_module._CaptureWriter(tmp_path / "capture", 1024) # noqa: SLF001
|
|
writer.open()
|
|
writer.record(_mqtt_message("RealtimePath", b"one"))
|
|
assert writer._raw is not None # noqa: SLF001
|
|
raw_fd = writer._raw.fileno() # noqa: SLF001
|
|
real_fsync = capture_module.os.fsync
|
|
|
|
def fail_raw_fsync(descriptor: int) -> None:
|
|
if descriptor == raw_fd:
|
|
raise OSError("synthetic raw fsync failure")
|
|
real_fsync(descriptor)
|
|
|
|
monkeypatch.setattr(capture_module.os, "fsync", fail_raw_fsync)
|
|
|
|
with pytest.raises(OSError, match="synthetic raw fsync failure"):
|
|
writer._commit_pending() # noqa: SLF001
|
|
|
|
assert writer.metadata_path.read_bytes() == b""
|
|
# Restore durability primitive so the fixture can close normally.
|
|
monkeypatch.setattr(capture_module.os, "fsync", real_fsync)
|
|
writer.close()
|