mdb-engine 0.5.1__py3-none-any.whl → 0.7.0__py3-none-any.whl

mdb_engine/routing/websockets.py

@@ -270,6 +270,31 @@ _manager_lock = asyncio.Lock()
 # Note: Registration happens synchronously during app startup, so no lock needed
 _message_handlers: dict[str, dict[str, Callable[[Any, dict[str, Any]], Awaitable[None]]]] = {}
 
+# Global WebSocket ticket store (shared across all apps in multi-app setups)
+# This is set when the engine initializes and is used for ticket authentication
+_global_websocket_ticket_store: Any | None = None
+
+
+def set_global_websocket_ticket_store(ticket_store: Any) -> None:
+    """
+    Set the global WebSocket ticket store.
+
+    This is called by the engine during initialization to make the ticket store
+    accessible to all WebSocket endpoints, even when routes are registered via routers.
+
+    Args:
+        ticket_store: WebSocketTicketStore instance
+    """
+    global _global_websocket_ticket_store
+    _global_websocket_ticket_store = ticket_store
+    if ticket_store:
+        logger.info(
+            f"✅ Global WebSocket ticket store set for multi-app WebSocket authentication "
+            f"(store type: {type(ticket_store).__name__})"
+        )
+    else:
+        logger.warning("⚠️  Global WebSocket ticket store set to None")
+
 
 async def get_websocket_manager(app_slug: str) -> WebSocketConnectionManager:
     """
@@ -359,18 +384,190 @@ def _get_cookies_from_websocket(websocket: Any) -> dict[str, str]:
     return cookies
 
 
+async def _validate_websocket_origin_in_handler(websocket: Any, app_slug: str) -> bool:
+    """
+    Validate WebSocket Origin header in handler (since middleware may not intercept upgrades).
+
+    This provides CSWSH (Cross-Site WebSocket Hijacking) protection by validating
+    the Origin header before accepting the WebSocket connection.
+
+    Args:
+        websocket: FastAPI WebSocket instance
+        app_slug: App slug for context
+
+    Returns:
+        True if origin is valid, False otherwise
+    """
+    try:
+        # Get origin from headers
+        origin = None
+        if hasattr(websocket, "headers"):
+            origin = websocket.headers.get("origin")
+        elif hasattr(websocket, "scope") and "headers" in websocket.scope:
+            # Extract from ASGI scope headers
+            headers_dict = dict(websocket.scope["headers"])
+            origin_bytes = headers_dict.get(b"origin")
+            if not origin_bytes:
+                # Try case-insensitive lookup
+                for key, value in headers_dict.items():
+                    if isinstance(key, bytes) and key.lower() == b"origin":
+                        origin_bytes = value
+                        break
+            if origin_bytes:
+                origin = origin_bytes.decode("utf-8")
+
+        logger.info(
+            f"🔍 WebSocket origin validation in handler for '{app_slug}': " f"origin={origin}"
+        )
+
+        # Get allowed origins from app state (CORS config) FIRST
+        # This allows us to check if wildcard is allowed before rejecting missing origin
+        allowed_origins = []
+        try:
+            app = getattr(websocket, "app", None)
+            if app:
+                # Traverse up to find parent app with CORS config
+                while app:
+                    cors_config = getattr(app.state, "cors_config", None)
+                    if cors_config and cors_config.get("allow_origins"):
+                        origins = cors_config["allow_origins"]
+                        allowed_origins = origins if isinstance(origins, list) else [origins]
+                        logger.debug(
+                            f"Found CORS config on app '{getattr(app, 'title', 'unknown')}' "
+                            f"for '{app_slug}': {allowed_origins}"
+                        )
+                        break
+                    parent_app = getattr(app, "app", None)
+                    if parent_app is app:
+                        break
+                    app = parent_app
+        except (AttributeError, TypeError, KeyError) as e:
+            logger.debug(f"Could not read CORS config: {e}")
+
+        # Check if CORS allows all origins - if so, allow missing origin
+        if allowed_origins and "*" in allowed_origins:
+            if not origin:
+                logger.info(
+                    f"WebSocket upgrade missing Origin header for '{app_slug}', "
+                    "but CORS allows all origins (*) - allowing connection"
+                )
+                return True
+
+        if not origin:
+            # In development, allow connections without Origin
+            import os
+
+            env = os.getenv("ENVIRONMENT", "").lower()
+            g_nome_env = os.getenv("G_NOME_ENV", "").lower()
+            is_dev = env in ["development", "dev"] or g_nome_env in ["development", "dev"]
+            if is_dev:
+                logger.warning(
+                    f"WebSocket upgrade missing Origin header in development mode: "
+                    f"allowing connection for '{app_slug}'"
+                )
+                return True
+            else:
+                logger.warning(f"WebSocket upgrade missing Origin header for '{app_slug}'")
+                return False
+
+        # Normalize origin for comparison
+        def normalize_origin(orig: str) -> str:
+            """Normalize origin handling localhost variants."""
+            if not orig:
+                return orig
+            import re
+
+            normalized = re.sub(
+                r"://(0\.0\.0\.0|127\.0\.0\.1|localhost|::1)",
+                "://localhost",
+                orig.lower(),
+                flags=re.IGNORECASE,
+            )
+            normalized = re.sub(r":80$", "", normalized)
+            normalized = re.sub(r":443$", "", normalized)
+            return normalized.rstrip("/")
+
+        normalized_origin = normalize_origin(origin)
+        logger.debug(f"Normalized origin: {origin} -> {normalized_origin}")
+
+        # If no CORS config, generate fallback origins from request
+        if not allowed_origins:
+            # Try to get host from websocket scope
+            try:
+                if hasattr(websocket, "scope"):
+                    host = websocket.scope.get("server")
+                    if host:
+                        hostname = host[0] if isinstance(host, tuple) else host
+                        scheme = websocket.scope.get("scheme", "http")
+                        port = host[1] if isinstance(host, tuple) and len(host) > 1 else None
+
+                        # Generate localhost variants if server binds to 0.0.0.0
+                        if hostname in ["localhost", "0.0.0.0", "127.0.0.1", "::1"]:
+                            for variant in ["localhost", "127.0.0.1"]:
+                                if port and port not in [80, 443]:
+                                    allowed_origins.append(f"{scheme}://{variant}:{port}")
+                                else:
+                                    allowed_origins.append(f"{scheme}://{variant}")
+                        else:
+                            if port and port not in [80, 443]:
+                                allowed_origins.append(f"{scheme}://{hostname}:{port}")
+                            else:
+                                allowed_origins.append(f"{scheme}://{hostname}")
+            except (AttributeError, TypeError, KeyError) as e:
+                logger.debug(f"Could not determine fallback origins: {e}")
+
+        logger.debug(f"Allowed origins for '{app_slug}': {allowed_origins}")
+
+        # Check if origin matches any allowed origin
+        for allowed in allowed_origins:
+            if allowed == "*":
+                logger.warning(
+                    "WebSocket Origin validation using wildcard '*' - "
+                    "not recommended for production"
+                )
+                return True
+
+            normalized_allowed = normalize_origin(allowed)
+            if normalized_origin == normalized_allowed:
+                logger.info(
+                    f"✅ WebSocket origin validated in handler for '{app_slug}': "
+                    f"{origin} matches {allowed}"
+                )
+                return True
+
+        logger.warning(
+            f"❌ WebSocket origin validation failed for '{app_slug}': "
+            f"origin={origin} (normalized: {normalized_origin}), "
+            f"allowed={allowed_origins} "
+            f"(normalized: {[normalize_origin(a) for a in allowed_origins]})"
+        )
+        return False
+
+    except (
+        AttributeError,
+        KeyError,
+        UnicodeDecodeError,
+        TypeError,
+        ValueError,
+        RuntimeError,
+    ) as e:
+        logger.error(f"❌ Error validating WebSocket origin for '{app_slug}': {e}", exc_info=True)
+        # Fail secure - reject if we can't validate
+        return False
+
+
 async def authenticate_websocket(
     websocket: Any,
     app_slug: str,
     require_auth: bool = True,
 ) -> tuple[str | None, str | None]:
     """
-    Authenticate a WebSocket connection via httpOnly cookies.
+    Authenticate a WebSocket connection using multiple methods with fallback.
 
-    Uses cookie-based authentication with CSRF protection:
-    - Token stored in httpOnly cookie (not accessible to JavaScript)
-    - CSRF token validated via double-submit cookie pattern
-    - Origin validation provides additional protection
+    Authentication methods (in order of preference):
+    1. Ticket (query param or header) - short-lived, single-use, secure for multi-app SSO
+    2. Session key (query param or header) - encrypted session-based authentication
+    3. Cookie (httpOnly JWT token) - backward compatibility fallback
 
     Args:
         websocket: FastAPI WebSocket instance (can access headers before accept)
@@ -394,50 +591,201 @@ async def authenticate_websocket(
         return None, None
 
     try:
-        # Extract token from httpOnly cookie
-        # Use same cookie name as SharedAuthMiddleware for consistency
-        from ..auth.shared_middleware import AUTH_COOKIE_NAME
-
-        cookies = _get_cookies_from_websocket(websocket)
-        token = cookies.get(AUTH_COOKIE_NAME)  # Use mdb_auth_token (same as shared middleware)
+        # Helper function to get app and traverse hierarchy
+        def get_app_and_state():
+            """Get app instance and traverse hierarchy to find state."""
+            app = getattr(websocket, "app", None)
+            apps_checked = []
+            max_iterations = 10  # Safety limit to prevent infinite loops
+            iteration = 0
+            while app and iteration < max_iterations:
+                iteration += 1
+                app_title = getattr(app, "title", "unknown")
+                apps_checked.append(app_title)
+                yield app, app_title
+                # Try to get parent app (FastAPI routers have .app attribute pointing to parent)
+                parent_app = getattr(app, "app", None)
+                if parent_app is app or parent_app is None:  # Prevent infinite loop
+                    break
+                app = parent_app
+
+        # Try to get services from app state
+        websocket_ticket_store = _global_websocket_ticket_store
+        websocket_session_manager = None
+        user_pool = None
+
+        apps_checked_list = []
+        for app, app_title in get_app_and_state():
+            apps_checked_list.append(app_title)
+            if not websocket_ticket_store:
+                websocket_ticket_store = getattr(app.state, "websocket_ticket_store", None)
+                if websocket_ticket_store:
+                    logger.debug(
+                        f"Found websocket_ticket_store on app '{app_title}' "
+                        f"for app_slug '{app_slug}'"
+                    )
+            if not websocket_session_manager:
+                websocket_session_manager = getattr(app.state, "websocket_session_manager", None)
+                if websocket_session_manager:
+                    logger.info(
+                        f"✅ Found websocket_session_manager on app '{app_title}' "
+                        f"for app_slug '{app_slug}' (checked: {apps_checked_list})"
+                    )
+            if not user_pool:
+                user_pool = getattr(app.state, "user_pool", None)
+                if user_pool:
+                    logger.debug(f"Found user_pool on app '{app_title}' for app_slug '{app_slug}'")
 
-        if not token:
-            logger.error(
-                f"❌ No token cookie found for WebSocket connection to app '{app_slug}' "
-                f"(require_auth={require_auth}). "
-                f"Available cookies: {list(cookies.keys()) if cookies else 'none'}. "
-                f"Ensure httpOnly cookie is set during authentication."
+        if not websocket_session_manager:
+            logger.warning(
+                f"⚠️ websocket_session_manager not found for '{app_slug}' "
+                f"(checked apps: {apps_checked_list})"
             )
-            if require_auth:
-                return None, None  # Signal auth failure
-            return None, None
 
-        logger.info(
-            f"WebSocket token found in cookie for app '{app_slug}' " "(cookie-based authentication)"
+        # Method 1: Ticket authentication (preferred)
+        ticket = None
+        try:
+            if hasattr(websocket, "query_params"):
+                ticket = websocket.query_params.get("ticket")
+            if not ticket and hasattr(websocket, "headers"):
+                ticket = websocket.headers.get("X-WebSocket-Ticket")
+        except (AttributeError, TypeError, KeyError):
+            pass
+
+        if ticket and websocket_ticket_store:
+            try:
+                ticket_data = await websocket_ticket_store.validate_and_consume_ticket(ticket)
+                if ticket_data:
+                    user_id = ticket_data.get("user_id")
+                    user_email = ticket_data.get("user_email")
+                    logger.info(
+                        f"WebSocket authenticated successfully for app '{app_slug}': {user_email} "
+                        f"(method: ticket)"
+                    )
+                    return user_id, user_email
+            except (ValueError, TypeError, AttributeError, KeyError, RuntimeError) as e:
+                logger.debug(f"Ticket validation failed: {e}")
+
+        # Method 2: Session key authentication (fallback)
+        session_key = None
+        try:
+            if hasattr(websocket, "query_params"):
+                # Try dict-like access first
+                if hasattr(websocket.query_params, "get"):
+                    session_key = websocket.query_params.get("session_key")
+                # Fallback: try dict access
+                elif isinstance(websocket.query_params, dict):
+                    session_key = websocket.query_params.get("session_key")
+                # Fallback: try attribute access
+                elif hasattr(websocket.query_params, "session_key"):
+                    session_key = websocket.query_params.session_key
+            if not session_key and hasattr(websocket, "headers"):
+                session_key = websocket.headers.get("X-WebSocket-Session-Key")
+        except (AttributeError, TypeError, KeyError) as e:
+            logger.debug(f"Error extracting session key from websocket: {e}")
+            pass
+
+        logger.debug(
+            f"Session key extraction for '{app_slug}': "
+            f"session_key={'present' if session_key else 'missing'}, "
+            f"websocket_session_manager={'present' if websocket_session_manager else 'missing'}"
         )
 
-        # Decode and validate token
-        import jwt
+        if session_key and websocket_session_manager:
+            try:
+                logger.debug(
+                    f"Validating session key for '{app_slug}': "
+                    f"session_key={session_key[:16]}... (truncated), "
+                    f"manager_type={type(websocket_session_manager).__name__}"
+                )
+                session_data = await websocket_session_manager.validate_session(session_key)
+                if session_data:
+                    user_id = session_data.get("user_id")
+                    user_email = session_data.get("user_email")
+                    logger.info(
+                        f"WebSocket authenticated successfully for app '{app_slug}': {user_email} "
+                        f"(method: session_key)"
+                    )
+                    return user_id, user_email
+                else:
+                    logger.warning(
+                        f"Session key validation returned None for '{app_slug}': "
+                        f"session_key={session_key[:16]}... (truncated)"
+                    )
+            except RuntimeError as e:
+                # Handle event loop conflicts (e.g., TestClient with anyio)
+                error_msg = str(e).lower()
+                if "attached to a different loop" in error_msg or "different loop" in error_msg:
+                    logger.warning(
+                        f"Event loop conflict during session key validation for '{app_slug}': {e}. "
+                        "This may occur in test environments with TestClient. "
+                        "Session key validation failed."
+                    )
+                    # Return None to indicate authentication failure
+                    # Don't re-raise - let the normal auth failure path handle it
+                else:
+                    # Re-raise other RuntimeErrors
+                    logger.warning(
+                        f"Session key validation failed for '{app_slug}': {e}",
+                        exc_info=True,
+                    )
+                    raise
+            except (ValueError, TypeError, AttributeError, KeyError) as e:
+                logger.warning(
+                    f"Session key validation failed for '{app_slug}': {e}",
+                    exc_info=True,
+                )
+        elif session_key and not websocket_session_manager:
+            logger.error(
+                f"Session key provided for '{app_slug}' but websocket_session_manager not found. "
+                "Cannot validate session key."
+            )
 
-        from ..auth.dependencies import SECRET_KEY
-        from ..auth.jwt import decode_jwt_token
+        # Method 3: Cookie-based JWT authentication (backward compatibility)
+        from ..auth.shared_middleware import AUTH_COOKIE_NAME
 
-        try:
-            payload = decode_jwt_token(token, str(SECRET_KEY))
-            user_id = payload.get("sub") or payload.get("user_id")
-            user_email = payload.get("email")
+        cookies = _get_cookies_from_websocket(websocket)
+        auth_token = cookies.get(AUTH_COOKIE_NAME) if cookies else None
 
-            logger.info(
-                f"WebSocket authenticated successfully for app '{app_slug}': {user_email} "
-                f"(method: cookie)"
-            )
-            return user_id, user_email
-        except (jwt.ExpiredSignatureError, jwt.InvalidTokenError):
-            logger.exception(
-                f"❌ JWT decode error for app '{app_slug}'. "
-                f"Token present: {bool(token)}, Token length: {len(token) if token else 0}"
+        if auth_token and user_pool:
+            try:
+                # For tests that expect JWT errors to be raised, we need to validate
+                # the token structure first. However, validate_token handles JWT errors
+                # internally, so we'll let it handle validation and only catch non-JWT errors.
+                user = await user_pool.validate_token(auth_token)
+                if user:
+                    user_id = str(user.get("_id") or user.get("sub") or user.get("user_id"))
+                    user_email = user.get("email")
+                    logger.info(
+                        f"WebSocket authenticated successfully for app '{app_slug}': {user_email} "
+                        f"(method: cookie)"
+                    )
+                    return user_id, user_email
+            except (
+                ValueError,
+                TypeError,
+                AttributeError,
+                KeyError,
+                RuntimeError,
+            ) as e:
+                # Handle errors from token validation or user data access
+                logger.debug(f"Cookie token validation failed: {e}")
+            # Note: JWT errors (DecodeError, ExpiredSignatureError) are caught
+            # internally by validate_token and return None. For tests that need JWT
+            # errors raised, they should mock validate_token to raise them directly.
+
+        # No authentication method succeeded
+        if require_auth:
+            logger.error(
+                f"❌ WebSocket authentication failed for app '{app_slug}'. "
+                "No valid ticket, session key, or cookie found. "
+                "Generate ticket via /auth/ticket endpoint, "
+                "session key via /auth/websocket-session, "
+                "or ensure JWT cookie is present."
             )
-            raise
+            return None, None
+
+        return None, None
 
     except WebSocketDisconnect:
         raise
@@ -662,13 +1010,46 @@ def create_websocket_endpoint(
         # This print should appear in server logs when a WebSocket connection is attempted
         import sys
 
-        print(
-            f"🔌 [WEBSOCKET HANDLER CALLED] App: '{app_slug}', Path: {path}",
-            file=sys.stderr,
-            flush=True,
-        )
-        print(f"🔌 [WEBSOCKET HANDLER CALLED] App: '{app_slug}', Path: {path}", flush=True)
-        logger.info(f"🔌 [WEBSOCKET HANDLER CALLED] App: '{app_slug}', Path: {path}")
+        # Log BEFORE any operations to catch if handler is called at all
+        try:
+            print(
+                f"🔌 [WEBSOCKET HANDLER CALLED] App: '{app_slug}', Path: {path}",
+                file=sys.stderr,
+                flush=True,
+            )
+            print(f"🔌 [WEBSOCKET HANDLER CALLED] App: '{app_slug}', Path: {path}", flush=True)
+            logger.info(f"🔌 [WEBSOCKET HANDLER CALLED] App: '{app_slug}', Path: {path}")
+
+            # Try to access websocket properties to see if we can even access it
+            try:
+                ws_path = (
+                    getattr(websocket, "url", {}).path if hasattr(websocket, "url") else "unknown"
+                )
+                ws_headers = (
+                    dict(getattr(websocket, "headers", {})) if hasattr(websocket, "headers") else {}
+                )
+                origin = ws_headers.get("origin", "missing")
+                print(
+                    f"🔌 [WEBSOCKET DETAILS] Path: {ws_path}, Origin: {origin}, "
+                    f"Headers: {list(ws_headers.keys())}",
+                    file=sys.stderr,
+                    flush=True,
+                )
+            except (AttributeError, RuntimeError, TypeError) as access_error:
+                print(
+                    f"⚠️ [WEBSOCKET ACCESS ERROR] "
+                    f"Could not access websocket properties: {access_error}",
+                    file=sys.stderr,
+                    flush=True,
+                )
+        except (RuntimeError, AttributeError) as log_error:
+            # Even logging failed - this is very bad
+            print(
+                f"❌ [CRITICAL] Failed to log WebSocket handler call: {log_error}",
+                file=sys.stderr,
+                flush=True,
+            )
+
         connection = None
         try:
             # Log connection attempt with query params (can access before accept)
@@ -693,6 +1074,42 @@ def create_websocket_endpoint(
                 f"(require_auth={require_auth}, query_params={query_str})"
             )
 
+            # CRITICAL: Accept FIRST, then validate (FastAPI requires immediate accept)
+            # If we don't accept immediately, FastAPI closes with code 1000
+            # We'll validate origin after accept, but before processing messages
+            try:
+                await websocket.accept()
+                logger.info(
+                    f"✅ WebSocket accepted for app '{app_slug}' (before origin validation)"
+                )
+                print(f"✅ [WEBSOCKET ACCEPTED] App: '{app_slug}'", flush=True)
+            except (RuntimeError, ConnectionError, OSError, ValueError) as accept_error:
+                logger.error(
+                    f"❌ Failed to accept WebSocket for app '{app_slug}': {accept_error}",
+                    exc_info=True,
+                )
+                print(
+                    f"❌ [WEBSOCKET ACCEPT FAILED] App: '{app_slug}', Error: {accept_error}",
+                    flush=True,
+                )
+                return
+
+            # CRITICAL: Validate origin AFTER accepting (CSWSH protection)
+            # We accept first to satisfy FastAPI's requirements, then validate
+            origin_valid = await _validate_websocket_origin_in_handler(websocket, app_slug)
+            if not origin_valid:
+                logger.error(
+                    f"❌ WebSocket origin validation FAILED for app '{app_slug}' - "
+                    f"closing connection after accept"
+                )
+                try:
+                    await websocket.close(code=1008, reason="Invalid origin")
+                except (RuntimeError, ConnectionError, OSError):
+                    # WebSocket may already be closed or in invalid state
+                    pass
+                # Raise WebSocketDisconnect so TestClient can detect the rejection
+                raise WebSocketDisconnect(code=1008, reason="Invalid origin")
+
             # CRITICAL: Authenticate BEFORE accepting connection
             # This prevents CSRF middleware from rejecting established connections
             # We can access headers/query_params before accept() is called
@@ -717,14 +1134,17 @@ def create_websocket_endpoint(
                     f"rejecting connection. require_auth={require_auth}, "
                     f"user_id={user_id}, user_email={user_email}"
                 )
-                # Reject without accepting - FastAPI will send 403 if accept() not called
-                # We can't call websocket.close() before accept(), so we just return
-                # The connection will be rejected by the server
-                return
-
-            # Accept connection
-            await _accept_websocket_connection(websocket, app_slug)
-
+                # Close connection since we already accepted it
+                try:
+                    await websocket.close(code=1008, reason="Authentication failed")
+                except (RuntimeError, ConnectionError, OSError):
+                    # WebSocket may already be closed or in invalid state
+                    pass
+                # Raise WebSocketDisconnect so TestClient can detect the close
+                # WebSocketDisconnect is already imported at module level
+                raise WebSocketDisconnect(code=1008, reason="Authentication failed")
+
+            # Connection already accepted at line 984 - no need to accept again
             # Connect with metadata (websocket already accepted)
             connection = await manager.connect(websocket, user_id=user_id, user_email=user_email)
 
@@ -788,8 +1208,11 @@ def create_websocket_endpoint(
                     logger.warning(f"WebSocket receive error for app '{app_slug}': {e}")
                     await asyncio.sleep(0.1)
 
+        except WebSocketDisconnect as e:
+            # Re-raise WebSocketDisconnect so TestClient can detect it
+            logger.warning(f"WebSocket connection rejected for app '{app_slug}': {e}")
+            raise
         except (
-            WebSocketDisconnect,
             RuntimeError,
             OSError,
             ValueError,