chore: import upstream snapshot with attribution

Server/tests/integration/test_plugin_hub_websocket_auth.py

@@ -0,0 +1,183 @@
+"""Tests for PluginHub WebSocket API key authentication gate."""
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from core.config import config
+from core.constants import API_KEY_HEADER
+from services.api_key_service import ApiKeyService, ValidationResult
+from transport.plugin_hub import PluginHub
+from transport.plugin_registry import PluginRegistry
+
+
+@pytest.fixture(autouse=True)
+def _reset_api_key_singleton():
+    ApiKeyService._instance = None
+    yield
+    ApiKeyService._instance = None
+
+
+@pytest.fixture(autouse=True)
+def _reset_plugin_hub():
+    """Ensure PluginHub class-level state doesn't leak between tests."""
+    old_registry = PluginHub._registry
+    old_connections = PluginHub._connections.copy()
+    old_pending = PluginHub._pending.copy()
+    old_lock = PluginHub._lock
+    old_loop = PluginHub._loop
+
+    yield
+
+    PluginHub._registry = old_registry
+    PluginHub._connections = old_connections
+    PluginHub._pending = old_pending
+    PluginHub._lock = old_lock
+    PluginHub._loop = old_loop
+
+
+def _make_mock_websocket(headers=None, state_attrs=None):
+    """Create a mock WebSocket with configurable headers and state."""
+    ws = AsyncMock()
+    ws.headers = headers or {}
+    ws.state = SimpleNamespace(**(state_attrs or {}))
+    ws.accept = AsyncMock()
+    ws.close = AsyncMock()
+    ws.send_json = AsyncMock()
+    return ws
+
+
+def _make_hub():
+    """Create a PluginHub instance with a minimal ASGI scope."""
+    scope = {"type": "websocket"}
+    return PluginHub(scope, receive=AsyncMock(), send=AsyncMock())
+
+
+def _init_api_key_service(validate_result=None):
+    """Initialize ApiKeyService with a mocked validate method."""
+    svc = ApiKeyService(validation_url="https://auth.example.com/validate")
+    if validate_result is not None:
+        svc.validate = AsyncMock(return_value=validate_result)
+    return svc
+
+
+class TestWebSocketAuthGate:
+    @pytest.mark.asyncio
+    async def test_no_api_key_remote_hosted_rejected(self, monkeypatch):
+        """WebSocket without API key in remote-hosted mode -> close 4401."""
+        monkeypatch.setattr(config, "http_remote_hosted", True)
+        _init_api_key_service(ValidationResult(valid=True, user_id="u1"))
+
+        ws = _make_mock_websocket(headers={})  # No X-API-Key header
+        hub = _make_hub()
+
+        await hub.on_connect(ws)
+
+        ws.close.assert_called_once_with(code=4401, reason="API key required")
+        ws.accept.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_invalid_api_key_rejected(self, monkeypatch):
+        """WebSocket with invalid API key -> close 4403."""
+        monkeypatch.setattr(config, "http_remote_hosted", True)
+        _init_api_key_service(ValidationResult(
+            valid=False, error="Invalid API key"))
+
+        ws = _make_mock_websocket(headers={API_KEY_HEADER: "sk-bad-key"})
+        hub = _make_hub()
+
+        await hub.on_connect(ws)
+
+        ws.close.assert_called_once_with(code=4403, reason="Invalid API key")
+        ws.accept.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_valid_api_key_accepted(self, monkeypatch):
+        """WebSocket with valid API key -> accepted, user_id stored in state."""
+        monkeypatch.setattr(config, "http_remote_hosted", True)
+        _init_api_key_service(
+            ValidationResult(valid=True, user_id="user-42",
+                             metadata={"plan": "pro"})
+        )
+
+        ws = _make_mock_websocket(headers={API_KEY_HEADER: "sk-valid-key"})
+        hub = _make_hub()
+
+        await hub.on_connect(ws)
+
+        ws.accept.assert_called_once()
+        ws.close.assert_not_called()
+        assert ws.state.user_id == "user-42"
+        assert ws.state.api_key_metadata == {"plan": "pro"}
+        # Should have sent welcome message
+        ws.send_json.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_auth_service_unavailable_close_1013(self, monkeypatch):
+        """Auth service error with 'unavailable' -> close 1013 (try again later)."""
+        monkeypatch.setattr(config, "http_remote_hosted", True)
+        _init_api_key_service(
+            ValidationResult(
+                valid=False, error="Auth service unavailable", cacheable=False)
+        )
+
+        ws = _make_mock_websocket(headers={API_KEY_HEADER: "sk-some-key"})
+        hub = _make_hub()
+
+        await hub.on_connect(ws)
+
+        ws.close.assert_called_once_with(code=1013, reason="Try again later")
+        ws.accept.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_not_remote_hosted_accepts_without_key(self, monkeypatch):
+        """When not remote-hosted, WebSocket accepted without API key."""
+        monkeypatch.setattr(config, "http_remote_hosted", False)
+
+        ws = _make_mock_websocket(headers={})
+        hub = _make_hub()
+
+        await hub.on_connect(ws)
+
+        ws.accept.assert_called_once()
+        ws.close.assert_not_called()
+
+
+class TestUserIdFlowsToRegistration:
+    @pytest.mark.asyncio
+    async def test_user_id_passed_to_registry_on_register(self, monkeypatch):
+        """After valid auth, the register message should pass user_id to registry."""
+        monkeypatch.setattr(config, "http_remote_hosted", True)
+        _init_api_key_service(
+            ValidationResult(valid=True, user_id="user-99")
+        )
+
+        registry = PluginRegistry()
+        loop = asyncio.get_running_loop()
+        PluginHub.configure(registry, loop)
+
+        # Simulate full flow: connect, then register
+        ws = _make_mock_websocket(headers={API_KEY_HEADER: "sk-valid-key"})
+        hub = _make_hub()
+
+        await hub.on_connect(ws)
+        assert ws.state.user_id == "user-99"
+
+        # Simulate register message
+        register_data = {
+            "type": "register",
+            "project_name": "TestProject",
+            "project_hash": "abc123",
+            "unity_version": "2022.3",
+        }
+        await hub.on_receive(ws, register_data)
+
+        # Verify registry has the user_id
+        sessions = await registry.list_sessions(user_id="user-99")
+        assert len(sessions) == 1
+        session = next(iter(sessions.values()))
+        assert session.user_id == "user-99"
+        assert session.project_name == "TestProject"
+        assert session.project_hash == "abc123"