fix: convert SecurityHeadersMiddleware to pure ASGI (#26924)
SecurityHeadersMiddleware was the last middleware in the stack still subclassing BaseHTTPMiddleware, after CommitSession, AuthToken, WebsocketUpgradeGuard and Redirect were all moved to pure ASGI in utils/asgi_middleware.py. BaseHTTPMiddleware re-buffers the response body through an anyio task group, which has known issues with streaming and Content-Length-bearing responses (e.g. the FileResponse returned by /api/v1/audio/speech). Reimplement it as a pure-ASGI middleware that stamps the configured security headers onto the http.response.start message via MutableHeaders and forwards all body chunks untouched, matching the pattern already used by its four siblings. set_security_headers() and all its helpers are unchanged. Co-authored-by: classic298 <classic298@users.noreply.github.com>
This commit is contained in:
@@ -2,15 +2,33 @@ import os
|
||||
import re
|
||||
from typing import Dict
|
||||
|
||||
from fastapi import Request
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.datastructures import MutableHeaders
|
||||
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
||||
|
||||
|
||||
class SecurityHeadersMiddleware(BaseHTTPMiddleware):
|
||||
async def dispatch(self, request: Request, call_next):
|
||||
response = await call_next(request)
|
||||
response.headers.update(set_security_headers())
|
||||
return response
|
||||
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
|
||||
|
||||
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
if scope['type'] != 'http':
|
||||
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 set_security_headers().items():
|
||||
headers[key] = value
|
||||
await send(message)
|
||||
|
||||
await self.app(scope, receive, send_with_security_headers)
|
||||
|
||||
|
||||
def set_security_headers() -> Dict[str, str]:
|
||||
|
||||
Reference in New Issue
Block a user