| 1 | import os |
| 2 | import threading |
| 3 | import uuid |
| 4 | from abc import abstractmethod |
| 5 | from dataclasses import dataclass |
| 6 | from pathlib import Path |
| 7 | from typing import Any, Iterable, Union, TYPE_CHECKING |
| 8 | from urllib.parse import urlparse |
| 9 | |
| 10 | import socketio |
| 11 | from flask import Flask, session, request |
| 12 | |
| 13 | from helpers import files, cache |
| 14 | from helpers.print_style import PrintStyle |
| 15 | from helpers.errors import format_error |
| 16 | from helpers.tunnel_origins import get_active_tunnel_origins, origin_key |
| 17 | |
| 18 | if TYPE_CHECKING: |
| 19 | from helpers.ws_manager import WsManager |
| 20 | |
| 21 | |
| 22 | # Shared types and utilities |
| 23 | |
| 24 | from helpers.network import is_loopback_address |
| 25 | |
| 26 | ConnectionIdentity = tuple[str, str] # (namespace, sid) |
| 27 | |
| 28 | |
| 29 | def _ws_debug_enabled() -> bool: |
| 30 | """Check A0_WS_DEBUG env var — lightweight, no heavy imports.""" |
| 31 | value = os.getenv("A0_WS_DEBUG", "").strip().lower() |
| 32 | return value in {"1", "true", "yes", "on"} |
| 33 | |
| 34 | |
| 35 | def ws_debug(message: str) -> None: |
| 36 | """Log *message* via :class:`PrintStyle` when ``A0_WS_DEBUG`` is active.""" |
| 37 | if _ws_debug_enabled(): |
| 38 | PrintStyle.debug(message) |
| 39 | |
| 40 | |
| 41 | class ConnectionNotFoundError(RuntimeError): |
| 42 | """Raised when attempting to emit to a non-existent WebSocket connection.""" |
| 43 | |
| 44 | def __init__(self, sid: str, *, namespace: str | None = None) -> None: |
| 45 | self.sid = sid |
| 46 | self.namespace = namespace |
| 47 | if namespace: |
| 48 | super().__init__(f"Connection not found: namespace={namespace} sid={sid}") |
| 49 | else: |
| 50 | super().__init__(f"Connection not found: {sid}") |
| 51 | |
| 52 | |
| 53 | def _default_port_for_scheme(scheme: str) -> int | None: |
| 54 | if scheme == "http": |
| 55 | return 80 |
| 56 | if scheme == "https": |
| 57 | return 443 |
| 58 | return None |
| 59 | |
| 60 | |
| 61 | def normalize_origin(value: Any) -> str | None: |
| 62 | """Normalize an Origin/Referer header value to scheme://host[:port].""" |
| 63 | if not isinstance(value, str) or not value.strip(): |
| 64 | return None |
| 65 | parsed = urlparse(value.strip()) |
| 66 | if not parsed.scheme or not parsed.hostname: |
| 67 | return None |
| 68 | origin = f"{parsed.scheme}://{parsed.hostname}" |
| 69 | if parsed.port: |
| 70 | origin += f":{parsed.port}" |
| 71 | return origin |
| 72 | |
| 73 | |
| 74 | def _parse_host_header(value: Any) -> tuple[str | None, int | None]: |
| 75 | if not isinstance(value, str) or not value.strip(): |
| 76 | return None, None |
| 77 | parsed = urlparse(f"http://{value.strip()}") |
| 78 | return parsed.hostname, parsed.port |
| 79 | |
| 80 | |
| 81 | def validate_ws_origin(environ: dict[str, Any]) -> tuple[bool, str | None]: |
| 82 | """Validate the browser Origin during the Socket.IO handshake. |
| 83 | |
| 84 | This is the minimum baseline recommended by RFC 6455 (Origin considerations) |
| 85 | and OWASP (CSWSH mitigation): reject cross-origin WebSocket handshakes when |
| 86 | the server is intended for a specific web UI origin. |
| 87 | """ |
| 88 | |
| 89 | raw_origin = environ.get("HTTP_ORIGIN") or environ.get("HTTP_REFERER") |
| 90 | origin = normalize_origin(raw_origin) |
| 91 | if origin is None: |
| 92 | return False, "missing_origin" |
| 93 | |
| 94 | origin_parsed = urlparse(origin) |
| 95 | origin_host = origin_parsed.hostname.lower() if origin_parsed.hostname else None |
| 96 | origin_port = origin_parsed.port or _default_port_for_scheme(origin_parsed.scheme) |
| 97 | if origin_host is None or origin_port is None: |
| 98 | return False, "invalid_origin" |
| 99 | |
| 100 | raw_host = environ.get("HTTP_HOST") |
| 101 | req_host, req_port = _parse_host_header(raw_host) |
| 102 | if not req_host: |
| 103 | req_host = environ.get("SERVER_NAME") |
| 104 | |
| 105 | if req_port is None: |
| 106 | server_port_raw = environ.get("SERVER_PORT") |
| 107 | try: |
| 108 | server_port = int(server_port_raw) if server_port_raw is not None else None |
| 109 | except (TypeError, ValueError): |
| 110 | server_port = None |
| 111 | if server_port is not None and server_port > 0: |
| 112 | req_port = server_port |
| 113 | |
| 114 | if req_host: |
| 115 | req_host = req_host.lower() |
| 116 | if req_port is None: |
| 117 | req_port = origin_port |
| 118 | |
| 119 | forwarded_host_raw = environ.get("HTTP_X_FORWARDED_HOST") |
| 120 | forwarded_host = None |
| 121 | forwarded_port = None |
| 122 | if isinstance(forwarded_host_raw, str) and forwarded_host_raw.strip(): |
| 123 | first = forwarded_host_raw.split(",")[0].strip() |
| 124 | forwarded_host, forwarded_port = _parse_host_header(first) |
| 125 | if forwarded_host: |
| 126 | forwarded_host = forwarded_host.lower() |
| 127 | |
| 128 | forwarded_proto_raw = environ.get("HTTP_X_FORWARDED_PROTO") |
| 129 | forwarded_scheme = None |
| 130 | if isinstance(forwarded_proto_raw, str) and forwarded_proto_raw.strip(): |
| 131 | forwarded_scheme = forwarded_proto_raw.split(",")[0].strip().lower() |
| 132 | forwarded_scheme = forwarded_scheme or origin_parsed.scheme |
| 133 | forwarded_port = ( |
| 134 | forwarded_port |
| 135 | if forwarded_port is not None |
| 136 | else _default_port_for_scheme(forwarded_scheme) or origin_port |
| 137 | ) |
| 138 | |
| 139 | candidates: list[tuple[str, int]] = [] |
| 140 | if req_host: |
| 141 | candidates.append((req_host, int(req_port))) |
| 142 | if forwarded_host: |
| 143 | candidates.append((forwarded_host, int(forwarded_port))) |
| 144 | |
| 145 | if not candidates: |
| 146 | return False, "missing_host" |
| 147 | |
| 148 | for host, port in candidates: |
| 149 | if origin_host == host and origin_port == port: |
| 150 | return True, None |
| 151 | |
| 152 | request_origin_key = (origin_parsed.scheme, origin_host, int(origin_port)) |
| 153 | for active_origin in get_active_tunnel_origins(): |
| 154 | if origin_key(active_origin) == request_origin_key: |
| 155 | return True, None |
| 156 | |
| 157 | if origin_host not in {host for host, _ in candidates}: |
| 158 | return False, "origin_host_mismatch" |
| 159 | return False, "origin_port_mismatch" |
| 160 | |
| 161 | |
| 162 | # Constants |
| 163 | |
| 164 | ThreadLockType = Union[threading.Lock, threading.RLock] |
| 165 | |
| 166 | NAMESPACE = "/ws" |
| 167 | CACHE_AREA = "ws_handlers(api)(plugins)" |
| 168 | # cache.toggle_area(CACHE_AREA, False) # cache off for now |
| 169 | |
| 170 | |
| 171 | @dataclass |
| 172 | class _SecurityContext: |
| 173 | auth_hash: str | None |
| 174 | csrf_token: str | None |
| 175 | client_csrf_token: str | None |
| 176 | csrf_cookie: str | None |
| 177 | remote_addr: str | None |
| 178 | api_key: str | None |
| 179 | |
| 180 | |
| 181 | _ws_contexts: dict[str, _SecurityContext] = {} |
| 182 | _active_handlers: dict[str, dict[str, "WsHandler"]] = {} |
| 183 | _contexts_lock = threading.Lock() |
| 184 | |
| 185 | |
| 186 | class WsHandler: |
| 187 | """Base class for WebSocket handlers loaded from api/ directories. |
| 188 | |
| 189 | Mirrors ApiHandler conventions: declarative security flags, dynamic file- |
| 190 | based loading, and a ``process(event, data, sid)`` entry point. Handlers |
| 191 | are activated per-connection based on the ``auth.handlers`` list sent by the |
| 192 | client during the Socket.IO connect handshake. |
| 193 | """ |
| 194 | |
| 195 | def __init__( |
| 196 | self, |
| 197 | socketio_server: socketio.AsyncServer, |
| 198 | lock: ThreadLockType, |
| 199 | *, |
| 200 | manager: "WsManager | None" = None, |
| 201 | namespace: str = NAMESPACE, |
| 202 | ): |
| 203 | self.socketio = socketio_server |
| 204 | self.lock = lock |
| 205 | self._manager = manager |
| 206 | self._namespace = namespace |
| 207 | |
| 208 | # Properties |
| 209 | |
| 210 | @property |
| 211 | def namespace(self) -> str: |
| 212 | return self._namespace |
| 213 | |
| 214 | @property |
| 215 | def manager(self) -> "WsManager": |
| 216 | if self._manager is None: |
| 217 | raise RuntimeError("WsHandler has no WsManager bound") |
| 218 | return self._manager |
| 219 | |
| 220 | @property |
| 221 | def identifier(self) -> str: |
| 222 | return f"{self.__class__.__module__}.{self.__class__.__name__}" |
| 223 | |
| 224 | def bind_manager( |
| 225 | self, manager: "WsManager", *, namespace: str | None = None |
| 226 | ) -> None: |
| 227 | """Late-bind (or rebind) the manager and optionally the namespace.""" |
| 228 | self._manager = manager |
| 229 | if namespace is not None: |
| 230 | self._namespace = namespace |
| 231 | |
| 232 | # Security flags (mirror ApiHandler) |
| 233 | |
| 234 | @classmethod |
| 235 | def requires_loopback(cls) -> bool: |
| 236 | return False |
| 237 | |
| 238 | @classmethod |
| 239 | def requires_api_key(cls) -> bool: |
| 240 | return False |
| 241 | |
| 242 | @classmethod |
| 243 | def requires_auth(cls) -> bool: |
| 244 | return True |
| 245 | |
| 246 | @classmethod |
| 247 | def requires_csrf(cls) -> bool: |
| 248 | return cls.requires_auth() |
| 249 | |
| 250 | # Lifecycle hooks |
| 251 | |
| 252 | async def on_connect(self, sid: str) -> None: |
| 253 | pass |
| 254 | |
| 255 | async def on_disconnect(self, sid: str) -> None: |
| 256 | pass |
| 257 | |
| 258 | # Event processing |
| 259 | |
| 260 | @abstractmethod |
| 261 | async def process(self, event: str, data: dict, sid: str) -> dict | None: |
| 262 | """Handle an incoming event. |
| 263 | |
| 264 | Return a dict to include in the acknowledgement, or ``None`` for |
| 265 | fire-and-forget semantics. |
| 266 | """ |
| 267 | |
| 268 | # Emit helpers (delegate to WsManager for envelope wrapping) |
| 269 | |
| 270 | async def emit_to( |
| 271 | self, |
| 272 | sid: str, |
| 273 | event: str, |
| 274 | data: dict, |
| 275 | *, |
| 276 | correlation_id: str | None = None, |
| 277 | ) -> None: |
| 278 | await self.manager.emit_to( |
| 279 | self._namespace, sid, event, data, |
| 280 | handler_id=self.identifier, |
| 281 | correlation_id=correlation_id, |
| 282 | ) |
| 283 | |
| 284 | async def broadcast( |
| 285 | self, |
| 286 | event: str, |
| 287 | data: dict, |
| 288 | *, |
| 289 | exclude_sids: str | Iterable[str] | None = None, |
| 290 | correlation_id: str | None = None, |
| 291 | ) -> None: |
| 292 | await self.manager.broadcast( |
| 293 | self._namespace, event, data, |
| 294 | exclude_sids=exclude_sids, |
| 295 | handler_id=self.identifier, |
| 296 | correlation_id=correlation_id, |
| 297 | ) |
| 298 | |
| 299 | # Aggregation helper |
| 300 | |
| 301 | async def dispatch_to_all_sids( |
| 302 | self, |
| 303 | event: str, |
| 304 | data: dict, |
| 305 | *, |
| 306 | correlation_id: str | None = None, |
| 307 | ) -> list[dict[str, Any]]: |
| 308 | """Dispatch *event* to every connected sid's activated handlers and |
| 309 | aggregate the results. |
| 310 | |
| 311 | Returns a list of ``{sid, correlationId, results}`` dicts – one per |
| 312 | connected sid. This mirrors the shape produced by |
| 313 | ``WsManager.route_event_all`` so that existing frontend |
| 314 | assertions remain valid. |
| 315 | """ |
| 316 | cid = correlation_id or uuid.uuid4().hex |
| 317 | |
| 318 | with _contexts_lock: |
| 319 | snapshot = { |
| 320 | sid: dict(handlers) |
| 321 | for sid, handlers in _active_handlers.items() |
| 322 | } |
| 323 | contexts_snapshot = dict(_ws_contexts) |
| 324 | |
| 325 | mgr = self._manager |
| 326 | aggregated: list[dict[str, Any]] = [] |
| 327 | for sid, handlers in snapshot.items(): |
| 328 | ctx = contexts_snapshot.get(sid) |
| 329 | # Skip sids whose security context was removed (concurrent disconnect). |
| 330 | if ctx is None: |
| 331 | continue |
| 332 | security_errors: list[dict[str, Any]] = [] |
| 333 | passing: list[WsHandler] = [] |
| 334 | for _path, instance in handlers.items(): |
| 335 | error = _check_security(type(instance), ctx) |
| 336 | if error is not None: |
| 337 | security_errors.append({ |
| 338 | "handlerId": instance.identifier, |
| 339 | "ok": False, |
| 340 | "correlationId": cid, |
| 341 | "error": error, |
| 342 | }) |
| 343 | continue |
| 344 | passing.append(instance) |
| 345 | |
| 346 | if mgr is not None and passing: |
| 347 | result = await mgr.process_client_event( |
| 348 | self._namespace, event, |
| 349 | dict(data, correlationId=cid), sid, |
| 350 | handlers=passing, |
| 351 | ) |
| 352 | sid_results = security_errors + result.get("results", []) |
| 353 | else: |
| 354 | # Fallback: inline processing |
| 355 | sid_results = list(security_errors) |
| 356 | for _path, instance in handlers.items(): |
| 357 | if instance not in passing: |
| 358 | continue |
| 359 | try: |
| 360 | result = await instance.process( |
| 361 | event, dict(data, correlationId=cid), sid, |
| 362 | ) |
| 363 | if result is not None: |
| 364 | sid_results.append({ |
| 365 | "handlerId": instance.identifier, |
| 366 | "ok": True, |
| 367 | "correlationId": cid, |
| 368 | "data": result, |
| 369 | }) |
| 370 | except Exception as e: |
| 371 | sid_results.append({ |
| 372 | "handlerId": instance.identifier, |
| 373 | "ok": False, |
| 374 | "correlationId": cid, |
| 375 | "error": {"code": "HANDLER_ERROR", "error": str(e)}, |
| 376 | }) |
| 377 | aggregated.append({ |
| 378 | "sid": sid, |
| 379 | "correlationId": cid, |
| 380 | "results": sid_results, |
| 381 | }) |
| 382 | return aggregated |
| 383 | |
| 384 | # Context helper (shared with ApiHandler) |
| 385 | |
| 386 | def use_context(self, ctxid: str, create_if_not_exists: bool = True): |
| 387 | from helpers.context_utils import use_context as _use_context |
| 388 | return _use_context(self.lock, ctxid, create_if_not_exists) |
| 389 | |
| 390 | |
| 391 | # Security check (aligned with api.py decorators) |
| 392 | |
| 393 | def _check_security(handler_cls: type[WsHandler], ctx: _SecurityContext) -> dict[str, Any] | None: |
| 394 | """Return an error payload dict if the check fails, or ``None`` on success.""" |
| 395 | |
| 396 | if handler_cls.requires_loopback(): |
| 397 | if not ctx.remote_addr or not is_loopback_address(ctx.remote_addr): |
| 398 | return {"code": "FORBIDDEN", "error": "Access denied"} |
| 399 | |
| 400 | if handler_cls.requires_auth(): |
| 401 | from helpers import login |
| 402 | user_pass_hash = login.get_credentials_hash() |
| 403 | if user_pass_hash and ctx.auth_hash != user_pass_hash: |
| 404 | return {"code": "AUTH_REQUIRED", "error": "Authentication required"} |
| 405 | |
| 406 | if handler_cls.requires_csrf(): |
| 407 | if not ctx.csrf_token: |
| 408 | return {"code": "CSRF_MISSING", "error": "CSRF token not initialised"} |
| 409 | if not ctx.client_csrf_token or ctx.client_csrf_token != ctx.csrf_token: |
| 410 | return {"code": "CSRF_INVALID", "error": "CSRF token missing or invalid"} |
| 411 | if ctx.csrf_cookie != ctx.csrf_token: |
| 412 | return {"code": "CSRF_COOKIE", "error": "CSRF cookie mismatch"} |
| 413 | |
| 414 | if handler_cls.requires_api_key(): |
| 415 | from helpers.settings import get_settings |
| 416 | valid_key = get_settings().get("mcp_server_token") |
| 417 | if not ctx.api_key or ctx.api_key != valid_key: |
| 418 | return {"code": "API_KEY_REQUIRED", "error": "API key required"} |
| 419 | |
| 420 | return None |
| 421 | |
| 422 | |
| 423 | # Namespace registration |
| 424 | |
| 425 | def register_ws_namespace( |
| 426 | socketio_server: socketio.AsyncServer, |
| 427 | webapp: Flask, |
| 428 | lock: ThreadLockType, |
| 429 | manager: "WsManager | None" = None, |
| 430 | ) -> None: |
| 431 | from helpers.modules import load_classes_from_file |
| 432 | from helpers import plugins, runtime |
| 433 | |
| 434 | def _resolve_handler(path: str) -> type[WsHandler] | None: |
| 435 | handler_cls: type[WsHandler] | None = None |
| 436 | |
| 437 | # Check built-in api/<path>.py |
| 438 | builtin_file = files.get_abs_path(f"api/{path}.py") |
| 439 | if files.is_in_dir(builtin_file, files.get_abs_path("api")) and files.exists(builtin_file): |
| 440 | classes = load_classes_from_file(builtin_file, WsHandler) |
| 441 | if classes: |
| 442 | handler_cls = classes[0] |
| 443 | |
| 444 | # Check user api/<path>.py |
| 445 | if handler_cls is None: |
| 446 | user_file = files.get_abs_path(files.USER_DIR, f"api/{path}.py") |
| 447 | if files.exists(user_file): |
| 448 | classes = load_classes_from_file(user_file, WsHandler) |
| 449 | if classes: |
| 450 | handler_cls = classes[0] |
| 451 | |
| 452 | # Check plugin api/<handler>.py — path format: plugins/<plugin_name>/<handler> |
| 453 | if handler_cls is None and path.startswith("plugins/"): |
| 454 | parts = path.split("/", 2) |
| 455 | if len(parts) == 3: |
| 456 | _, plugin_name, handler_name = parts |
| 457 | plugin_dir = plugins.find_plugin_dir(plugin_name) |
| 458 | if plugin_dir: |
| 459 | plugin_file = Path(plugin_dir) / "api" / f"{handler_name}.py" |
| 460 | if plugin_file.is_file(): |
| 461 | classes = load_classes_from_file(str(plugin_file), WsHandler) |
| 462 | if classes: |
| 463 | handler_cls = classes[0] |
| 464 | |
| 465 | return handler_cls |
| 466 | |
| 467 | def _resolve_cached(path: str) -> type[WsHandler] | None: |
| 468 | cached = cache.get(CACHE_AREA, path) |
| 469 | if cached is not None: |
| 470 | return cached |
| 471 | handler_cls = _resolve_handler(path) |
| 472 | if handler_cls is not None: |
| 473 | cache.add(CACHE_AREA, path, handler_cls) |
| 474 | return handler_cls |
| 475 | |
| 476 | @socketio_server.on("connect", namespace=NAMESPACE) # type: ignore |
| 477 | async def _on_connect(sid, environ, auth): |
| 478 | with webapp.request_context(environ): |
| 479 | origin_ok, origin_reason = validate_ws_origin(environ) |
| 480 | if not origin_ok: |
| 481 | PrintStyle.warning( |
| 482 | f"WS connect rejected for {sid}: {origin_reason or 'invalid'}" |
| 483 | ) |
| 484 | return False |
| 485 | |
| 486 | ctx = _SecurityContext( |
| 487 | auth_hash=session.get("authentication"), |
| 488 | csrf_token=session.get("csrf_token"), |
| 489 | client_csrf_token=( |
| 490 | (auth.get("csrf_token") or auth.get("csrfToken")) |
| 491 | if isinstance(auth, dict) else None |
| 492 | ), |
| 493 | csrf_cookie=request.cookies.get( |
| 494 | f"csrf_token_{runtime.get_runtime_id()}" |
| 495 | ), |
| 496 | remote_addr=str(request.remote_addr) if request.remote_addr else None, |
| 497 | api_key=( |
| 498 | (auth.get("api_key") or auth.get("apiKey")) |
| 499 | if isinstance(auth, dict) else None |
| 500 | ), |
| 501 | ) |
| 502 | user_id = session.get("user_id") or "single_user" |
| 503 | |
| 504 | with _contexts_lock: |
| 505 | _ws_contexts[sid] = ctx |
| 506 | |
| 507 | # Register with WsManager first so that the dispatcher loop and |
| 508 | # connection tracking are available before handler on_connect runs |
| 509 | # (extensions like StateSync depend on manager._dispatcher_loop). |
| 510 | if manager is not None: |
| 511 | await manager.handle_connect(NAMESPACE, sid, user_id=user_id) |
| 512 | |
| 513 | # Activate handlers declared in auth.handlers |
| 514 | handler_paths: list[str] = [] |
| 515 | if isinstance(auth, dict): |
| 516 | raw = auth.get("handlers") |
| 517 | if isinstance(raw, list): |
| 518 | handler_paths = [p for p in raw if isinstance(p, str)] |
| 519 | |
| 520 | activated: dict[str, WsHandler] = {} |
| 521 | for path in handler_paths: |
| 522 | try: |
| 523 | handler_cls = _resolve_cached(path) |
| 524 | if handler_cls is None: |
| 525 | continue |
| 526 | error = _check_security(handler_cls, ctx) |
| 527 | if error is not None: |
| 528 | continue |
| 529 | instance = handler_cls( |
| 530 | socketio_server, lock, |
| 531 | manager=manager, namespace=NAMESPACE, |
| 532 | ) |
| 533 | await instance.on_connect(sid) |
| 534 | activated[path] = instance |
| 535 | except Exception as e: |
| 536 | PrintStyle.error(f"WS on_connect error ({path}): {format_error(e)}") |
| 537 | |
| 538 | with _contexts_lock: |
| 539 | _active_handlers[sid] = activated |
| 540 | |
| 541 | return True |
| 542 | |
| 543 | @socketio_server.on("disconnect", namespace=NAMESPACE) # type: ignore |
| 544 | async def _on_disconnect(sid, reason=None): |
| 545 | with _contexts_lock: |
| 546 | activated = _active_handlers.pop(sid, {}) |
| 547 | _ws_contexts.pop(sid, None) |
| 548 | |
| 549 | for path, instance in activated.items(): |
| 550 | try: |
| 551 | await instance.on_disconnect(sid) |
| 552 | except Exception as e: |
| 553 | PrintStyle.error(f"WS on_disconnect error ({path}): {format_error(e)}") |
| 554 | |
| 555 | if manager is not None: |
| 556 | await manager.handle_disconnect(NAMESPACE, sid) |
| 557 | |
| 558 | @socketio_server.on("*", namespace=NAMESPACE) # type: ignore |
| 559 | async def _dispatch(event, sid, data): |
| 560 | incoming = data if isinstance(data, dict) else {} |
| 561 | |
| 562 | try: |
| 563 | with _contexts_lock: |
| 564 | ctx = _ws_contexts.get(sid) |
| 565 | activated = dict(_active_handlers.get(sid, {})) |
| 566 | |
| 567 | correlation_id = incoming.get("correlationId") or uuid.uuid4().hex |
| 568 | |
| 569 | if ctx is None: |
| 570 | return _error_response("AUTH_REQUIRED", |
| 571 | "No security context", correlation_id) |
| 572 | if not activated: |
| 573 | return _error_response("NO_HANDLERS", |
| 574 | "No handlers activated", correlation_id) |
| 575 | |
| 576 | # Pre-filter handlers through security checks |
| 577 | passing_handlers: list[WsHandler] = [] |
| 578 | security_errors: list[dict[str, Any]] = [] |
| 579 | for path, instance in activated.items(): |
| 580 | error = _check_security(type(instance), ctx) |
| 581 | if error is not None: |
| 582 | security_errors.append({ |
| 583 | "handlerId": instance.identifier, |
| 584 | "ok": False, |
| 585 | "correlationId": correlation_id, |
| 586 | "error": error, |
| 587 | }) |
| 588 | else: |
| 589 | passing_handlers.append(instance) |
| 590 | |
| 591 | # Delegate to WsManager for unified processing pipeline |
| 592 | # (worker thread isolation, diagnostic events, WsResult support) |
| 593 | if manager is not None and passing_handlers: |
| 594 | result = await manager.process_client_event( |
| 595 | NAMESPACE, event, incoming, sid, |
| 596 | handlers=passing_handlers, |
| 597 | ) |
| 598 | if security_errors: |
| 599 | result["results"] = security_errors + result.get("results", []) |
| 600 | return result |
| 601 | |
| 602 | # All handlers failed security or no manager — return collected errors |
| 603 | if not passing_handlers: |
| 604 | return {"correlationId": correlation_id, "results": security_errors} |
| 605 | |
| 606 | # Fallback: inline processing (no manager — should not happen in practice) |
| 607 | handler_payload: dict[str, Any] |
| 608 | if "data" in incoming and isinstance(incoming.get("data"), dict): |
| 609 | handler_payload = dict(incoming["data"]) |
| 610 | else: |
| 611 | handler_payload = dict(incoming) |
| 612 | handler_payload["correlationId"] = correlation_id |
| 613 | |
| 614 | results: list[dict[str, Any]] = list(security_errors) |
| 615 | for path, instance in activated.items(): |
| 616 | if instance not in passing_handlers: |
| 617 | continue |
| 618 | try: |
| 619 | result = await instance.process(event, handler_payload, sid) |
| 620 | if result is not None: |
| 621 | results.append({ |
| 622 | "handlerId": instance.identifier, |
| 623 | "ok": True, |
| 624 | "correlationId": correlation_id, |
| 625 | "data": result, |
| 626 | }) |
| 627 | except Exception as e: |
| 628 | error_text = format_error(e) |
| 629 | PrintStyle.error(f"WS handler error ({path}/{event}): {error_text}") |
| 630 | results.append({ |
| 631 | "handlerId": instance.identifier, |
| 632 | "ok": False, |
| 633 | "correlationId": correlation_id, |
| 634 | "error": {"code": "HANDLER_ERROR", "error": "Internal server error"}, |
| 635 | }) |
| 636 | |
| 637 | return {"correlationId": correlation_id, "results": results} |
| 638 | |
| 639 | except Exception as e: |
| 640 | error_text = format_error(e) |
| 641 | PrintStyle.error(f"WS dispatch error ({event}): {error_text}") |
| 642 | return _error_response( |
| 643 | "INTERNAL_ERROR", "Internal server error", |
| 644 | incoming.get("correlationId", ""), |
| 645 | ) |
| 646 | |
| 647 | |
| 648 | def _error_response(code: str, message: str, |
| 649 | correlation_id: str) -> dict[str, Any]: |
| 650 | return { |
| 651 | "correlationId": correlation_id, |
| 652 | "results": [{ |
| 653 | "handlerId": "ws.dispatch", |
| 654 | "ok": False, |
| 655 | "error": {"code": code, "error": message}, |
| 656 | }], |
| 657 | } |