fix(websocket): harden security and robustness

- Use hmac.compare_digest for timing-safe static token comparison - Add issued token capacity limit (_MAX_ISSUED_TOKENS=10000) with 429 response - Use atomic pop in _take_issued_token_if_valid to eliminate TOCTOU window - Enforce TLSv1.2 minimum version for SSL connections - Extract _safe_send helper for consistent ConnectionClosed handling - Move connection registration after ready send to prevent out-of-order delivery - Add HTTP-level allow_from check and client_id truncation in process_request - Make stop() idempotent with graceful shutdown error handling - Normalize path via validator instead of leaving raw value - Default websocket_requires_token to True for secure-by-default behavior - Add integration tests and ws_test_client helper - Refactor tests to use shared _ch factory and bus fixture
2026-04-09 18:22:35 +08:00
parent e0ccc401c0
commit 56a5906db5
4 changed files with 1102 additions and 99 deletions
@@ -3,11 +3,15 @@
 import asyncio
 import functools
 import json
+import time
+from typing import Any
 from unittest.mock import AsyncMock, MagicMock

 import httpx
 import pytest
 import websockets
+from websockets.exceptions import ConnectionClosed
+from websockets.frames import Close

 from nanobot.bus.events import OutboundMessage
 from nanobot.channels.websocket import (
@@ -21,6 +25,30 @@ from nanobot.channels.websocket import (
    _parse_request_path,
 )

+# -- Shared helpers (aligned with test_websocket_integration.py) ---------------
+
+_PORT = 29876
+
+
+def _ch(bus: Any, **kw: Any) -> WebSocketChannel:
+    cfg: dict[str, Any] = {
+        "enabled": True,
+        "allowFrom": ["*"],
+        "host": "127.0.0.1",
+        "port": _PORT,
+        "path": "/ws",
+        "websocketRequiresToken": False,
+    }
+    cfg.update(kw)
+    return WebSocketChannel(cfg, bus)
+
+
+@pytest.fixture()
+def bus() -> MagicMock:
+    b = MagicMock()
+    b.publish_inbound = AsyncMock()
+    return b
+

 async def _http_get(url: str, headers: dict[str, str] | None = None) -> httpx.Response:
    """Run GET in a thread to avoid blocking the asyncio loop shared with websockets."""
@@ -71,6 +99,21 @@ def test_parse_inbound_invalid_json_falls_back_to_raw_string() -> None:
    assert _parse_inbound_payload("{not json") == "{not json"


+@pytest.mark.parametrize(
+    ("raw", "expected"),
+    [
+        ('{"content": ""}', None),           # empty string content
+        ('{"content": 123}', None),          # non-string content
+        ('{"content": "  "}', None),         # whitespace-only content
+        ('["hello"]', '["hello"]'),           # JSON array: not a dict, treated as plain text
+        ('{"unknown_key": "val"}', None),    # unrecognized key
+        ('{"content": null}', None),         # null content
+    ],
+)
+def test_parse_inbound_payload_edge_cases(raw: str, expected: str | None) -> None:
+    assert _parse_inbound_payload(raw) == expected
+
+
 def test_web_socket_config_path_must_start_with_slash() -> None:
    with pytest.raises(ValueError, match='path must start with "/"'):
        WebSocketConfig(path="bad")
@@ -112,6 +155,14 @@ def test_issue_route_secret_matches_bearer_and_header() -> None:
    assert _issue_route_secret_matches(wrong, secret) is False


+def test_issue_route_secret_matches_empty_secret() -> None:
+    from websockets.datastructures import Headers
+
+    # Empty secret always returns True regardless of headers
+    assert _issue_route_secret_matches(Headers([]), "") is True
+    assert _issue_route_secret_matches(Headers([("Authorization", "Bearer anything")]), "") is True
+
+
@pytest.mark.asyncio
 async def test_send_delivers_json_message_with_media_and_reply() -> None:
    bus = MagicMock()
@@ -144,6 +195,33 @@ async def test_send_missing_connection_is_noop_without_error() -> None:
    await channel.send(msg)


+@pytest.mark.asyncio
+async def test_send_removes_connection_on_connection_closed() -> None:
+    bus = MagicMock()
+    channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus)
+    mock_ws = AsyncMock()
+    mock_ws.send.side_effect = ConnectionClosed(Close(1006, ""), Close(1006, ""), True)
+    channel._connections["chat-1"] = mock_ws
+
+    msg = OutboundMessage(channel="websocket", chat_id="chat-1", content="hello")
+    await channel.send(msg)
+
+    assert "chat-1" not in channel._connections
+
+
+@pytest.mark.asyncio
+async def test_send_delta_removes_connection_on_connection_closed() -> None:
+    bus = MagicMock()
+    channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"], "streaming": True}, bus)
+    mock_ws = AsyncMock()
+    mock_ws.send.side_effect = ConnectionClosed(Close(1006, ""), Close(1006, ""), True)
+    channel._connections["chat-1"] = mock_ws
+
+    await channel.send_delta("chat-1", "chunk", {"_stream_delta": True, "_stream_id": "s1"})
+
+    assert "chat-1" not in channel._connections
+
+
@pytest.mark.asyncio
 async def test_send_delta_emits_delta_and_stream_end() -> None:
    bus = MagicMock()
@@ -165,20 +243,39 @@ async def test_send_delta_emits_delta_and_stream_end() -> None:


@pytest.mark.asyncio
-async def test_end_to_end_client_receives_ready_and_agent_sees_inbound() -> None:
+async def test_send_non_connection_closed_exception_is_raised() -> None:
    bus = MagicMock()
-    bus.publish_inbound = AsyncMock()
+    channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus)
+    mock_ws = AsyncMock()
+    mock_ws.send.side_effect = RuntimeError("unexpected")
+    channel._connections["chat-1"] = mock_ws
+
+    msg = OutboundMessage(channel="websocket", chat_id="chat-1", content="hello")
+    with pytest.raises(RuntimeError, match="unexpected"):
+        await channel.send(msg)
+
+
+@pytest.mark.asyncio
+async def test_send_delta_missing_connection_is_noop() -> None:
+    bus = MagicMock()
+    channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"], "streaming": True}, bus)
+    # No exception, no error — just a no-op
+    await channel.send_delta("nonexistent", "chunk", {"_stream_delta": True, "_stream_id": "s1"})
+
+
+@pytest.mark.asyncio
+async def test_stop_is_idempotent() -> None:
+    bus = MagicMock()
+    channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus)
+    # stop() before start() should not raise
+    await channel.stop()
+    await channel.stop()
+
+
+@pytest.mark.asyncio
+async def test_end_to_end_client_receives_ready_and_agent_sees_inbound(bus: MagicMock) -> None:
    port = 29876
-    channel = WebSocketChannel(
-        {
-            "enabled": True,
-            "allowFrom": ["*"],
-            "host": "127.0.0.1",
-            "port": port,
-            "path": "/ws",
-        },
-        bus,
-    )
+    channel = _ch(bus, port=port)

    server_task = asyncio.create_task(channel.start())
    await asyncio.sleep(0.3)
@@ -212,20 +309,9 @@ async def test_end_to_end_client_receives_ready_and_agent_sees_inbound() -> None


@pytest.mark.asyncio
-async def test_token_rejects_handshake_when_mismatch() -> None:
-    bus = MagicMock()
+async def test_token_rejects_handshake_when_mismatch(bus: MagicMock) -> None:
    port = 29877
-    channel = WebSocketChannel(
-        {
-            "enabled": True,
-            "allowFrom": ["*"],
-            "host": "127.0.0.1",
-            "port": port,
-            "path": "/",
-            "token": "secret",
-        },
-        bus,
-    )
+    channel = _ch(bus, port=port, path="/", token="secret")

    server_task = asyncio.create_task(channel.start())
    await asyncio.sleep(0.3)
@@ -241,19 +327,9 @@ async def test_token_rejects_handshake_when_mismatch() -> None:


@pytest.mark.asyncio
-async def test_wrong_path_returns_404() -> None:
-    bus = MagicMock()
+async def test_wrong_path_returns_404(bus: MagicMock) -> None:
    port = 29878
-    channel = WebSocketChannel(
-        {
-            "enabled": True,
-            "allowFrom": ["*"],
-            "host": "127.0.0.1",
-            "port": port,
-            "path": "/ws",
-        },
-        bus,
-    )
+    channel = _ch(bus, port=port)

    server_task = asyncio.create_task(channel.start())
    await asyncio.sleep(0.3)
@@ -276,22 +352,13 @@ def test_registry_discovers_websocket_channel() -> None:


@pytest.mark.asyncio
-async def test_http_route_issues_token_then_websocket_requires_it() -> None:
-    bus = MagicMock()
-    bus.publish_inbound = AsyncMock()
+async def test_http_route_issues_token_then_websocket_requires_it(bus: MagicMock) -> None:
    port = 29879
-    channel = WebSocketChannel(
-        {
-            "enabled": True,
-            "allowFrom": ["*"],
-            "host": "127.0.0.1",
-            "port": port,
-            "path": "/ws",
-            "tokenIssuePath": "/auth/token",
-            "tokenIssueSecret": "route-secret",
-            "websocketRequiresToken": True,
-        },
-        bus,
+    channel = _ch(
+        bus, port=port,
+        tokenIssuePath="/auth/token",
+        tokenIssueSecret="route-secret",
+        websocketRequiresToken=True,
    )

    server_task = asyncio.create_task(channel.start())
@@ -327,3 +394,205 @@ async def test_http_route_issues_token_then_websocket_requires_it() -> None:
    finally:
        await channel.stop()
        await server_task
+
+
+@pytest.mark.asyncio
+async def test_end_to_end_server_pushes_streaming_deltas_to_client(bus: MagicMock) -> None:
+    port = 29880
+    channel = _ch(bus, port=port, streaming=True)
+
+    server_task = asyncio.create_task(channel.start())
+    await asyncio.sleep(0.3)
+
+    try:
+        async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=stream-tester") as client:
+            ready_raw = await client.recv()
+            ready = json.loads(ready_raw)
+            chat_id = ready["chat_id"]
+
+            # Server pushes deltas directly
+            await channel.send_delta(
+                chat_id, "Hello ", {"_stream_delta": True, "_stream_id": "s1"}
+            )
+            await channel.send_delta(
+                chat_id, "world", {"_stream_delta": True, "_stream_id": "s1"}
+            )
+            await channel.send_delta(
+                chat_id, "", {"_stream_end": True, "_stream_id": "s1"}
+            )
+
+            delta1 = json.loads(await client.recv())
+            assert delta1["event"] == "delta"
+            assert delta1["text"] == "Hello "
+            assert delta1["stream_id"] == "s1"
+
+            delta2 = json.loads(await client.recv())
+            assert delta2["event"] == "delta"
+            assert delta2["text"] == "world"
+            assert delta2["stream_id"] == "s1"
+
+            end = json.loads(await client.recv())
+            assert end["event"] == "stream_end"
+            assert end["stream_id"] == "s1"
+    finally:
+        await channel.stop()
+        await server_task
+
+
+@pytest.mark.asyncio
+async def test_token_issue_rejects_when_at_capacity(bus: MagicMock) -> None:
+    port = 29881
+    channel = _ch(bus, port=port, tokenIssuePath="/auth/token", tokenIssueSecret="s")
+
+    server_task = asyncio.create_task(channel.start())
+    await asyncio.sleep(0.3)
+
+    try:
+        # Fill issued tokens to capacity
+        channel._issued_tokens = {
+            f"nbwt_fill_{i}": time.monotonic() + 300 for i in range(channel._MAX_ISSUED_TOKENS)
+        }
+
+        resp = await _http_get(
+            f"http://127.0.0.1:{port}/auth/token",
+            headers={"Authorization": "Bearer s"},
+        )
+        assert resp.status_code == 429
+        data = resp.json()
+        assert "error" in data
+    finally:
+        await channel.stop()
+        await server_task
+
+
+@pytest.mark.asyncio
+async def test_allow_from_rejects_unauthorized_client_id(bus: MagicMock) -> None:
+    port = 29882
+    channel = _ch(bus, port=port, allowFrom=["alice", "bob"])
+
+    server_task = asyncio.create_task(channel.start())
+    await asyncio.sleep(0.3)
+
+    try:
+        with pytest.raises(websockets.exceptions.InvalidStatus) as exc_info:
+            async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=eve"):
+                pass
+        assert exc_info.value.response.status_code == 403
+    finally:
+        await channel.stop()
+        await server_task
+
+
+@pytest.mark.asyncio
+async def test_client_id_truncation(bus: MagicMock) -> None:
+    port = 29883
+    channel = _ch(bus, port=port)
+
+    server_task = asyncio.create_task(channel.start())
+    await asyncio.sleep(0.3)
+
+    try:
+        long_id = "x" * 200
+        async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id={long_id}") as client:
+            ready = json.loads(await client.recv())
+            assert ready["client_id"] == "x" * 128
+            assert len(ready["client_id"]) == 128
+    finally:
+        await channel.stop()
+        await server_task
+
+
+@pytest.mark.asyncio
+async def test_non_utf8_binary_frame_ignored(bus: MagicMock) -> None:
+    port = 29884
+    channel = _ch(bus, port=port)
+
+    server_task = asyncio.create_task(channel.start())
+    await asyncio.sleep(0.3)
+
+    try:
+        async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=bin-test") as client:
+            await client.recv()  # consume ready
+            # Send non-UTF-8 bytes
+            await client.send(b"\xff\xfe\xfd")
+            await asyncio.sleep(0.05)
+            # publish_inbound should NOT have been called
+            bus.publish_inbound.assert_not_awaited()
+    finally:
+        await channel.stop()
+        await server_task
+
+
+@pytest.mark.asyncio
+async def test_static_token_accepts_issued_token_as_fallback(bus: MagicMock) -> None:
+    port = 29885
+    channel = _ch(
+        bus, port=port,
+        token="static-secret",
+        tokenIssuePath="/auth/token",
+        tokenIssueSecret="route-secret",
+    )
+
+    server_task = asyncio.create_task(channel.start())
+    await asyncio.sleep(0.3)
+
+    try:
+        # Get an issued token
+        resp = await _http_get(
+            f"http://127.0.0.1:{port}/auth/token",
+            headers={"Authorization": "Bearer route-secret"},
+        )
+        assert resp.status_code == 200
+        issued_token = resp.json()["token"]
+
+        # Connect using issued token (not the static one)
+        async with websockets.connect(f"ws://127.0.0.1:{port}/ws?token={issued_token}&client_id=caller") as client:
+            ready = json.loads(await client.recv())
+            assert ready["event"] == "ready"
+    finally:
+        await channel.stop()
+        await server_task
+
+
+@pytest.mark.asyncio
+async def test_allow_from_empty_list_denies_all(bus: MagicMock) -> None:
+    port = 29886
+    channel = _ch(bus, port=port, allowFrom=[])
+
+    server_task = asyncio.create_task(channel.start())
+    await asyncio.sleep(0.3)
+
+    try:
+        with pytest.raises(websockets.exceptions.InvalidStatus) as exc_info:
+            async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=anyone"):
+                pass
+        assert exc_info.value.response.status_code == 403
+    finally:
+        await channel.stop()
+        await server_task
+
+
+@pytest.mark.asyncio
+async def test_websocket_requires_token_without_issue_path(bus: MagicMock) -> None:
+    """When websocket_requires_token is True but no token or issue path configured, all connections are rejected."""
+    port = 29887
+    channel = _ch(bus, port=port, websocketRequiresToken=True)
+
+    server_task = asyncio.create_task(channel.start())
+    await asyncio.sleep(0.3)
+
+    try:
+        # No token at all → 401
+        with pytest.raises(websockets.exceptions.InvalidStatus) as exc_info:
+            async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=u"):
+                pass
+        assert exc_info.value.response.status_code == 401
+
+        # Wrong token → 401
+        with pytest.raises(websockets.exceptions.InvalidStatus) as exc_info:
+            async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=u&token=wrong"):
+                pass
+        assert exc_info.value.response.status_code == 401
+    finally:
+        await channel.stop()
+        await server_task