File size: 4,053 Bytes
fd4e724 03cabce fd4e724 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | import hashlib
import hmac
import time
from collections import deque
from fastapi.responses import JSONResponse
class Guard:
"""Authenticate before multipart parsing; bound chunked and fixed-length bodies."""
def __init__(self, app, token, max_bytes, rate=60):
self.app = app
self.token = token
self.max_bytes = max_bytes + 64 * 1024
self.rate = rate
self.calls = deque()
async def __call__(self, scope, receive, send):
if scope["type"] != "http":
return await self.app(scope, receive, send)
headers = dict(scope.get("headers", []))
path = scope["path"]
# Only /health and documentation endpoints are unauthenticated
if not any(path.rstrip("/").endswith(p) for p in ("/health", "/api-docs", "/docs", "/openapi.json", "/redoc")):
if not self.token:
return await JSONResponse({"code": "service_unconfigured"}, status_code=503)(
scope, receive, send
)
auth = headers.get(b"authorization", b"")
expected = b"Bearer " + self.token.encode("utf-8")
if not hmac.compare_digest(auth, expected):
return await JSONResponse({"code": "unauthorized"}, status_code=401)(
scope, receive, send
)
if scope["method"] == "POST":
now = time.monotonic()
while self.calls and now - self.calls[0] > 60:
self.calls.popleft()
if len(self.calls) >= self.rate:
return await JSONResponse({"code": "rate_limit"}, status_code=429)(
scope, receive, send
)
self.calls.append(now)
try:
if int(headers.get(b"content-length", b"0")) > self.max_bytes:
return await JSONResponse({"code": "upload_limit"}, status_code=413)(
scope, receive, send
)
except ValueError:
return await JSONResponse({"code": "invalid_length"}, status_code=400)(
scope, receive, send
)
# Buffer at most the bounded body before Starlette sees it; no spool files on oversize.
body = bytearray()
if scope["method"] in ("POST", "PUT", "PATCH"):
while True:
message = await receive()
if message["type"] == "http.disconnect":
return
body.extend(message.get("body", b""))
if len(body) > self.max_bytes:
return await JSONResponse({"code": "upload_limit"}, status_code=413)(
scope, receive, send
)
if not message.get("more_body", False):
break
delivered = False
async def bounded_receive():
nonlocal delivered
if not delivered:
delivered = True
return {
"type": "http.request",
"body": bytes(body),
"more_body": False,
}
return await receive()
return await self._invoke(scope, bounded_receive, send)
return await self._invoke(scope, receive, send)
async def _invoke(self, scope, receive, send):
started = False
async def safe_send(message):
nonlocal started
if message["type"] == "http.response.start":
started = True
await send(message)
try:
await self.app(scope, receive, safe_send)
except Exception:
# Avoid leaking unexpected parser/provider exceptions into server tracebacks.
if not started:
await JSONResponse({"code": "internal_error"}, status_code=500)(
scope, receive, send
)
def owner(request):
return hashlib.sha256(request.headers.get("authorization", "").encode()).hexdigest()
|