|
7 | 7 | from contextlib import asynccontextmanager |
8 | 8 | from pydantic import BaseModel, Field |
9 | 9 | from typing import Any, Dict, List, Optional |
| 10 | +import json |
10 | 11 | import logging |
11 | 12 |
|
12 | 13 | from fastapi import FastAPI, HTTPException, BackgroundTasks, Request |
|
39 | 40 | ) |
40 | 41 | from services.humanitarian_verification import HumanitarianVerificationService |
41 | 42 |
|
| 43 | +class HTTPBodyTooLarge(Exception): |
| 44 | + """Internal signal raised when an incoming request body exceeds the |
| 45 | + configured `max_request_body_bytes` limit. Caught and converted to a |
| 46 | + 413 response by :class:`MaxRequestBodySizeMiddleware`.""" |
| 47 | + |
| 48 | + def __init__(self, limit: int, observed: int): |
| 49 | + super().__init__( |
| 50 | + f"Request body of {observed} bytes exceeds limit of {limit} bytes" |
| 51 | + ) |
| 52 | + self.limit = limit |
| 53 | + self.observed = observed |
| 54 | + |
| 55 | + |
| 56 | +class MaxRequestBodySizeMiddleware: |
| 57 | + """Reject HTTP requests whose body would exceed ``max_bytes``. |
| 58 | +
|
| 59 | + The middleware sits at the outer edge of the ASGI stack so that oversized |
| 60 | + requests are rejected *before* any other middleware (redirects, |
| 61 | + observability, rate limiting) or the application itself buffers the body. |
| 62 | + It is DoS-grade protection: clients can trip the limit either by sending a |
| 63 | + ``Content-Length`` header that exceeds the cap, or by streaming more bytes |
| 64 | + than the cap via chunked transfer encoding. |
| 65 | +
|
| 66 | + The middleware intentionally wraps the raw ASGI ``receive`` callable rather |
| 67 | + than using Starlette's ``BaseHTTPMiddleware`` — ``BaseHTTPMiddleware`` |
| 68 | + buffers the body in-memory which defeats the point of the limit. |
| 69 | + """ |
| 70 | + |
| 71 | + METHODS_WITH_BODY = ("POST", "PUT", "PATCH") |
| 72 | + |
| 73 | + def __init__(self, app, max_bytes: int, bypass_prefixes: Optional[List[str]] = None): |
| 74 | + self.app = app |
| 75 | + # Treat non-positive values as "disabled" — useful for tests that |
| 76 | + # don't want the limit to interfere. |
| 77 | + self.max_bytes = max_bytes if max_bytes and max_bytes > 0 else None |
| 78 | + # Always skip health/metrics/docs endpoints to match the pattern used |
| 79 | + # by monitor_requests. Allow additional prefixes via settings. |
| 80 | + default_bypass = [ |
| 81 | + "/health", |
| 82 | + "/", |
| 83 | + "/ai/metrics", |
| 84 | + "/docs", |
| 85 | + "/redoc", |
| 86 | + "/openapi.json", |
| 87 | + ] |
| 88 | + self.bypass_prefixes = tuple({*(default_bypass), *(bypass_prefixes or [])}) |
| 89 | + |
| 90 | + def _is_bypassed(self, path: str) -> bool: |
| 91 | + if path in self.bypass_prefixes: |
| 92 | + return True |
| 93 | + # Prefix matching only applies to entries that explicitly opt in |
| 94 | + # via a trailing '/'. The root '/' is intentionally excluded: |
| 95 | + # otherwise every HTTP path (which all begin with '/') would be |
| 96 | + # bypassed. |
| 97 | + return any( |
| 98 | + path.startswith(p) |
| 99 | + for p in self.bypass_prefixes |
| 100 | + if p.endswith("/") and p != "/" |
| 101 | + ) |
| 102 | + |
| 103 | + async def __call__(self, scope, receive, send): |
| 104 | + # Only operate on HTTP requests; pass through WebSocket / lifespan. |
| 105 | + if scope["type"] != "http": |
| 106 | + return await self.app(scope, receive, send) |
| 107 | + |
| 108 | + # No limit configured or no body expected — no-op. |
| 109 | + if self.max_bytes is None or scope["method"] not in self.METHODS_WITH_BODY: |
| 110 | + return await self.app(scope, receive, send) |
| 111 | + |
| 112 | + path = scope.get("path", "") |
| 113 | + if self._is_bypassed(path): |
| 114 | + return await self.app(scope, receive, send) |
| 115 | + |
| 116 | + # Eager check on Content-Length. If the client declared a body |
| 117 | + # larger than the limit, reject immediately without consuming any |
| 118 | + # bytes off the wire. |
| 119 | + try: |
| 120 | + content_length_hdr = None |
| 121 | + for name, value in scope.get("headers", []): |
| 122 | + if name == b"content-length": |
| 123 | + content_length_hdr = value.decode("latin-1") |
| 124 | + break |
| 125 | + if content_length_hdr is not None: |
| 126 | + declared = int(content_length_hdr) |
| 127 | + if declared > self.max_bytes: |
| 128 | + await self._log_rejection( |
| 129 | + scope, |
| 130 | + declared_or_observed=declared, |
| 131 | + reason="declared_size", |
| 132 | + ) |
| 133 | + return await self._send_413( |
| 134 | + send, |
| 135 | + observed=declared, |
| 136 | + reason="declared_size", |
| 137 | + ) |
| 138 | + except (ValueError, TypeError): |
| 139 | + # Malformed Content-Length — fall through to stream counting. |
| 140 | + pass |
| 141 | + |
| 142 | + total = 0 |
| 143 | + |
| 144 | + async def wrapped_receive(): |
| 145 | + nonlocal total |
| 146 | + message = await receive() |
| 147 | + mtype = message.get("type") |
| 148 | + if mtype == "http.request": |
| 149 | + chunk = message.get("body", b"") |
| 150 | + total += len(chunk) |
| 151 | + if total > self.max_bytes: |
| 152 | + # Signal the exception so that the outer __call__ can |
| 153 | + # emit a 413 even if the application has already started |
| 154 | + # producing a response. |
| 155 | + raise HTTPBodyTooLarge(self.max_bytes, total) |
| 156 | + return message |
| 157 | + |
| 158 | + try: |
| 159 | + await self.app(scope, wrapped_receive, send) |
| 160 | + except HTTPBodyTooLarge as exc: |
| 161 | + await self._log_rejection( |
| 162 | + scope, |
| 163 | + declared_or_observed=exc.observed, |
| 164 | + reason="streamed_size", |
| 165 | + ) |
| 166 | + await self._send_413( |
| 167 | + send, |
| 168 | + observed=exc.observed, |
| 169 | + reason="streamed_size", |
| 170 | + ) |
| 171 | + |
| 172 | + async def _send_413(self, send, observed: int, reason: str): |
| 173 | + """Emit a JSON 413 response using the project's ErrorEnvelope shape. |
| 174 | +
|
| 175 | + ``reason`` distinguishes eager (Content-Length) rejection from |
| 176 | + streamed rejection; the message is worded accordingly so the |
| 177 | + response is precise and not misleading. |
| 178 | + """ |
| 179 | + if reason == "declared_size": |
| 180 | + msg = ( |
| 181 | + f"Declared request body of {observed} bytes exceeds the " |
| 182 | + f"maximum allowed size of {self.max_bytes} bytes." |
| 183 | + ) |
| 184 | + else: |
| 185 | + msg = ( |
| 186 | + f"Request body streamed so far ({observed} bytes) exceeds " |
| 187 | + f"the maximum allowed size of {self.max_bytes} bytes." |
| 188 | + ) |
| 189 | + |
| 190 | + envelope = ErrorEnvelope( |
| 191 | + error=ErrorDetail( |
| 192 | + code="PAYLOAD_TOO_LARGE", |
| 193 | + message=msg, |
| 194 | + ) |
| 195 | + ).model_dump() |
| 196 | + body = json.dumps(envelope).encode("utf-8") |
| 197 | + |
| 198 | + await send( |
| 199 | + { |
| 200 | + "type": "http.response.start", |
| 201 | + "status": 413, |
| 202 | + "headers": [ |
| 203 | + (b"content-type", b"application/json"), |
| 204 | + (b"content-length", str(len(body)).encode("ascii")), |
| 205 | + ], |
| 206 | + } |
| 207 | + ) |
| 208 | + await send({"type": "http.response.body", "body": body}) |
| 209 | + |
| 210 | + async def _log_rejection( |
| 211 | + self, |
| 212 | + scope, |
| 213 | + declared_or_observed: int, |
| 214 | + reason: str, |
| 215 | + ) -> None: |
| 216 | + """Emit a structured warning so operators can correlate DoS attempts. |
| 217 | +
|
| 218 | + ``reason`` is either ``"declared_size"`` (Content-Length spoofing) |
| 219 | + or ``"streamed_size"`` (chunked transfer smuggling), so logs |
| 220 | + differentiate between attack classes. |
| 221 | + """ |
| 222 | + client = scope.get("client") |
| 223 | + client_str = f"{client[0]}:{client[1]}" if client else "unknown" |
| 224 | + logger.warning( |
| 225 | + "request body rejected: method=%s path=%s bytes=%d limit=%d " |
| 226 | + "client=%s reason=%s", |
| 227 | + scope.get("method"), |
| 228 | + scope.get("path"), |
| 229 | + declared_or_observed, |
| 230 | + self.max_bytes, |
| 231 | + client_str, |
| 232 | + reason, |
| 233 | + ) |
| 234 | + |
| 235 | + |
42 | 236 | limiter = Limiter(key_func=get_remote_address) |
43 | 237 |
|
44 | 238 | log_level_name = settings.log_level.upper() if hasattr(settings, "log_level") else "INFO" |
@@ -93,6 +287,20 @@ async def lifespan(app: FastAPI): |
93 | 287 | lifespan=lifespan, |
94 | 288 | ) |
95 | 289 |
|
| 290 | +# Register the body-size limit at the outermost layer so it short-circuits |
| 291 | +# before legacy redirects, observability middleware, or any handler buffers |
| 292 | +# the request body. |
| 293 | +_bypass_paths = [ |
| 294 | + p.strip() |
| 295 | + for p in (settings.request_body_bypass_paths or "").split(",") |
| 296 | + if p.strip() |
| 297 | +] |
| 298 | +app.add_middleware( |
| 299 | + MaxRequestBodySizeMiddleware, |
| 300 | + max_bytes=settings.max_request_body_bytes, |
| 301 | + bypass_prefixes=_bypass_paths, |
| 302 | +) |
| 303 | + |
96 | 304 | proof_of_life_analyzer = ProofOfLifeAnalyzer( |
97 | 305 | config=ProofOfLifeConfig( |
98 | 306 | confidence_threshold=settings.proof_of_life_confidence_threshold, |
|
0 commit comments