SPB Git forge

spb/ai-atlas

Public
41commits 1branches 0releases
4.6 MBsize
maindefault branch
12 days agolast push
HTML 77.2% TypeScript 10.5% Python 9.6% JavaScript 2.5%
4.4 KB · 104 lines python
Raw Blame History
1"""Weak ETag + Cache-Control on public GET responses (API 1.1).23Pure ASGI middleware so it can sit *inside* GZip (the hash is computed on the uncompressed JSON body — gzip output embeds a4timestamp and would never be stable). `If-None-Match` matching the body hash → `304 Not Modified` with an empty body.5Admin, docs and non-200 responses are left untouched."""6from __future__ import annotations78import hashlib9from typing import Any1011from starlette.types import ASGIApp, Message, Receive, Scope, Send1213PUBLIC_PREFIX = "/api/v1/"14SKIP_PREFIXES = ("/api/v1/admin", "/api/v1/docs", "/api/v1/openapi.json", "/api/v1/api-keys")15CACHE_CONTROL = "public, max-age=60, stale-while-revalidate=300"16MAX_BUFFER = 8 * 1024 * 1024  # bodies above this are passed through unhashed171819def weak_etag(body: bytes) -> str:20    return 'W/"' + hashlib.sha1(body).hexdigest() + '"'212223def _etags_match(header: str | None, etag: str) -> bool:24    if not header:25        return False26    if header.strip() == "*":27        return True28    wanted = {t.strip() for t in header.split(",")}29    strong = etag.removeprefix("W/")30    return etag in wanted or strong in wanted or any((t.removeprefix("W/")) == strong for t in wanted)313233class ETagMiddleware:34    def __init__(self, app: ASGIApp) -> None:35        self.app = app3637    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:38        if scope["type"] != "http" or scope.get("method") != "GET":39            await self.app(scope, receive, send)40            return41        path: str = scope.get("path", "")42        if not path.startswith(PUBLIC_PREFIX) or path.startswith(SKIP_PREFIXES):43            await self.app(scope, receive, send)44            return45        headers = {k.decode("latin-1").lower(): v.decode("latin-1") for k, v in scope.get("headers", [])}46        inm = headers.get("if-none-match")4748        start: Message | None = None49        chunks: list[bytes] = []50        passthrough = False51        size = 05253        async def send_wrapper(message: Message) -> None:54            nonlocal start, passthrough, size55            if passthrough:56                await send(message)57                return58            if message["type"] == "http.response.start":59                start = message60                if message.get("status") != 200:61                    passthrough = True62                    await send(message)63                return64            if message["type"] == "http.response.body":65                body = message.get("body", b"")66                size += len(body)67                if size > MAX_BUFFER:68                    passthrough = True69                    assert start is not None70                    await send(start)71                    for c in chunks:72                        await send({"type": "http.response.body", "body": c, "more_body": True})73                    await send(message)74                    return75                chunks.append(body)76                if message.get("more_body"):77                    return78                assert start is not None79                full = b"".join(chunks)80                etag = weak_etag(full)81                raw_headers: list[tuple[bytes, bytes]] = [(k, v) for k, v in start.get("headers", []) if k.lower() not in (b"etag", b"cache-control")]82                raw_headers.append((b"etag", etag.encode("latin-1")))83                raw_headers.append((b"cache-control", CACHE_CONTROL.encode("latin-1")))84                if _etags_match(inm, etag):85                    keep = {b"etag", b"cache-control", b"vary", b"x-content-type-options", b"referrer-policy", b"access-control-allow-origin"}86                    hdrs = [(k, v) for k, v in raw_headers if k.lower() in keep]87                    await send({"type": "http.response.start", "status": 304, "headers": hdrs})88                    await send({"type": "http.response.body", "body": b"", "more_body": False})89                    return90                await send({"type": "http.response.start", "status": 200, "headers": raw_headers})91                await send({"type": "http.response.body", "body": full, "more_body": False})92                return93            await send(message)9495        await self.app(scope, receive, send_wrapper)969798def etag_of(value: Any) -> str:99    """Helper for tests and callers holding an already-rendered body."""100    return weak_etag(value if isinstance(value, bytes) else str(value).encode())101102103__all__ = ["CACHE_CONTROL", "ETagMiddleware", "etag_of", "weak_etag"]104