diff --git a/README.md b/README.md index 5cba1fb..85e6fab 100644 --- a/README.md +++ b/README.md @@ -207,13 +207,13 @@ Runtime limits are configured with environment variables: | `MEMORIES_RATE_LIMIT_WRITES_PER_MINUTE` | `30` | Applies to REST create, update, and delete requests. | | `MEMORIES_RATE_LIMIT_BATCH_PER_MINUTE` | `10` | Applies to `POST /memories/batch`. | | `MEMORIES_RATE_LIMIT_MCP_PER_MINUTE` | `240` | Applies to MCP streamable HTTP traffic under `/mcp`. | -| `MEMORIES_REQUEST_BODY_MAX_BYTES` | `1048576` | Applies to `POST` and `PATCH` REST/MCP HTTP requests with `Content-Length`. | +| `MEMORIES_REQUEST_BODY_MAX_BYTES` | `1048576` | Applies to `POST` and `PATCH` REST/MCP HTTP request bodies. | Rate limiting uses a fixed 60-second in-memory window per process. Clients are identified by a trimmed `X-Client-Id` header, capped at 128 characters, when present; otherwise the client IP is used. A limited request returns `429` with `Retry-After`, `X-RateLimit-Limit`, `X-RateLimit-Remaining`, and `X-RateLimit-Reset` headers. `/health` and `/ready` are exempt from normal rate limiting. -Request body-size enforcement uses the `Content-Length` header and returns `413` with `X-Request-Body-Limit`. This covers normal local agent clients but does not fully cover oversized chunked or missing-length bodies. +Request body-size enforcement returns `413` with `X-Request-Body-Limit` when the declared or streamed request body exceeds the configured limit. For behavioral load checks, use the standalone stress harness at [scripts/rate_limit_stress.py](scripts/rate_limit_stress.py). It exercises mixed agent usage, fast bursts, concurrent floods, and two-agent isolation across REST and MCP surfaces, then writes an HTML report. See [scripts/README.md](scripts/README.md) for setup, tuning options, and result interpretation. diff --git a/app/main.py b/app/main.py index ad2b4f8..dd26f7c 100644 --- a/app/main.py +++ b/app/main.py @@ -18,6 +18,7 @@ FixedWindowRateLimiter, reject_request_body_if_too_large, reject_request_if_rate_limited, + reject_request_stream_if_too_large, ) from app.request_logging import current_duration_ms, log_http_request from app.schemas import ( @@ -152,12 +153,12 @@ def complete_response(response): ) return set_request_id_header(response, request_id) - body_limit_response = reject_request_body_if_too_large( - request, - safety_config.request_body_max_bytes, - ) - if body_limit_response is not None: - return complete_response(body_limit_response) + origin = request.headers.get("origin") + if request.url.path.startswith("/mcp") and origin is not None: + if origin not in browser_client_config.allowed_origins: + return complete_response( + JSONResponse(status_code=403, content={"detail": "Origin not allowed"}), + ) rate_limit_response = reject_request_if_rate_limited( request, @@ -167,12 +168,19 @@ def complete_response(response): if rate_limit_response is not None: return complete_response(rate_limit_response) - origin = request.headers.get("origin") - if request.url.path.startswith("/mcp") and origin is not None: - if origin not in browser_client_config.allowed_origins: - return complete_response( - JSONResponse(status_code=403, content={"detail": "Origin not allowed"}), - ) + body_limit_response = reject_request_body_if_too_large( + request, + safety_config.request_body_max_bytes, + ) + if body_limit_response is not None: + return complete_response(body_limit_response) + + body_limit_response = await reject_request_stream_if_too_large( + request, + safety_config.request_body_max_bytes, + ) + if body_limit_response is not None: + return complete_response(body_limit_response) try: response = await call_next(request) diff --git a/app/request_limits.py b/app/request_limits.py index eebb38a..ee68d58 100644 --- a/app/request_limits.py +++ b/app/request_limits.py @@ -38,10 +38,13 @@ class _RateLimitWindow: def request_body_limit_applies(request: Request) -> bool: - if request.method.upper() not in BODY_LIMITED_METHODS: + return _body_limit_applies(request.method, request.url.path) + + +def _body_limit_applies(method: str, path: str) -> bool: + if method.upper() not in BODY_LIMITED_METHODS: return False - path = request.url.path return path == "/memories" or path.startswith("/memories/") or path.startswith("/mcp") @@ -68,12 +71,36 @@ def reject_request_body_if_too_large(request: Request, max_body_bytes: int) -> J and request_body_limit_applies(request) and content_length > max_body_bytes ): - return JSONResponse( - status_code=413, - content={"detail": "Request body too large"}, - headers={REQUEST_BODY_LIMIT_HEADER: str(max_body_bytes)}, - ) + return request_body_too_large_response(max_body_bytes) + + return None + + +def request_body_too_large_response(max_body_bytes: int) -> JSONResponse: + return JSONResponse( + status_code=413, + content={"detail": "Request body too large"}, + headers={REQUEST_BODY_LIMIT_HEADER: str(max_body_bytes)}, + ) + + +async def reject_request_stream_if_too_large( + request: Request, + max_body_bytes: int, +) -> JSONResponse | None: + if not request_body_limit_applies(request): + return None + + chunks: list[bytes] = [] + total_bytes = 0 + async for chunk in request.stream(): + total_bytes += len(chunk) + if total_bytes > max_body_bytes: + return request_body_too_large_response(max_body_bytes) + if chunk: + chunks.append(chunk) + request._body = b"".join(chunks) return None diff --git a/tests/contract/test_http_contract.py b/tests/contract/test_http_contract.py index a611633..8d00655 100644 --- a/tests/contract/test_http_contract.py +++ b/tests/contract/test_http_contract.py @@ -1,3 +1,4 @@ +import asyncio import json import logging import sqlite3 @@ -40,6 +41,76 @@ def request_log_payloads(caplog): return payloads +def asgi_request( + application, + method: str, + path: str, + chunks: list[bytes], + headers: dict[str, str] | None = None, +): + return asyncio.run(_asgi_request(application, method, path, chunks, headers or {})) + + +async def _asgi_request( + application, + method: str, + path: str, + chunks: list[bytes], + headers: dict[str, str], +): + response_messages = [] + request_messages = [ + { + "type": "http.request", + "body": chunk, + "more_body": index < len(chunks) - 1, + } + for index, chunk in enumerate(chunks) + ] + request_index = 0 + + async def receive(): + nonlocal request_index + if request_index < len(request_messages): + message = request_messages[request_index] + request_index += 1 + return message + return {"type": "http.disconnect"} + + async def send(message): + response_messages.append(message) + + await application( + { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": method, + "scheme": "http", + "path": path, + "raw_path": path.encode(), + "query_string": b"", + "headers": [(name.lower().encode(), value.encode()) for name, value in headers.items()], + "client": ("testclient", 50000), + "server": ("testserver", 80), + "root_path": "", + }, + receive, + send, + ) + + response_start = next( + message for message in response_messages if message["type"] == "http.response.start" + ) + response_body = b"".join( + message.get("body", b"") + for message in response_messages + if message["type"] == "http.response.body" + ) + response_headers = {name.decode(): value.decode() for name, value in response_start["headers"]} + return response_start["status"], response_headers, response_body + + def test_health_check_returns_ping_without_initializing_database( client: TestClient, data_file: Path ): @@ -189,18 +260,24 @@ def test_post_memory_rejects_oversized_body_before_json_parsing( ): monkeypatch.setenv("MEMORIES_REQUEST_BODY_MAX_BYTES", "10") caplog.set_level(logging.INFO, logger="app.main") - test_client = TestClient(app_main.create_app()) + test_app = app_main.create_app() - response = test_client.post( + status_code, headers, body = asgi_request( + test_app, + "POST", "/memories", - content="x" * 11, - headers={"Content-Type": "application/json", "X-Request-Id": "oversized-request"}, + [b"x" * 11], + { + "Content-Type": "application/json", + "Content-Length": "11", + "X-Request-Id": "oversized-request", + }, ) - assert response.status_code == 413 - assert response.json() == {"detail": "Request body too large"} - assert response.headers["X-Request-Body-Limit"] == "10" - assert response.headers["X-Request-Id"] == "oversized-request" + assert status_code == 413 + assert json.loads(body) == {"detail": "Request body too large"} + assert headers["x-request-body-limit"] == "10" + assert headers["x-request-id"] == "oversized-request" request_logs = request_log_payloads(caplog) assert len(request_logs) == 1 record, payload = request_logs[0] @@ -212,6 +289,74 @@ def test_post_memory_rejects_oversized_body_before_json_parsing( assert read_database(data_file) == [] +def test_post_memory_rejects_oversized_body_without_content_length( + monkeypatch, data_file: Path, caplog +): + monkeypatch.setenv("MEMORIES_REQUEST_BODY_MAX_BYTES", "10") + caplog.set_level(logging.INFO, logger="app.main") + test_app = app_main.create_app() + + status_code, headers, body = asgi_request( + test_app, + "POST", + "/memories", + [b"x" * 6, b"x" * 5], + {"Content-Type": "application/json", "X-Request-Id": "streamed-request"}, + ) + + assert status_code == 413 + assert json.loads(body) == {"detail": "Request body too large"} + assert headers["x-request-body-limit"] == "10" + assert headers["x-request-id"] == "streamed-request" + request_logs = request_log_payloads(caplog) + assert len(request_logs) == 1 + record, payload = request_logs[0] + assert record.levelno == logging.INFO + assert payload["method"] == "POST" + assert payload["path"] == "/memories" + assert payload["status"] == 413 + assert payload["request_id"] == "streamed-request" + assert read_database(data_file) == [] + + +def test_post_memory_replays_body_without_content_length_under_limit(monkeypatch, data_file: Path): + monkeypatch.setenv("MEMORIES_REQUEST_BODY_MAX_BYTES", "100") + test_app = app_main.create_app() + + status_code, _headers, body = asgi_request( + test_app, + "POST", + "/memories", + [b'{"content": "Learning FastAPI testing"}'], + {"Content-Type": "application/json"}, + ) + + assert status_code == 422 + assert any( + error["loc"] == ["body", "tags"] and error["type"] == "missing" + for error in json.loads(body)["detail"] + ) + assert read_database(data_file) == [] + + +def test_post_memory_accepts_chunked_body_exactly_at_limit(monkeypatch, data_file: Path): + body = b'{"content":"x","tags":["boundary"]}' + monkeypatch.setenv("MEMORIES_REQUEST_BODY_MAX_BYTES", str(len(body))) + test_app = app_main.create_app() + + status_code, _headers, response_body = asgi_request( + test_app, + "POST", + "/memories", + [body[:10], body[10:]], + {"Content-Type": "application/json"}, + ) + + assert status_code == 200 + assert json.loads(response_body)["content"] == "x" + assert len(read_database(data_file)) == 1 + + def test_post_memory_rate_limit_uses_client_id_identity(monkeypatch, data_file: Path): monkeypatch.setenv("MEMORIES_RATE_LIMIT_WRITES_PER_MINUTE", "1") test_client = TestClient(app_main.create_app()) @@ -295,38 +440,52 @@ def test_rate_limiting_can_be_disabled(monkeypatch, data_file: Path): assert len(read_database(data_file)) == 2 -def test_body_size_limit_runs_before_rate_limit(monkeypatch, data_file: Path): +def test_oversized_request_counts_against_rate_limit(monkeypatch, data_file: Path): monkeypatch.setenv("MEMORIES_REQUEST_BODY_MAX_BYTES", "10") monkeypatch.setenv("MEMORIES_RATE_LIMIT_WRITES_PER_MINUTE", "1") - test_client = TestClient(app_main.create_app()) + test_app = app_main.create_app() - oversized_response = test_client.post( + oversized_status_code, _oversized_headers, oversized_body = asgi_request( + test_app, + "POST", "/memories", - content="x" * 11, - headers={"Content-Type": "application/json"}, + [b"x" * 11], + {"Content-Type": "application/json", "Content-Length": "11"}, + ) + second_status_code, second_headers, second_body = asgi_request( + test_app, + "POST", + "/memories", + [b"{}"], + {"Content-Type": "application/json", "Content-Length": "2"}, ) - first_counted_response = test_client.post("/memories", json={}) - second_counted_response = test_client.post("/memories", json={}) - assert oversized_response.status_code == 413 - assert first_counted_response.status_code == 422 - assert_rate_limited(second_counted_response, 1) + assert oversized_status_code == 413 + assert json.loads(oversized_body) == {"detail": "Request body too large"} + assert second_status_code == 429 + assert json.loads(second_body) == {"detail": "Rate limit exceeded"} + assert second_headers["retry-after"].isdigit() + assert second_headers["x-ratelimit-limit"] == "1" + assert second_headers["x-ratelimit-remaining"] == "0" + assert second_headers["x-ratelimit-reset"].isdigit() assert read_database(data_file) == [] def test_patch_memory_rejects_oversized_body_before_json_parsing(monkeypatch, data_file: Path): monkeypatch.setenv("MEMORIES_REQUEST_BODY_MAX_BYTES", "10") - test_client = TestClient(app_main.create_app()) + test_app = app_main.create_app() - response = test_client.patch( + status_code, headers, body = asgi_request( + test_app, + "PATCH", "/memories/1", - content="x" * 11, - headers={"Content-Type": "application/json"}, + [b"x" * 11], + {"Content-Type": "application/json", "Content-Length": "11"}, ) - assert response.status_code == 413 - assert response.json() == {"detail": "Request body too large"} - assert response.headers["X-Request-Body-Limit"] == "10" + assert status_code == 413 + assert json.loads(body) == {"detail": "Request body too large"} + assert headers["x-request-body-limit"] == "10" assert read_database(data_file) == [] diff --git a/tests/contract/test_mcp_http_transport.py b/tests/contract/test_mcp_http_transport.py index 00e8faa..a249f25 100644 --- a/tests/contract/test_mcp_http_transport.py +++ b/tests/contract/test_mcp_http_transport.py @@ -1,3 +1,4 @@ +import asyncio import importlib import json import logging @@ -31,6 +32,88 @@ def request_log_payloads(caplog): return payloads +def asgi_post( + application, + path: str, + chunks: list[bytes], + headers: dict[str, str] | None = None, + *, + fail_if_body_read: bool = False, +): + return asyncio.run( + _asgi_post( + application, + path, + chunks, + headers or {}, + fail_if_body_read=fail_if_body_read, + ) + ) + + +async def _asgi_post( + application, + path: str, + chunks: list[bytes], + headers: dict[str, str], + *, + fail_if_body_read: bool, +): + response_messages = [] + request_messages = [ + { + "type": "http.request", + "body": chunk, + "more_body": index < len(chunks) - 1, + } + for index, chunk in enumerate(chunks) + ] + request_index = 0 + + async def receive(): + nonlocal request_index + if fail_if_body_read: + raise AssertionError("request body was read before the request was rejected") + if request_index < len(request_messages): + message = request_messages[request_index] + request_index += 1 + return message + return {"type": "http.disconnect"} + + async def send(message): + response_messages.append(message) + + await application( + { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": path, + "raw_path": path.encode(), + "query_string": b"", + "headers": [(name.lower().encode(), value.encode()) for name, value in headers.items()], + "client": ("testclient", 50000), + "server": ("testserver", 80), + "root_path": "", + }, + receive, + send, + ) + + response_start = next( + message for message in response_messages if message["type"] == "http.response.start" + ) + response_body = b"".join( + message.get("body", b"") + for message in response_messages + if message["type"] == "http.response.body" + ) + response_headers = {name.decode(): value.decode() for name, value in response_start["headers"]} + return response_start["status"], response_headers, response_body + + def build_client_with_fresh_mcp(): importlib.reload(mcp_server_module) reloaded_main = importlib.reload(main_module) @@ -82,6 +165,25 @@ def test_mcp_http_rejects_browser_requests_when_no_local_allowlist_exists( assert payload["request_id"] == "mcp-origin-check" +def test_mcp_http_rejects_disallowed_origin_without_reading_streamed_body( + monkeypatch, tmp_path: Path +): + config_path = tmp_path / "missing_mcp_browser_clients.local.json" + monkeypatch.setattr(config_module, "MCP_BROWSER_CLIENTS_LOCAL_FILE", config_path) + test_app = main_module.create_app() + + status_code, _headers, body = asgi_post( + test_app, + "/mcp", + [b"x" * 100], + {"Content-Type": "application/json", "Origin": "http://localhost:3000"}, + fail_if_body_read=True, + ) + + assert status_code == 403 + assert json.loads(body) == {"detail": "Origin not allowed"} + + def test_mcp_http_allows_configured_browser_origin(monkeypatch, tmp_path: Path): with build_client_with_browser_config( monkeypatch, @@ -104,18 +206,51 @@ def test_mcp_http_allows_configured_browser_origin(monkeypatch, tmp_path: Path): def test_mcp_http_rejects_oversized_body_before_transport_parsing(monkeypatch): monkeypatch.setenv("MEMORIES_REQUEST_BODY_MAX_BYTES", "10") - client = TestClient(main_module.create_app()) + test_app = main_module.create_app() + + status_code, headers, body = asgi_post( + test_app, + "/mcp", + [b"x" * 11], + {"Content-Type": "application/json", "Content-Length": "11"}, + ) + + assert status_code == 413 + assert json.loads(body) == {"detail": "Request body too large"} + assert headers["x-request-body-limit"] == "10" + assert_uuid(headers["x-request-id"]) + + +def test_mcp_http_rejects_chunked_oversized_body_without_content_length(monkeypatch): + monkeypatch.setenv("MEMORIES_REQUEST_BODY_MAX_BYTES", "10") + test_app = main_module.create_app() + + status_code, headers, body = asgi_post( + test_app, + "/mcp", + [b"x" * 6, b"x" * 5], + {"Content-Type": "application/json"}, + ) + + assert status_code == 413 + assert json.loads(body) == {"detail": "Request body too large"} + assert headers["x-request-body-limit"] == "10" + + +def test_mcp_http_accepts_chunked_body_exactly_at_limit(monkeypatch): + body = b"{}" + monkeypatch.setenv("MEMORIES_REQUEST_BODY_MAX_BYTES", str(len(body))) + test_app = main_module.create_app() - response = client.post( + status_code, headers, _body = asgi_post( + test_app, "/mcp", - content="x" * 11, - headers={"Content-Type": "application/json"}, + [body[:1], body[1:]], + {"Content-Type": "application/json"}, ) - assert response.status_code == 413 - assert response.json() == {"detail": "Request body too large"} - assert response.headers["X-Request-Body-Limit"] == "10" - assert_uuid(response.headers["X-Request-Id"]) + assert status_code == 307 + assert headers["location"] == "http://testserver/mcp/" def test_mcp_http_rate_limit_returns_stable_429(monkeypatch):