refac

backend/open_webui/main.py

@@ -205,12 +205,7 @@ from open_webui.utils import logger
 from open_webui.utils.access_control import has_permission
 from open_webui.utils.access_control.folders import has_folder_write_access
 from open_webui.utils.actions import chat_action as chat_action_handler
-from open_webui.utils.asgi_middleware import (
+from open_webui.utils.asgi_middleware import AppHTTPMiddleware
-    AuthTokenMiddleware,
-    CommitSessionMiddleware,
-    RedirectMiddleware,
-    WebsocketUpgradeGuardMiddleware,
-)
 from open_webui.utils.audit import AuditLevel, AuditLoggingMiddleware
 from open_webui.utils.auth import (
     create_admin_user,
@@ -267,7 +262,6 @@ from open_webui.utils.oauth import (
 )
 from open_webui.utils.plugin import install_tool_and_function_dependencies
 from open_webui.utils.redis import get_redis_client
-from open_webui.utils.security_headers import SecurityHeadersMiddleware
 from open_webui.utils.session_pool import cleanup_response, get_client_timeout, get_session, stream_wrapper
 from open_webui.utils.tool_approval import (
     ResolveToolCallForm,
@@ -807,11 +801,7 @@ if ENABLE_COMPRESSION_MIDDLEWARE:
 # `terminate_force_close` tracebacks under aiosqlite and as random
 # CancelledError storms across the request path. See
 # `open_webui.utils.asgi_middleware` for the rationale.
-app.add_middleware(RedirectMiddleware)
+app.add_middleware(AppHTTPMiddleware)
-app.add_middleware(SecurityHeadersMiddleware)
-app.add_middleware(CommitSessionMiddleware)
-app.add_middleware(AuthTokenMiddleware, fastapi_app=app)
-app.add_middleware(WebsocketUpgradeGuardMiddleware)
 
 
 app.add_middleware(
@@ -2906,7 +2896,7 @@ def _sync_db_ping() -> None:
     """Verify the database is reachable with a simple SELECT 1.
 
     Uses a raw connection from the engine pool instead of the thread-local
-    ScopedSession.  This is necessary because CommitSessionMiddleware
+    ScopedSession.  This is necessary because AppHTTPMiddleware
     deliberately skips healthcheck paths (/health, /ready, /health/db),
     so any ScopedSession opened on a healthcheck worker thread is never
     rolled back or removed.  If the session ever enters an invalid state

backend/open_webui/utils/asgi_middleware.py

@@ -40,6 +40,7 @@ from fastapi.security import HTTPAuthorizationCredentials
 from open_webui.env import CUSTOM_API_KEY_HEADER
 from open_webui.internal.db import ScopedSession
 from open_webui.utils.auth import get_http_authorization_cred
+from open_webui.utils.security_headers import set_security_headers
 from starlette.datastructures import MutableHeaders
 from starlette.requests import Request
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
@@ -47,9 +48,17 @@ from starlette.types import ASGIApp, Message, Receive, Scope, Send
 log = logging.getLogger(__name__)
 
 
-class CommitSessionMiddleware:
+class AppHTTPMiddleware:
-    """Commit and release the thread-local sync `ScopedSession` after each
+    """Open WebUI's pure-ASGI HTTP middleware.
-    HTTP request.
+
+    Keeps the app's request-wide behavior in one middleware layer without
+    hiding the old concerns behind a stack of wrappers:
+
+    * reject malformed `/ws/socket.io` upgrade requests
+    * stash bearer/cookie/API-key credentials on `request.state.token`
+    * stamp `X-Process-Time` and configured security headers
+    * serve the legacy `/watch` and `?shared=` redirects
+    * commit and release the thread-local sync `ScopedSession`
 
     Most requests now use the async session; the sync ScopedSession is
     only touched by startup, healthchecks, and a handful of legacy
@@ -81,156 +90,80 @@ class CommitSessionMiddleware:
 
     def __init__(self, app: ASGIApp) -> None:
         self.app = app
+        # Headers derive only from env vars, which are static for the process
+        # lifetime — compute them once instead of per response.
+        self._security_headers = list(set_security_headers().items())
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         if scope['type'] != 'http':
             await self.app(scope, receive, send)
             return
 
-        path = scope.get('path', '')
+        if await self._reject_invalid_websocket(scope, receive, send):
-        # Keep health probes independent from sync session commit/remove
-        # so DB pressure cannot delay or fail probe responses.
-        if path in {'/health', '/ready', '/health/db'}:
-            await self.app(scope, receive, send)
-            return
-
-        try:
-            await self.app(scope, receive, send)
-        except BaseException:
-            # Downstream did not complete successfully. Roll back any
-            # pending sync writes, release the connection, and let the
-            # exception propagate.
-            if ScopedSession.registry.has():
-                try:
-                    ScopedSession.rollback()
-                except Exception:
-                    log.exception('CommitSessionMiddleware: rollback failed after downstream error')
-                finally:
-                    ScopedSession.remove()
-            raise
-
-        # Nothing in this request touched the sync session: committing would
-        # only instantiate one to run an empty transaction.
-        if not ScopedSession.registry.has():
-            return
-
-        # Downstream completed. Commit pending sync work.
-        try:
-            ScopedSession.commit()
-        except Exception:
-            log.exception('CommitSessionMiddleware: post-request commit failed; response was already sent to client')
-            try:
-                ScopedSession.rollback()
-            except Exception:
-                log.exception('CommitSessionMiddleware: rollback failed after commit failure')
-            raise
-        finally:
-            # CRITICAL: remove() returns the connection to the pool.
-            # Without this, connections remain "checked out" and
-            # accumulate as "idle in transaction" in PostgreSQL.
-            ScopedSession.remove()
-
-
-class AuthTokenMiddleware:
-    """Extract the bearer/cookie/API-key credential and stash it on
-    `request.state.token`.
-
-    The header used for API-key transport is controlled by the
-    ``CUSTOM_API_KEY_HEADER`` environment variable (default ``x-api-key``).
-    This is useful when Open WebUI sits behind a reverse proxy that
-    consumes the ``Authorization`` header for its own authentication —
-    set the env var to a unique header (e.g. ``X-OpenWebUI-Key``) so
-    the middleware checks that instead and avoids the 401 short-circuit.
-
-    Routes that depend on `get_verified_user` etc. read this state.
-    Also stamps an `X-Process-Time` response header.
-    """
-
-    def __init__(self, app: ASGIApp, *, fastapi_app) -> None:
-        self.app = app
-        self._fastapi_app = fastapi_app
-
-    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
-        if scope['type'] != 'http':
-            await self.app(scope, receive, send)
             return
 
         start_time = time.monotonic()
         request = Request(scope)
+        self._set_token(request)
+        send_with_headers = self._send_with_headers(send, start_time)
 
+        try:
+            if await self._redirect_legacy_url(scope, receive, send_with_headers):
+                pass
+            # Keep health probes independent from sync session commit/remove so DB
+            # pressure cannot delay or fail probe responses.
+            elif scope.get('path', '') in {'/health', '/ready', '/health/db'}:
+                await self.app(scope, receive, send_with_headers)
+                return
+            else:
+                await self.app(scope, receive, send_with_headers)
+        except BaseException:
+            self._rollback_session('AppHTTPMiddleware: rollback failed after downstream error')
+            raise
+
+        self._commit_session()
+
+    def _set_token(self, request: Request) -> None:
         token = get_http_authorization_cred(request.headers.get('Authorization'))
-        if token is None:
+        if token is None and (cookie_token := request.cookies.get('token')):
-            cookie_token = request.cookies.get('token')
+            token = HTTPAuthorizationCredentials(scheme='Bearer', credentials=cookie_token)
-            if cookie_token:
+        if token is None and (api_key := request.headers.get(CUSTOM_API_KEY_HEADER)):
-                token = HTTPAuthorizationCredentials(scheme='Bearer', credentials=cookie_token)
+            token = HTTPAuthorizationCredentials(scheme='Bearer', credentials=api_key)
-        if token is None:
-            api_key = request.headers.get(CUSTOM_API_KEY_HEADER)
-            if api_key:
-                token = HTTPAuthorizationCredentials(scheme='Bearer', credentials=api_key)
-
         request.state.token = token
 
-        async def send_with_timing(message: Message) -> None:
+    def _send_with_headers(self, send: Send, start_time: float) -> Send:
+        async def send_with_headers(message: Message) -> None:
             if message['type'] == 'http.response.start':
-                process_time = time.monotonic() - start_time
                 headers = MutableHeaders(scope=message)
-                headers['X-Process-Time'] = f'{process_time:.6f}'
+                headers['X-Process-Time'] = f'{time.monotonic() - start_time:.6f}'
+                for key, value in self._security_headers:
+                    headers[key] = value
             await send(message)
 
-        await self.app(scope, receive, send_with_timing)
+        return send_with_headers
-
-
-class WebsocketUpgradeGuardMiddleware:
-    """Reject HTTP requests to `/ws/socket.io` that claim
-    `transport=websocket` but lack the proper `Upgrade`/`Connection`
-    headers.
-
-    Works around https://github.com/miguelgrinberg/python-engineio/issues/367
-    where engineio mishandles such requests.
-    """
-
-    def __init__(self, app: ASGIApp) -> None:
-        self.app = app
-
-    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
-        if scope['type'] != 'http':
-            await self.app(scope, receive, send)
-            return
 
+    async def _reject_invalid_websocket(self, scope: Scope, receive: Receive, send: Send) -> bool:
         path = scope.get('path', '')
-        if '/ws/socket.io' in path:
+        if '/ws/socket.io' not in path:
-            query_string = scope.get('query_string', b'').decode('latin-1', errors='replace')
+            return False
-            query_params = parse_qs(query_string)
-            if query_params.get('transport', [''])[0] == 'websocket':
-                headers = _scope_headers(scope)
-                upgrade = headers.get('upgrade', '').lower()
-                connection_tokens = [token.strip() for token in headers.get('connection', '').lower().split(',')]
-                if upgrade != 'websocket' or 'upgrade' not in connection_tokens:
-                    response = JSONResponse(
-                        status_code=400,
-                        content={'detail': 'Invalid WebSocket upgrade request'},
-                    )
-                    await response(scope, receive, send)
-                    return
 
-        await self.app(scope, receive, send)
+        query_params = parse_qs(scope.get('query_string', b'').decode('latin-1', errors='replace'))
+        if query_params.get('transport', [''])[0] != 'websocket':
+            return False
 
+        headers = _scope_headers(scope)
+        upgrade = headers.get('upgrade', '').lower()
+        connection_tokens = [token.strip() for token in headers.get('connection', '').lower().split(',')]
+        if upgrade == 'websocket' and 'upgrade' in connection_tokens:
+            return False
 
-class RedirectMiddleware:
+        response = JSONResponse(status_code=400, content={'detail': 'Invalid WebSocket upgrade request'})
-    """Rewrites a couple of legacy entry-points to the SPA's own routes:
+        await response(scope, receive, send)
+        return True
 
-    * ``GET /watch?v=ID`` (YouTube) → ``/?youtube=ID``
+    async def _redirect_legacy_url(self, scope: Scope, receive: Receive, send: Send) -> bool:
-    * ``GET /?shared=…`` (PWA share-target) → ``/?youtube=…`` /
+        if scope.get('method', '').upper() != 'GET':
-      ``/?load-url=…`` / ``/?q=…``
+            return False
-    """
-
-    def __init__(self, app: ASGIApp) -> None:
-        self.app = app
-
-    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
-        if scope['type'] != 'http' or scope.get('method', '').upper() != 'GET':
-            await self.app(scope, receive, send)
-            return
 
         path = scope.get('path', '')
         raw_query = scope.get('query_string', b'')
@@ -238,11 +171,9 @@ class RedirectMiddleware:
         # decode + parse_qs work for every other GET. (A false positive on the
         # substring check just falls through to the full parse below.)
         if not (path.endswith('/watch') or b'shared' in raw_query):
-            await self.app(scope, receive, send)
+            return False
-            return
 
-        query_string = raw_query.decode('latin-1', errors='replace')
+        query_params = parse_qs(raw_query.decode('latin-1', errors='replace'))
-        query_params = parse_qs(query_string)
 
         redirect_params: dict[str, str] = {}
         if path.endswith('/watch') and 'v' in query_params and query_params['v']:
@@ -270,9 +201,41 @@ class RedirectMiddleware:
             redirect_url = f'/?{urlencode(redirect_params)}'
             response = RedirectResponse(url=redirect_url)
             await response(scope, receive, send)
+            return True
+
+        return False
+
+    def _rollback_session(self, message: str) -> None:
+        if not ScopedSession.registry.has():
             return
 
-        await self.app(scope, receive, send)
+        try:
+            ScopedSession.rollback()
+        except Exception:
+            log.exception(message)
+        finally:
+            ScopedSession.remove()
+
+    def _commit_session(self) -> None:
+        # Nothing in this request touched the sync session: committing would
+        # only instantiate one to run an empty transaction.
+        if not ScopedSession.registry.has():
+            return
+
+        try:
+            ScopedSession.commit()
+        except Exception:
+            log.exception('AppHTTPMiddleware: post-request commit failed; response was already sent to client')
+            try:
+                ScopedSession.rollback()
+            except Exception:
+                log.exception('AppHTTPMiddleware: rollback failed after commit failure')
+            raise
+        finally:
+            # CRITICAL: remove() returns the connection to the pool.
+            # Without this, connections remain "checked out" and
+            # accumulate as "idle in transaction" in PostgreSQL.
+            ScopedSession.remove()
 
 
 def _scope_headers(scope: Scope) -> dict[str, str]:

backend/open_webui/utils/security_headers.py

@@ -2,37 +2,6 @@ import os
 import re
 from typing import Dict
 
-from starlette.datastructures import MutableHeaders
-from starlette.types import ASGIApp, Message, Receive, Scope, Send
-
-
-class SecurityHeadersMiddleware:
-    """Apply configured security headers to every HTTP response.
-
-    Pure ASGI to avoid BaseHTTPMiddleware's response re-buffering. See
-    open_webui.utils.asgi_middleware for the rationale.
-    """
-
-    def __init__(self, app: ASGIApp) -> None:
-        self.app = app
-        # Headers derive only from env vars, which are static for the process
-        # lifetime — compute them once instead of per response.
-        self._headers = list(set_security_headers().items())
-
-    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
-        if scope['type'] != 'http' or not self._headers:
-            await self.app(scope, receive, send)
-            return
-
-        async def send_with_security_headers(message: Message) -> None:
-            if message['type'] == 'http.response.start':
-                headers = MutableHeaders(scope=message)
-                for key, value in self._headers:
-                    headers[key] = value
-            await send(message)
-
-        await self.app(scope, receive, send_with_security_headers)
-
 
 def set_security_headers() -> Dict[str, str]:
     """