chongming-core/test_mqtt_auth.py

"""
MQTT Broker 认证集成测试

验证 [<env>.mqtt].username/password 配置生效：
  A. 内嵌 Broker 启用认证（FileAuthPlugin）：
     1. 正确凭据 → 连接/订阅/发布正常
     2. 错误密码 → 连接被拒（ConnectError）
     3. 匿名连接 → 连接被拒
  B. 未配置凭据（默认）→ 匿名连接正常（原行为）
  C. MQTTBridge 客户端凭据传递（含特殊字符密码的 URI 编码）
"""
import asyncio
import json
import multiprocessing
import socket
import sys
import time
from pathlib import Path
from urllib.parse import quote

sys.path.insert(0, str(Path(__file__).parent))

from amqtt.client import MQTTClient
from amqtt.errors import ConnectError

from src.chongming.app.core.lifespan_setup import _run_mqtt_broker
from src.chongming.app.core.mqtt_bridge import MQTTBridge

HOST = "127.0.0.1"
USERNAME = "cm_user"
PASSWORD = "p@ss w0rd!%#"


def cred_uri(port: int, user: str | None = None, pwd: str | None = None) -> str:
    """构建带凭据的 Broker URI（百分号编码，与 MQTTBridge._broker_uri 同语义）"""
    if user:
        return f"mqtt://{quote(user, safe='')}:{quote(pwd or '', safe='')}@{HOST}:{port}/"
    return f"mqtt://{HOST}:{port}/"


def free_port() -> int:
    with socket.socket() as s:
        s.bind((HOST, 0))
        return s.getsockname()[1]


def wait_port(port: int, timeout: float = 20.0) -> bool:
    deadline = time.time() + timeout
    while time.time() < deadline:
        with socket.socket() as s:
            s.settimeout(0.5)
            if s.connect_ex((HOST, port)) == 0:
                return True
        time.sleep(0.1)
    return False


def start_broker(port: int, username=None, password=None):
    proc = multiprocessing.Process(
        target=_run_mqtt_broker, args=(HOST, port, username, password), daemon=True
    )
    proc.start()
    assert wait_port(port), f"Broker 未在 {port} 端口就绪"
    return proc


async def try_connect(uri: str, expect_success: bool, label: str):
    client = MQTTClient(config={"auto_reconnect": False})
    try:
        await client.connect(uri)
    except ConnectError as e:
        assert not expect_success, f"{label}: 应连接成功，实际被拒 ({e})"
        return None
    assert expect_success, f"{label}: 应被拒绝，实际连接成功"
    return client


async def scenario_auth(port: int):
    """配置了账号密码 → 正确凭据通过，错误/匿名被拒"""
    proc = start_broker(port, USERNAME, PASSWORD)
    try:
        # 1. 正确凭据
        c = await try_connect(cred_uri(port, USERNAME, PASSWORD), True, "正确凭据")
        await c.subscribe([("auth/t/#", 1)])
        await c.publish("auth/t/1", b"hello", 1)
        m = await asyncio.wait_for(c.deliver_message(), timeout=10)
        assert m.data == b"hello", f"收到数据异常: {m.data}"
        await c.disconnect()
        print("PASS [A1] 正确凭据: 连接/订阅/发布/接收正常")

        # 2. 错误密码
        await try_connect(cred_uri(port, USERNAME, "wrong"), False, "错误密码")
        print("PASS [A2] 错误密码: 连接被拒")

        # 3. 匿名
        await try_connect(cred_uri(port), False, "匿名连接")
        print("PASS [A3] 匿名连接: 连接被拒")
    finally:
        proc.terminate()
        proc.join(timeout=5)


async def scenario_anonymous(port: int):
    """未配置凭据（默认）→ 匿名可用"""
    proc = start_broker(port)
    try:
        c = await try_connect(cred_uri(port), True, "匿名连接（默认）")
        await c.disconnect()
        print("PASS [B] 默认匿名模式: 匿名连接正常")
    finally:
        proc.terminate()
        proc.join(timeout=5)


class _FakeWorkerManager:
    async def call(self, *a, **kw):
        return {"ok": True}


async def scenario_bridge(port: int):
    """MQTTBridge 客户端携带凭据（特殊字符密码经 URI 百分号编码）"""
    proc = start_broker(port, USERNAME, PASSWORD)
    sub = None
    try:
        sub = MQTTClient(config={"auto_reconnect": False})
        await sub.connect(cred_uri(port, USERNAME, PASSWORD))
        await sub.subscribe([("chongming/#", 1)])

        bridge = MQTTBridge(worker_manager=_FakeWorkerManager(),
                            host=HOST, port=port,
                            username=USERNAME, password=PASSWORD)
        await bridge.start()
        assert bridge.is_running, "Bridge 未连接"

        # 发布请求 → Bridge 路由到 worker → 响应回发（验证完整链路带认证）
        pub = MQTTClient(config={"auto_reconnect": False})
        await pub.connect(cred_uri(port, USERNAME, PASSWORD))
        await pub.subscribe([("chongming/demo/op/response", 1)])
        await pub.publish(
            "chongming/demo/op",
            b'{"meta": {"timeout": 10}, "data": {"x": 1}}', 1)
        m = await asyncio.wait_for(sub.deliver_message(), timeout=10)
        assert m.topic == "chongming/demo/op", f"未收到请求: {m.topic}"
        m2 = await asyncio.wait_for(sub.deliver_message(), timeout=10)
        assert m2.topic == "chongming/demo/op/response", f"未收到响应: {m2.topic}"
        assert json.loads(m2.data) == {"ok": True}, f"响应内容异常: {m2.data}"
        await bridge.stop()
        await pub.disconnect()
        print("PASS [C] MQTTBridge 凭据传递: 连接/路由/响应链路正常（含特殊字符密码）")
    finally:
        if sub and sub.session is not None:
            try:
                await sub.disconnect()
            except Exception:
                pass
        proc.terminate()
        proc.join(timeout=5)


async def main():
    p1, p2, p3 = free_port(), free_port(), free_port()
    await scenario_auth(p1)
    await scenario_anonymous(p2)
    await scenario_bridge(p3)
    print("\n全部通过 ✔")


if __name__ == "__main__":
    multiprocessing.set_start_method("fork", force=True)
    asyncio.run(main())