svfs/tests/test_client_mqtt.py

#!/usr/bin/env python3
"""MQTT 推送通道的测试。

分两层：

* **装配层**：注入假客户端，确定性地验证"连接参数、订阅哪两条主题、状态机、
  重连语义、缺 paho 时的降级"——不依赖网络与 broker；
* **端到端**：自建一个最小 MQTT broker（3.1.1、QoS 0，只实现 CONNECT /
  SUBSCRIBE / PUBLISH / PINGREQ / DISCONNECT），用**真 paho 客户端**连上去，
  让 broker 主动推消息，验证消息真的落进总线：主机消息进消息中心、插件消息
  只进插件队列。没装 paho 时这一层自动跳过。
"""

import json
import socket
import struct
import sys
import threading
import time
import unittest

from helpers import VFSTestCase

from svfs_client import bus as bus_mod
from svfs_client import mqtt as mqtt_mod
from svfs_client import settings, topics

assert sys.modules.get("helpers") is not None

try:
    import paho.mqtt.client as _paho  # noqa: F401 - 端到端测试需要
    HAVE_PAHO = True
except Exception:
    HAVE_PAHO = False


def config(**overrides):
    base = {"enabled": True, "host": "127.0.0.1", "port": 1883, "tls": False,
            "prefix": "svfs", "org": "acme", "username": "alice",
            "password": "tok", "keepalive": 30}
    base.update(overrides)
    return base


def wait_until(predicate, timeout=10.0, interval=0.02):
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if predicate():
            return True
        time.sleep(interval)
    return False


class _FakeClient:
    """paho 客户端的替身：记录调用，手动触发回调。"""

    def __init__(self):
        self.calls = []
        self.subscriptions = []
        self.on_connect = None
        self.on_disconnect = None
        self.on_message = None
        self.username_pw = None
        self.tls = None
        self.reconnect_delay = None

    # paho 的接口
    def connect_async(self, host, port, keepalive):
        self.calls.append(("connect_async", host, port, keepalive))

    def loop_start(self):
        self.calls.append(("loop_start",))

    def loop_stop(self):
        self.calls.append(("loop_stop",))

    def disconnect(self):
        self.calls.append(("disconnect",))

    def subscribe(self, topic):
        self.subscriptions.append(topic)
        self.calls.append(("subscribe", topic))

    def username_pw_set(self, username, password=None):
        self.username_pw = (username, password)

    def tls_set(self, **kwargs):
        self.tls = kwargs

    def reconnect_delay_set(self, min_delay, max_delay):
        self.reconnect_delay = (min_delay, max_delay)

    # 测试侧的触发点
    def fire_connect(self, code=0):
        self.on_connect(self, None, {}, code, None)

    def fire_disconnect(self, code=0):
        self.on_disconnect(self, None, None, code, None)

    def fire_message(self, topic, payload):
        self.on_message(self, None, _Message(topic, payload))


class _Message:
    def __init__(self, topic, payload):
        self.topic = topic
        self.payload = payload if isinstance(payload, bytes) else payload.encode()


class TransportWiringTest(VFSTestCase):
    """装配层：参数、订阅、状态机（不需要 broker）。"""

    def setUp(self):
        super().setUp()
        self.states = []
        self.messages = []
        self.client = _FakeClient()

    def transport(self, **overrides):
        return mqtt_mod.MqttTransport(
            config(**overrides),
            on_message=lambda topic, payload: self.messages.append((topic, payload)),
            on_state=lambda state, detail: self.states.append((state, detail)),
            client=self.client,
        )

    def test_disabled_config_does_not_connect(self):
        transport = mqtt_mod.MqttTransport({"enabled": False},
                                           client=self.client)
        self.assertEqual(transport.start(), mqtt_mod.STATE_DISABLED)
        self.assertEqual(self.client.calls, [])

    def test_missing_paho_degrades_with_hint(self):
        original = mqtt_mod._import_paho
        mqtt_mod._import_paho = lambda: None
        self.addCleanup(setattr, mqtt_mod, "_import_paho", original)
        transport = mqtt_mod.MqttTransport(config(), on_state=lambda *a: None)
        self.assertEqual(transport.start(), mqtt_mod.STATE_MISSING)
        self.assertIn("paho", transport.detail)

    def test_start_connects_subscribes_and_reports_state(self):
        transport = self.transport()
        self.assertEqual(transport.start(), mqtt_mod.STATE_CONNECTING)
        self.assertIn(("connect_async", "127.0.0.1", 1883, 30), self.client.calls)
        self.assertIn(("loop_start",), self.client.calls)
        self.assertEqual(self.client.username_pw, ("alice", "tok"))
        self.assertFalse(self.client.tls)          # 非 TLS 不调 tls_set
        self.assertEqual(self.client.reconnect_delay,
                         (mqtt_mod.RECONNECT_MIN_DELAY, mqtt_mod.RECONNECT_MAX_DELAY))

        self.client.fire_connect(0)
        self.assertTrue(transport.is_connected())
        # 订阅的是"自己的 + 组织广播"两条，不是整个组织
        self.assertEqual(self.client.subscriptions,
                         ["svfs/acme/user/alice/#", "svfs/acme/broadcast/#"])
        self.assertEqual(self.states[-1][0], mqtt_mod.STATE_CONNECTED)

    def test_reconnect_resubscribes(self):
        transport = self.transport()
        transport.start()
        self.client.fire_connect(0)
        self.client.subscriptions.clear()
        # 掉线 → paho 自动重连 → 重新订阅（broker 不记得旧会话的订阅）
        self.client.fire_disconnect(7)
        self.assertEqual(transport.state, mqtt_mod.STATE_CONNECTING)
        self.client.fire_connect(0)
        self.assertEqual(len(self.client.subscriptions), 2)
        self.assertEqual(transport.connected_count, 2)

    def test_refused_connection_reports_error(self):
        transport = self.transport()
        transport.start()
        self.client.fire_connect(5)   # 5 = 未授权
        self.assertEqual(transport.state, mqtt_mod.STATE_ERROR)
        self.assertIn("拒绝", transport.detail)

    def test_message_callback_hands_payload_to_bus(self):
        transport = self.transport()
        transport.start()
        self.client.fire_connect(0)
        self.client.fire_message("svfs/acme/user/alice/message", '{"body": "hi"}')
        self.assertEqual(self.messages,
                         [("svfs/acme/user/alice/message", b'{"body": "hi"}')])
        self.assertEqual(transport.received_count, 1)

    def test_stop_is_idempotent_and_marks_stopped(self):
        transport = self.transport()
        transport.start()
        self.client.fire_connect(0)
        transport.stop()
        self.assertEqual(transport.state, mqtt_mod.STATE_STOPPED)
        self.assertIn(("disconnect",), self.client.calls)
        self.assertIn(("loop_stop",), self.client.calls)
        # 主动断开后的 on_disconnect 不该被当成"掉线重连中"
        self.client.fire_disconnect(0)
        self.assertEqual(transport.state, mqtt_mod.STATE_STOPPED)

    def test_tls_config_sets_tls(self):
        transport = self.transport(tls=True, port=8883)
        transport.start()
        self.assertIsNotNone(self.client.tls)
        self.assertIn(("connect_async", "127.0.0.1", 8883, 30), self.client.calls)


class BrokerTest(VFSTestCase):
    """端到端：真 paho ↔ 自建最小 broker ↔ 统一总线。"""

    ORG = "acme"
    USER = "alice"

    def setUp(self):
        super().setUp()
        if not HAVE_PAHO:
            self.skipTest("没装 paho-mqtt，跳过真实 MQTT 端到端")
        self.broker = _MiniBroker()
        self.addCleanup(self.broker.close)
        self.bus = bus_mod.MessageBus(org=self.ORG, user=self.USER)
        self.transport = mqtt_mod.MqttTransport(
            config(host="127.0.0.1", port=self.broker.port, keepalive=5),
            on_message=self.bus.deliver)
        self.addCleanup(self.transport.stop)
        self.assertEqual(self.transport.start(), mqtt_mod.STATE_CONNECTING)
        self.assertTrue(wait_until(lambda: self.transport.is_connected()),
                        "客户端没能连上自建 broker")

    def test_server_push_lands_in_bus(self):
        # broker 记录客户端订阅的主题（验证最小权限订阅）
        self.assertTrue(wait_until(lambda: len(self.broker.subscriptions) >= 2))
        self.assertEqual(sorted(self.broker.subscriptions),
                         ["svfs/acme/broadcast/#", "svfs/acme/user/alice/#"])

        # 服务端主动推一条主机消息 + 一条插件消息
        self.broker.publish("svfs/acme/user/alice/message",
                            json.dumps({"id": "m1", "title": "备份完成",
                                        "body": "3 个包"}))
        self.broker.publish("svfs/acme/user/alice/plugin/com.example.dash",
                            json.dumps({"cpu": 0.42}))
        self.assertTrue(wait_until(lambda: len(self.bus.log) == 1))
        self.assertTrue(wait_until(lambda: self.bus.pending("com.example.dash") == 1))

        message = self.bus.log.items()[0]
        self.assertEqual((message["title"], message["body"]), ("备份完成", "3 个包"))
        self.assertTrue(message["unread"])
        # 插件消息没有混进消息中心
        self.assertEqual(len(self.bus.log), 1)
        self.assertEqual(self.bus.drain("com.example.dash")[0]["payload"],
                         json.dumps({"cpu": 0.42}))

    def test_qos1_style_redelivery_is_deduped_end_to_end(self):
        payload = json.dumps({"id": "m2", "body": "同一条"})
        self.broker.publish("svfs/acme/user/alice/message", payload)
        self.broker.publish("svfs/acme/user/alice/message", payload)
        self.assertTrue(wait_until(lambda: len(self.bus.log) == 1))
        time.sleep(0.2)
        self.assertEqual(len(self.bus.log), 1, "重投的两条应该去重成一条")
        self.assertEqual(self.bus.stats()["duplicate"], 1)

    def test_other_users_message_is_dropped_by_bus(self):
        """即使 broker 的 ACL 配错、把别人的消息推过来，总线也挡掉。

        正常情况根本到不了这一步：客户端只订阅 ``user/<自己>/#``，broker 不会
        投别的用户的消息（上一条测试验证的正是这个）。这里用 publish_raw 模拟
        配错的 broker，验证第二道防线。
        """
        self.broker.publish_raw("svfs/acme/user/bob/message", '{"body": "别人的"}')
        self.assertTrue(wait_until(lambda: self.bus.stats().get("forbidden", 0) == 1))
        self.assertEqual(len(self.bus.log), 0)

    def test_other_orgs_message_is_dropped_by_bus(self):
        self.broker.publish_raw("svfs/othercorp/user/alice/message",
                                '{"body": "别的租户"}')
        self.assertTrue(wait_until(lambda: self.bus.stats().get("forbidden", 0) == 1))
        self.assertEqual(len(self.bus.log), 0)


# ----------------------------------------------------------------------
# 最小 MQTT broker（只为测试：MQTT 3.1.1、QoS 0、CONNECT/SUBSCRIBE/PUBLISH/
# PINGREQ/DISCONNECT）。故意不引第三方 broker —— 测试要能离线跑。
# ----------------------------------------------------------------------


class _MiniBroker:
    """够用的 MQTT broker：接连接、记订阅、按通配符投递 PUBLISH。"""

    def __init__(self):
        self._server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._server.bind(("127.0.0.1", 0))
        self._server.listen(4)
        self.port = self._server.getsockname()[1]
        self.subscriptions = []          # broker 收到的订阅（测试断言用）
        self.published = []              # broker 收到的客户端 PUBLISH（本测试没用到）
        self._clients = []               # [(sock, [topics])]
        self._lock = threading.Lock()
        self._closed = False
        threading.Thread(target=self._accept_loop, daemon=True).start()

    # ---- 对外 -------------------------------------------------------

    def publish(self, topic, payload):
        """broker 主动推一条（QoS 0），按订阅投递。"""
        self._send(topic, payload, ignore_subscriptions=False)

    def publish_raw(self, topic, payload):
        """无视订阅、推给所有连接（模拟 ACL 配错的 broker）。

        正常 broker 不会这么干（订阅本身就是 ``user/<自己>/#``），但"broker
        配错"是真实存在的运维事故，总线的归属校验就是为它准备的兜底。
        """
        self._send(topic, payload, ignore_subscriptions=True)

    def _send(self, topic, payload, ignore_subscriptions):
        body = payload.encode("utf-8") if isinstance(payload, str) else bytes(payload)
        packet = _publish_packet(topic, body)
        with self._lock:
            if ignore_subscriptions:
                targets = [sock for sock, _filters in self._clients]
            else:
                targets = [sock for sock, filters in self._clients
                           if any(_topic_matches(f, topic) for f in filters)]
        for sock in targets:
            try:
                sock.sendall(packet)
            except OSError:
                pass

    def close(self):
        self._closed = True
        with self._lock:
            for sock, _filters in self._clients:
                try:
                    sock.close()
                except OSError:
                    pass
            self._clients = []
        try:
            self._server.close()
        except OSError:
            pass

    # ---- 内部 -------------------------------------------------------

    def _accept_loop(self):
        while not self._closed:
            try:
                sock, _addr = self._server.accept()
            except OSError:
                return
            with self._lock:
                self._clients.append((sock, []))
            threading.Thread(target=self._client_loop, args=(sock,),
                             daemon=True).start()

    def _client_loop(self, sock):
        try:
            while not self._closed:
                header = sock.recv(1)
                if not header:
                    return
                packet_type = header[0] >> 4
                length, _ = _read_length(sock)
                body = _read_exact(sock, length) if length else b""
                if packet_type == 1:            # CONNECT
                    sock.sendall(b"\x20\x02\x00\x00")
                elif packet_type == 8:          # SUBSCRIBE
                    packet_id = body[:2]
                    filters = _parse_subscribe(body[2:])
                    with self._lock:
                        for entry in self._clients:
                            if entry[0] is sock:
                                entry[1].extend(filters)
                        self.subscriptions.extend(filters)
                    sock.sendall(b"\x90" + bytes([2 + len(filters)]) + packet_id
                                 + b"\x00" * len(filters))
                elif packet_type == 3:          # PUBLISH（客户端发的）→ 转投订阅者
                    topic, payload = _parse_publish(body)
                    self.published.append(topic)
                    self._send(topic, payload, ignore_subscriptions=False)
                elif packet_type == 12:         # PINGREQ
                    sock.sendall(b"\xd0\x00")
                elif packet_type == 14:         # DISCONNECT
                    return
        except OSError:
            return
        finally:
            with self._lock:
                self._clients = [entry for entry in self._clients if entry[0] is not sock]
            try:
                sock.close()
            except OSError:
                pass


def _read_length(sock):
    """MQTT 的可变长度字段（剩下长度）。"""
    multiplier = 1
    value = 0
    for _ in range(4):
        byte = sock.recv(1)
        if not byte:
            raise OSError("连接已关闭")
        digit = byte[0]
        value += (digit & 0x7F) * multiplier
        if not digit & 0x80:
            break
        multiplier *= 128
    return value, multiplier


def _read_exact(sock, count):
    chunks = []
    remaining = count
    while remaining > 0:
        chunk = sock.recv(remaining)
        if not chunk:
            raise OSError("连接已关闭")
        chunks.append(chunk)
        remaining -= len(chunk)
    return b"".join(chunks)


def _parse_subscribe(body):
    """SUBSCRIBE 载荷：包 id(2) + 若干个 [主题长度(2) + 主题 + QoS(1)]。"""
    filters = []
    offset = 0
    while offset + 2 <= len(body):
        length = struct.unpack(">H", body[offset:offset + 2])[0]
        start = offset + 2
        filters.append(body[start:start + length].decode("utf-8"))
        offset = start + length + 1  # 跳过 QoS 字节（别忘了，否则后面全错位）
    return filters


def _parse_publish(body):
    length = struct.unpack(">H", body[:2])[0]
    topic = body[2:2 + length].decode("utf-8")
    return topic, body[2 + length:]


def _publish_packet(topic, payload):
    encoded = topic.encode("utf-8")
    body = struct.pack(">H", len(encoded)) + encoded + payload
    return b"\x30" + _encode_length(len(body)) + body


def _encode_length(length):
    out = bytearray()
    while True:
        digit = length % 128
        length //= 128
        if length:
            digit |= 0x80
        out.append(digit)
        if not length:
            return bytes(out)


def _topic_matches(topic_filter, topic):
    """主题过滤器匹配（支持 ``+`` 与末尾的 ``#``）。"""
    filter_parts = topic_filter.split("/")
    topic_parts = topic.split("/")
    for index, part in enumerate(filter_parts):
        if part == "#":
            return True
        if index >= len(topic_parts):
            return False
        if part != "+" and part != topic_parts[index]:
            return False
    return len(filter_parts) == len(topic_parts)


if __name__ == "__main__":
    unittest.main()