perf: convert APIKeyRestrictionMiddleware to pure ASGI (#22188)

backend/open_webui/main.py

@@ -1398,46 +1398,52 @@ app.add_middleware(RedirectMiddleware)
 app.add_middleware(SecurityHeadersMiddleware)
 
 
-class APIKeyRestrictionMiddleware(BaseHTTPMiddleware):
+class APIKeyRestrictionMiddleware:
-    async def dispatch(self, request: Request, call_next):
+    def __init__(self, app):
-        auth_header = request.headers.get("Authorization")
+        self.app = app
-        token = None
 
-        if auth_header:
+    async def __call__(self, scope, receive, send):
-            parts = auth_header.split(" ", 1)
+        if scope["type"] == "http":
-            if len(parts) == 2:
+            request = Request(scope)
-                token = parts[1]
+            auth_header = request.headers.get("Authorization")
+            token = None
 
-        # Only apply restrictions if an sk- API key is used
+            if auth_header:
-        if token and token.startswith("sk-"):
+                parts = auth_header.split(" ", 1)
-            # Check if restrictions are enabled
+                if len(parts) == 2:
-            if request.app.state.config.ENABLE_API_KEYS_ENDPOINT_RESTRICTIONS:
+                    token = parts[1]
-                allowed_paths = [
-                    path.strip()
-                    for path in str(
-                        request.app.state.config.API_KEYS_ALLOWED_ENDPOINTS
-                    ).split(",")
-                    if path.strip()
-                ]
 
-                request_path = request.url.path
+            # Only apply restrictions if an sk- API key is used
+            if token and token.startswith("sk-"):
+                # Check if restrictions are enabled
+                if app.state.config.ENABLE_API_KEYS_ENDPOINT_RESTRICTIONS:
+                    allowed_paths = [
+                        path.strip()
+                        for path in str(
+                            app.state.config.API_KEYS_ALLOWED_ENDPOINTS
+                        ).split(",")
+                        if path.strip()
+                    ]
 
-                # Match exact path or prefix path
+                    request_path = request.url.path
-                is_allowed = any(
-                    request_path == allowed or request_path.startswith(allowed + "/")
-                    for allowed in allowed_paths
-                )
 
-                if not is_allowed:
+                    # Match exact path or prefix path
-                    return JSONResponse(
+                    is_allowed = any(
-                        status_code=status.HTTP_403_FORBIDDEN,
+                        request_path == allowed
-                        content={
+                        or request_path.startswith(allowed + "/")
-                            "detail": "API key not allowed to access this endpoint."
+                        for allowed in allowed_paths
-                        },
                     )
 
-        response = await call_next(request)
+                    if not is_allowed:
-        return response
+                        await JSONResponse(
+                            status_code=status.HTTP_403_FORBIDDEN,
+                            content={
+                                "detail": "API key not allowed to access this endpoint."
+                            },
+                        )(scope, receive, send)
+                        return
+
+        await self.app(scope, receive, send)
 
 
 app.add_middleware(APIKeyRestrictionMiddleware)