Coverage for src/local_deep_research/web/services/socket_service.py: 96%
231 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-19 23:35 +0000
1from threading import Lock
2from typing import Any
4from flask import Flask, request, session
5from flask_socketio import SocketIO, join_room
6from loguru import logger
8from ...constants import ResearchStatus
9from ...database.encrypted_db import db_manager
10from ...database.session_passwords import session_password_store
11from ..routes.globals import get_active_research_snapshot
14def _install_origin_rejection_logging(socketio: SocketIO) -> bool:
15 """Re-emit engine.io's silenced WebSocket origin rejections via loguru.
17 engine.io validates the Origin at handshake and calls
18 ``_log_error_once('<origin> is not an accepted origin.', 'bad-origin')``,
19 but the server runs with ``logger=False`` so that message never surfaces —
20 the only symptom of a misconfigured WebSocket origin is a frozen progress
21 UI. Wrap that one call to log a WARNING (deduped per origin) pointing at the
22 fix. An Origin is a scheme+host, not PII. Best-effort: a no-op (returns
23 False) if engine.io internals change, so it can never break startup.
25 The dedup set is capped: the handshake is pre-auth and ``Origin`` is
26 attacker-controlled, so an unbounded set would be a memory-growth + log-
27 amplification vector. After ``cap`` distinct origins we stop tracking/warning
28 (an operator has more than enough signal by then).
29 """
30 try:
31 eio = socketio.server.eio
32 original = eio._log_error_once
33 except AttributeError:
34 logger.debug(
35 "Socket.IO: origin-rejection logging not installed "
36 "(engine.io internals changed); handshake rejections stay silent"
37 )
38 return False
40 warned: set[str] = set()
41 cap = 100
43 def _log_error_once(message, message_key):
44 if (
45 message_key == "bad-origin"
46 and len(warned) < cap
47 and message not in warned
48 ):
49 warned.add(message)
50 logger.warning(
51 f"Socket.IO rejected a WebSocket handshake: {message} Set "
52 "LDR_SECURITY_WEBSOCKET_ALLOWED_ORIGINS to this origin if it is "
53 "your front-end; behind a TLS-terminating proxy, also forward "
54 "X-Forwarded-Proto so the same-origin check sees https."
55 )
56 return original(message, message_key)
58 eio._log_error_once = _log_error_once
59 return True
62class SocketIOService:
63 """
64 Singleton class for managing SocketIO connections and subscriptions.
65 """
67 _instance = None
69 def __new__(cls, *args: Any, app: Flask | None = None, **kwargs: Any):
70 """
71 Args:
72 app: The Flask app to bind this service to. It must be specified
73 the first time this is called and the singleton instance is
74 created, but will be ignored after that.
75 *args: Arguments to pass to the superclass's __new__ method.
76 **kwargs: Keyword arguments to pass to the superclass's __new__ method.
77 """
78 if not cls._instance:
79 if app is None:
80 raise ValueError(
81 "Flask app must be specified to create a SocketIOService instance."
82 )
83 cls._instance = super(SocketIOService, cls).__new__(
84 cls, *args, **kwargs
85 )
86 cls._instance.__init_singleton(app)
87 return cls._instance
89 def __init_singleton(self, app: Flask) -> None:
90 """
91 Initializes the singleton instance.
93 Args:
94 app: The app to bind this service to.
96 """
97 self.__app = app # Store the Flask app reference
99 # Determine WebSocket CORS policy from env var or default
100 from ...settings.env_registry import get_env_setting
102 ws_origins_env = get_env_setting("security.websocket.allowed_origins")
103 socketio_cors: str | list[str] | None
104 if ws_origins_env is not None:
105 if ws_origins_env == "*":
106 socketio_cors = "*"
107 elif ws_origins_env:
108 socketio_cors = [o.strip() for o in ws_origins_env.split(",")]
109 else:
110 socketio_cors = None
111 else:
112 # No env var set — fail closed to same-origin only, matching HTTP CORS default
113 socketio_cors = None
115 if socketio_cors is None:
116 logger.info(
117 "Socket.IO CORS: same-origin only (set LDR_SECURITY_WEBSOCKET_ALLOWED_ORIGINS to configure)"
118 )
119 elif socketio_cors == "*":
120 logger.debug("Socket.IO CORS: all origins allowed")
121 else:
122 logger.info(f"Socket.IO CORS: restricted to {socketio_cors}")
124 self.__socketio = SocketIO(
125 app,
126 cors_allowed_origins=socketio_cors,
127 async_mode="threading",
128 path="/socket.io",
129 logger=False,
130 engineio_logger=False,
131 ping_timeout=20,
132 ping_interval=5,
133 )
135 # Make a rejected WebSocket origin diagnosable (otherwise it is a silent
136 # frozen progress UI). Skipped for the allow-all case, which rejects
137 # nothing.
138 if socketio_cors != "*":
139 _install_origin_rejection_logging(self.__socketio)
141 # Socket subscription tracking.
142 self.__socket_subscriptions: dict[str, Any] = {}
143 # Set to false to disable logging in the event handlers. This can
144 # be necessary because it will sometimes run the handlers directly
145 # during a call to `emit` that was made in a logging handler.
146 self.__logging_enabled = True
147 # Protects access to shared state.
148 self.__lock = Lock()
150 # Register events.
151 @self.__socketio.on("connect")
152 def on_connect():
153 return self.__handle_connect(request)
155 @self.__socketio.on("disconnect")
156 def on_disconnect(reason: str):
157 self.__handle_disconnect(request, reason)
159 @self.__socketio.on("subscribe_to_research")
160 def on_subscribe(data):
161 self.__handle_subscribe(data, request)
163 # Backwards-compatible alias: the JS client emits 'join' on subscribe.
164 # Without this, the catch-up snapshot in __handle_subscribe never
165 # fires and per-client targeting falls through to broadcast.
166 @self.__socketio.on("join")
167 def on_join(data):
168 self.__handle_subscribe(data, request)
170 @self.__socketio.on("leave")
171 def on_leave(data):
172 self.__handle_unsubscribe(data, request)
174 @self.__socketio.on("unsubscribe_from_research")
175 def on_unsubscribe(data):
176 self.__handle_unsubscribe(data, request)
178 @self.__socketio.on_error
179 def on_error(e):
180 return self.__handle_socket_error(e)
182 @self.__socketio.on_error_default
183 def on_default_error(e):
184 return self.__handle_default_error(e)
186 def __log_info(self, message: str, *args: Any, **kwargs: Any) -> None:
187 """Log an info message."""
188 if self.__logging_enabled:
189 logger.info(message, *args, **kwargs)
191 def __log_error(self, message: str, *args: Any, **kwargs: Any) -> None:
192 """Log an error message."""
193 if self.__logging_enabled:
194 logger.error(message, *args, **kwargs)
196 def __log_exception(self, message: str, *args: Any, **kwargs: Any) -> None:
197 """Log an exception."""
198 if self.__logging_enabled:
199 logger.exception(message, *args, **kwargs)
201 @staticmethod
202 def user_room(username: str) -> str:
203 """Socket.IO room name that every one of a user's connected tabs joins.
205 Used to scope user-private events to a single account. Kept here so the
206 connect handler and event emitters share one definition and cannot
207 drift apart.
208 """
209 return f"user:{username}"
211 def emit_socket_event(self, event, data, room=None):
212 """
213 Emit a socket event to clients.
215 Args:
216 event: The event name to emit
217 data: The data to send with the event
218 room: Optional room ID to send to specific client
220 Returns:
221 bool: True if emission was successful, False otherwise
222 """
223 try:
224 # If room is specified, only emit to that room
225 if room:
226 self.__socketio.emit(event, data, room=room)
227 else:
228 # Otherwise broadcast to all
229 self.__socketio.emit(event, data)
230 return True
231 except Exception:
232 logger.exception(f"Error emitting socket event {event}")
233 return False
235 def emit_to_subscribers(
236 self, event_base, research_id, data, enable_logging: bool = True
237 ):
238 """
239 Emit an event to all subscribers of a specific research.
241 Args:
242 event_base: Base event name (will be formatted with research_id)
243 research_id: ID of the research
244 data: The data to send with the event
245 enable_logging: If set to false, this will disable all logging,
246 which is useful if we are calling this inside of a logging
247 handler.
249 Returns:
250 bool: True if emission was successful, False otherwise
252 """
253 if not enable_logging:
254 self.__logging_enabled = False
256 try:
257 full_event = f"{event_base}_{research_id}"
259 # Emit only to specific subscribers (no broadcast) to avoid
260 # duplicate messages and reduce server load under concurrency
261 with self.__lock:
262 subscriptions = self.__socket_subscriptions.get(research_id)
263 if subscriptions:
264 subscriptions = (
265 subscriptions.copy()
266 ) # snapshot avoids RuntimeError
267 else:
268 subscriptions = None
269 if subscriptions is not None:
270 for sid in subscriptions:
271 try:
272 self.__socketio.emit(full_event, data, room=sid)
273 except Exception:
274 self.__log_exception(
275 f"Error emitting to subscriber {sid}"
276 )
277 # When no targeted subscribers exist yet, drop the event.
278 # The catch-up snapshot in __handle_subscribe replays the
279 # latest progress on subscribe, so early-arriving events
280 # are recovered correctly without a cross-user broadcast.
282 return True
283 except Exception:
284 self.__log_exception(
285 f"Error emitting to subscribers for research {research_id}"
286 )
287 return False
288 finally:
289 self.__logging_enabled = True
291 def remove_subscriptions_for_research(self, research_id: str) -> None:
292 """Remove all socket subscriptions for a completed research."""
293 with self.__lock:
294 removed = self.__socket_subscriptions.pop(research_id, None)
295 if removed is not None:
296 self.__log_info(
297 f"Removed {len(removed)} subscription(s) for research {research_id}"
298 )
300 def __handle_connect(self, request):
301 """Handle client connection"""
302 username = session.get("username")
303 if not username:
304 self.__log_info(
305 f"Rejected unauthenticated WebSocket connection from {request.sid}"
306 )
307 return False
308 if not db_manager.is_user_connected(username):
309 # Cookie is valid but the per-user DB engine isn't open yet (race vs first
310 # XHR after page load, gunicorn worker restart, or idle eviction). Lazily
311 # open it using the password the user authenticated with at login.
312 session_id = session.get("session_id")
313 password = (
314 session_password_store.get_session_password(
315 username, session_id
316 )
317 if session_id
318 else None
319 )
320 if not password:
321 self.__log_info(
322 f"Rejected WebSocket connection for {username}: no active DB session and no stored password"
323 )
324 return False
325 try:
326 db_manager.open_user_database(username, password)
327 except Exception as e:
328 # Use __log_error (not __log_exception) so loguru cannot include
329 # the `password` local in a diagnose=True traceback.
330 self.__log_error(
331 f"Lazy DB open failed for {username} at WebSocket connect: {type(e).__name__}"
332 )
333 return False
334 # Join a per-user room so user-scoped events (e.g. settings_changed,
335 # which carries raw setting values including plaintext API keys) reach
336 # only this user's own browser tabs and are never broadcast to every
337 # connected client. Flask-SocketIO auto-removes the socket from the
338 # room on disconnect.
339 join_room(self.user_room(username))
340 self.__log_info(f"Client connected: {request.sid} (user: {username})")
341 return True
343 def __handle_disconnect(self, request, reason: str):
344 """Handle client disconnection"""
345 try:
346 self.__log_info(
347 f"Client {request.sid} disconnected because: {reason}"
348 )
349 # Clean up subscriptions for this client.
350 # __socket_subscriptions is keyed by research_id → set of sids,
351 # so we iterate all entries and discard the disconnecting sid.
352 with self.__lock:
353 empty_keys = []
354 for research_id, sids in self.__socket_subscriptions.items():
355 sids.discard(request.sid)
356 if not sids:
357 empty_keys.append(research_id)
358 for key in empty_keys:
359 del self.__socket_subscriptions[key]
360 self.__log_info(f"Removed subscription for client {request.sid}")
362 # Clean up any thread-local database sessions that may have been
363 # created during socket handler execution. This prevents file
364 # descriptor leaks from unclosed SQLAlchemy sessions.
365 try:
366 from ...database.thread_local_session import (
367 cleanup_current_thread,
368 )
370 cleanup_current_thread()
371 except ImportError:
372 pass # Module not available, skip cleanup
373 except Exception:
374 self.__log_exception(
375 "Error cleaning up thread session on disconnect"
376 )
377 except Exception as e:
378 self.__log_exception(f"Error handling disconnect: {e}")
380 def __handle_subscribe(self, data, request):
381 """Handle client subscription to research updates."""
382 research_id = data.get("research_id")
383 if not research_id:
384 return
386 # Verify the connected user actually owns this research before
387 # subscribing. The in-memory `_active_research` snapshot is keyed
388 # only by research_id (no user tuple), so without this guard any
389 # logged-in user could subscribe to any guessed/leaked research
390 # UUID and receive its progress events. The per-user encrypted DB
391 # is the ownership boundary: if the research row doesn't exist in
392 # the user's DB, they don't own it.
393 username = session.get("username")
394 if not username or not self._user_owns_research(username, research_id): 394 ↛ 395line 394 didn't jump to line 395 because the condition on line 394 was never true
395 self.__log_info(
396 f"Rejected subscribe from {request.sid}: user does not own research {research_id}"
397 )
398 return
400 with self.__lock:
401 if research_id not in self.__socket_subscriptions:
402 self.__socket_subscriptions[research_id] = set()
403 self.__socket_subscriptions[research_id].add(request.sid)
404 self.__log_info(
405 f"Client {request.sid} subscribed to research {research_id}"
406 )
408 # Send current status immediately if available in active research
409 snapshot = get_active_research_snapshot(research_id)
410 if snapshot is not None:
411 progress = snapshot["progress"]
412 latest_log = snapshot["log"][-1] if snapshot["log"] else None
414 if latest_log:
415 self.emit_socket_event(
416 f"progress_{research_id}",
417 {
418 "progress": progress,
419 "message": latest_log.get("message", "Processing..."),
420 "status": ResearchStatus.IN_PROGRESS,
421 "log_entry": latest_log,
422 },
423 room=request.sid,
424 )
426 @staticmethod
427 def _user_owns_research(username: str, research_id: str) -> bool:
428 """Return True if the given user owns this research / benchmark id.
430 Used as the authorization boundary for WebSocket subscriptions —
431 ownership is checked against the user's encrypted SQLite database,
432 which is the per-user data partition. A static helper so unit
433 tests can exercise the authz logic without standing up the
434 singleton/Flask app.
436 Recognizes both normal research (``ResearchHistory``, UUID id) and
437 benchmark runs (``BenchmarkRun``, integer id) — the benchmark page
438 subscribes with its ``BenchmarkRun.id``, which lives in the same
439 per-user DB. Both checks stay scoped to the caller's own database,
440 so no cross-user access is introduced.
441 """
442 try:
443 from ...database.session_context import get_user_db_session
444 from ...database.models import ResearchHistory
446 with get_user_db_session(username) as db:
447 if (
448 db.query(ResearchHistory.id)
449 .filter(ResearchHistory.id == research_id)
450 .first()
451 is not None
452 ):
453 return True
455 # Benchmark pages subscribe with their BenchmarkRun.id.
456 # Recognize the user's own benchmark runs so the ownership
457 # gate doesn't drop benchmark live progress (regression vs.
458 # the removed cross-user broadcast). research_id stays a
459 # string (never coerced to int — IDs are strings/UUIDs
460 # repo-wide); SQLite applies numeric affinity to match the
461 # Integer column. Only attempt this for numeric ids.
462 if str(research_id).isdigit():
463 from ...database.models.benchmark import BenchmarkRun
465 return (
466 db.query(BenchmarkRun.id)
467 .filter(BenchmarkRun.id == research_id)
468 .first()
469 is not None
470 )
471 return False
472 except Exception:
473 # Conservative: deny on any DB-open or query failure so a
474 # transient infra error never silently widens authz.
475 logger.opt(exception=True).warning(
476 "Failed to verify research ownership for socket subscribe"
477 )
478 return False
480 def __handle_unsubscribe(self, data, request):
481 """Handle client unsubscribe from research updates."""
482 research_id = (
483 data.get("research_id") if isinstance(data, dict) else None
484 )
485 if not research_id:
486 return
488 # Symmetric with __handle_subscribe: require the caller to own the
489 # research before mutating the per-research subscription set. The
490 # practical impact of an unguarded unsubscribe is small (no data
491 # exfiltration; subscribe is already guarded), but it keeps the
492 # authz boundary consistent and avoids log spam from spoofed sids.
493 username = session.get("username")
494 if not username or not self._user_owns_research(username, research_id):
495 self.__log_info(
496 f"Rejected unsubscribe from {request.sid}: user does not own research {research_id}"
497 )
498 return
500 with self.__lock:
501 subs = self.__socket_subscriptions.get(research_id)
502 if subs:
503 subs.discard(request.sid)
504 # Prune empty sets so the dict doesn't grow unbounded with
505 # stale research_ids over long server runtimes.
506 if not subs:
507 self.__socket_subscriptions.pop(research_id, None)
508 self.__log_info(
509 f"Client {request.sid} unsubscribed from research {research_id}"
510 )
512 def __handle_socket_error(self, e):
513 """Handle Socket.IO errors"""
514 self.__log_exception(f"Socket.IO error: {str(e)}")
515 # Don't propagate exceptions to avoid crashing the server
516 return False
518 def __handle_default_error(self, e):
519 """Handle unhandled Socket.IO errors"""
520 self.__log_exception(f"Unhandled Socket.IO error: {str(e)}")
521 # Don't propagate exceptions to avoid crashing the server
522 return False
524 def run(self, host: str, port: int, debug: bool = False) -> None:
525 """
526 Runs the SocketIO server.
528 Args:
529 host: The hostname to bind the server to.
530 port: The port number to listen on.
531 debug: Whether to run in debug mode. Defaults to False.
533 """
534 # Suppress Server header to prevent version information disclosure
535 # This must be done before starting the server because Werkzeug adds
536 # the header at the HTTP layer, not WSGI layer
537 try:
538 from werkzeug.serving import WSGIRequestHandler
540 WSGIRequestHandler.version_string = lambda self: "" # type: ignore[method-assign]
541 logger.debug("Suppressed Server header for security")
542 except ImportError:
543 logger.warning(
544 "Could not suppress Server header - werkzeug not found"
545 )
547 logger.info(f"Starting web server on {host}:{port} (debug: {debug})")
548 self.__socketio.run(
549 self.__app, # Use the stored Flask app reference
550 debug=debug,
551 host=host,
552 port=port,
553 allow_unsafe_werkzeug=True,
554 use_reloader=False,
555 )