96 lines
3.2 KiB
Python
96 lines
3.2 KiB
Python
"""Pure-ASGI middleware that decompresses compressed request bodies.
|
|
|
|
Gated on `SGLANG_ENABLE_REQUEST_DECOMPRESSION` and request header
|
|
`x-body-compressed`, whose value names the method. For example, a caller that
|
|
compressed the body with zstd sets the `x-body-compressed: zstd` header.
|
|
"""
|
|
|
|
import asyncio
|
|
import io
|
|
import logging
|
|
|
|
import zstandard
|
|
from fastapi.responses import Response
|
|
from starlette.datastructures import Headers
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _zstd_decompress(raw: bytes) -> bytes:
|
|
return zstandard.ZstdDecompressor().stream_reader(io.BytesIO(raw)).read()
|
|
|
|
|
|
_DECOMPRESSORS = {"zstd": _zstd_decompress}
|
|
|
|
|
|
def _rewrite_headers(headers, new_len):
|
|
"""Update headers to reflect body status after decompression."""
|
|
out = [
|
|
(k, v)
|
|
for (k, v) in headers
|
|
if k not in (b"content-length", b"x-body-compressed")
|
|
]
|
|
out.append((b"content-length", str(new_len).encode()))
|
|
return out
|
|
|
|
|
|
class RequestDecompressionMiddleware:
|
|
"""Decompress request body per request header `x-body-compressed`."""
|
|
|
|
def __init__(self, app):
|
|
self.app = app
|
|
|
|
async def __call__(self, scope, receive, send):
|
|
# No-op passthrough for any request without the compression header.
|
|
if scope["type"] != "http":
|
|
return await self.app(scope, receive, send)
|
|
method = Headers(scope=scope).get("x-body-compressed")
|
|
if method is None:
|
|
return await self.app(scope, receive, send)
|
|
|
|
# Fail loud on an unsupported compression method.
|
|
decompress = _DECOMPRESSORS.get(method)
|
|
if decompress is None:
|
|
return await Response(
|
|
f"unsupported x-body-compressed {method!r}; "
|
|
f"supported: {sorted(_DECOMPRESSORS)}",
|
|
status_code=400,
|
|
)(scope, receive, send)
|
|
|
|
# Collect request body.
|
|
body = b""
|
|
more_body = True
|
|
while more_body:
|
|
message = await receive()
|
|
# Incomplete body (e.g. client disconnect); hand off to later stages.
|
|
if message["type"] != "http.request":
|
|
return await self.app(scope, receive, send)
|
|
body += message.get("body", b"")
|
|
more_body = message.get("more_body", False)
|
|
|
|
# Decompress off the event loop by releasing the GIL around the C decompress.
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
body = await loop.run_in_executor(None, decompress, body)
|
|
except Exception as e:
|
|
logger.warning("request body decompress failed: %s", e)
|
|
return await Response("decompress failed", status_code=400)(
|
|
scope, receive, send
|
|
)
|
|
|
|
# Update the headers after decompression
|
|
scope = dict(scope)
|
|
scope["headers"] = _rewrite_headers(scope["headers"], len(body))
|
|
|
|
# Fake receiver to let later stages see the decompressed body.
|
|
body_sent = False
|
|
|
|
async def wrapped_receive():
|
|
nonlocal body_sent
|
|
if not body_sent:
|
|
body_sent = True
|
|
return {"type": "http.request", "body": body, "more_body": False}
|
|
return await receive()
|
|
|
|
await self.app(scope, wrapped_receive, send)
|