diff --git a/python/openai/README.md b/python/openai/README.md index 4598a5a43f..d32ce97356 100644 --- a/python/openai/README.md +++ b/python/openai/README.md @@ -949,3 +949,28 @@ curl -H "api-key: my-secret-key" \ # Multiple APIs in single argument with shared authentication --openai-restricted-api "inference,model-repository shared-key shared-secret" ``` + +## HTTP Request Body Size Limit + +The frontend enforces a maximum request body size prior to JSON parsing. Requests that exceed this limit are rejected with an error response. + +Use `--http-max-input-size` to configure the limit (default: `67108864` bytes / 64 MiB): + +```bash +python3 openai_frontend/main.py \ + --model-repository /path/to/models \ + --tokenizer meta-llama/Meta-Llama-3.1-8B-Instruct \ + --http-max-input-size 67108864 +``` + +The limit applies to all endpoints. Example error response: + +```json +{ + "error": { + "message": "Request content size exceeds the maximum allowed input size of 67108864 bytes. Use --http-max-input-size to increase the limit.", + "type": "invalid_request_error", + "code": "content_too_large" + } +} +``` diff --git a/python/openai/openai_frontend/frontend/fastapi/middleware/request_size.py b/python/openai/openai_frontend/frontend/fastapi/middleware/request_size.py new file mode 100644 index 0000000000..333a7968f4 --- /dev/null +++ b/python/openai/openai_frontend/frontend/fastapi/middleware/request_size.py @@ -0,0 +1,163 @@ +# Copyright 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions +# are met: +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of NVIDIA CORPORATION nor the names of its +# contributors may be used to endorse or promote products derived +# from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY +# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY +# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from fastapi.responses import JSONResponse +from starlette.types import ASGIApp, Message, Receive, Scope, Send +from utils.utils import StatusCode, validate_positive_int + + +async def _disconnect_receive() -> Message: + return {"type": "http.disconnect"} + + +class RequestSizeLimitMiddleware: + """ + Reject HTTP requests whose body exceeds ``http_max_input_size`` bytes. + First validation rejects on the Content-Length header before any body bytes are + read. Second validation streams the body chunks, counting bytes as they arrive, + and rejects as soon as the running total crosses the limit. Driving + receive() from the middleware protects every endpoint, including + handlers that never read the body. The buffered body is released the + moment the application consumes it, so the middleware contributes no + sustained memory overhead. + """ + + def __init__(self, app: ASGIApp, http_max_input_size: int) -> None: + self.app = app + self.http_max_input_size = validate_positive_int(http_max_input_size) + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + # Stage 1: reject on Content-Length before reading any body bytes. + for name, value in scope["headers"]: + if name != b"content-length": + continue + try: + content_length = int(value) + except ValueError: + await self._send_error( + scope, + send, + StatusCode.CLIENT_ERROR, + "invalid_content_length", + "Invalid Content-Length header: not an integer.", + ) + return + if content_length < 0: + await self._send_error( + scope, + send, + StatusCode.CLIENT_ERROR, + "invalid_content_length", + "Invalid Content-Length header: must be non-negative.", + ) + return + if content_length > self.http_max_input_size: + await self._send_error( + scope, + send, + StatusCode.CONTENT_TOO_LARGE, + "content_too_large", + self._oversized_request_message( + content_length, self.http_max_input_size + ), + ) + return + break + + # Stage 2: count chunks as they arrive, reject if total exceeds limit. + body_chunks: list[bytes] = [] + total = 0 + while True: + message = await receive() + if message["type"] != "http.request": + return + chunk = message.get("body", b"") + total += len(chunk) + if total > self.http_max_input_size: + await self._send_error( + scope, + send, + StatusCode.CONTENT_TOO_LARGE, + "content_too_large", + self._oversized_request_message(total, self.http_max_input_size), + ) + return + body_chunks.append(chunk) + if not message.get("more_body", False): + break + + # Assemble the buffered body and replay it to the app. + body_message: Message = { + "type": "http.request", + "body": b"".join(body_chunks), + "more_body": False, + } + del body_chunks + + async def replay_receive() -> Message: + nonlocal body_message + if body_message is not None: + # Drop the reference on hand-off so the body is freed while + # the app processes it, instead of being held by this closure. + message, body_message = body_message, None + return message + # Body already delivered — delegate to the original receive() so + # streaming responses can wait for the real client disconnect. + return await receive() + + await self.app(scope, replay_receive, send) + + @staticmethod + def _oversized_request_message(actual_bytes: int, max_bytes: int) -> str: + return ( + f"Request size of {actual_bytes} bytes exceeds the maximum allowed " + f"input size of {max_bytes} bytes. " + f"Use --http-max-input-size to increase the limit." + ) + + async def _send_error( + self, + scope: Scope, + send: Send, + status_code: int, + code: str, + message: str, + ) -> None: + response = JSONResponse( + status_code=status_code, + content={ + "error": { + "message": message, + "type": "invalid_request_error", + "code": code, + } + }, + ) + await response(scope, _disconnect_receive, send) diff --git a/python/openai/openai_frontend/frontend/fastapi_frontend.py b/python/openai/openai_frontend/frontend/fastapi_frontend.py index f7fa8aab3c..752befd8bc 100644 --- a/python/openai/openai_frontend/frontend/fastapi_frontend.py +++ b/python/openai/openai_frontend/frontend/fastapi_frontend.py @@ -34,6 +34,7 @@ APIRestrictionMiddleware, RestrictedFeatures, ) +from frontend.fastapi.middleware.request_size import RequestSizeLimitMiddleware from frontend.fastapi.routers import ( chat, completions, @@ -43,6 +44,7 @@ observability, ) from frontend.frontend import OpenAIFrontend +from utils.utils import HTTP_DEFAULT_MAX_INPUT_SIZE class FastApiFrontend(OpenAIFrontend): @@ -53,10 +55,12 @@ def __init__( port: int = 8000, log_level: str = "info", restricted_apis: list = None, + http_max_input_size: int = HTTP_DEFAULT_MAX_INPUT_SIZE, ): self.host: str = host self.port: int = port self.log_level: str = log_level + self.http_max_input_size: int = http_max_input_size if restricted_apis: self.restricted_apis: RestrictedFeatures = RestrictedFeatures( restricted_apis @@ -111,6 +115,7 @@ def _create_app(self): self._add_cors_middleware(app) if self.restricted_apis != None: self._add_api_restriction_middleware(app) + self._add_request_size_limit_middleware(app) return app @@ -137,3 +142,10 @@ def _add_api_restriction_middleware(self, app: FastAPI): print( f"[INFO] API restrictions enabled. Restricted API endpoints: {self.restricted_apis.RestrictionDict()}" ) + + def _add_request_size_limit_middleware(self, app: FastAPI): + app.add_middleware( + RequestSizeLimitMiddleware, + http_max_input_size=self.http_max_input_size, + ) + print(f"[INFO] HTTP request size limit set to {self.http_max_input_size} bytes") diff --git a/python/openai/openai_frontend/main.py b/python/openai/openai_frontend/main.py index ac310e662c..cf9228b9fc 100755 --- a/python/openai/openai_frontend/main.py +++ b/python/openai/openai_frontend/main.py @@ -34,6 +34,7 @@ import tritonserver from engine.triton_engine import TritonLLMEngine from frontend.fastapi_frontend import FastApiFrontend +from utils.utils import HTTP_DEFAULT_MAX_INPUT_SIZE, validate_positive_int def signal_handler( @@ -197,6 +198,14 @@ def parse_args(): action="append", help="Restrict access to specific OpenAI API endpoints. Format: ',,... ' (e.g., 'inference,model-repository admin-key admin-value'). If not specified, all endpoints are allowed.", ) + openai_group.add_argument( + "--http-max-input-size", + type=validate_positive_int, + default=HTTP_DEFAULT_MAX_INPUT_SIZE, + help=f"Maximum allowed HTTP request input size in bytes for the OpenAI " + f"frontend (default: {HTTP_DEFAULT_MAX_INPUT_SIZE}, i.e. 64 MiB). " + "Requests exceeding this limit will be rejected.", + ) # KServe Predict v2 Frontend kserve_group = parser.add_argument_group("Triton KServe Frontend") @@ -269,6 +278,7 @@ def main(): port=args.openai_port, log_level=args.uvicorn_log_level, restricted_apis=args.openai_restricted_api, + http_max_input_size=args.http_max_input_size, ) except ValueError as e: print( diff --git a/python/openai/openai_frontend/utils/utils.py b/python/openai/openai_frontend/utils/utils.py index c8fd3609f6..86b4d921e4 100644 --- a/python/openai/openai_frontend/utils/utils.py +++ b/python/openai/openai_frontend/utils/utils.py @@ -1,4 +1,4 @@ -# Copyright 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions @@ -26,6 +26,10 @@ from enum import IntEnum +# Default value for the --http-max-input-size CLI flag (64 MiB). +# Same as HTTP_DEFAULT_MAX_INPUT_SIZE in src/common.h. +HTTP_DEFAULT_MAX_INPUT_SIZE: int = 1 << 26 + class ServerError(Exception): """Exception raised for server errors.""" @@ -44,4 +48,15 @@ class StatusCode(IntEnum): CLIENT_ERROR = 400 AUTHORIZATION_ERROR = 401 NOT_FOUND = 404 + CONTENT_TOO_LARGE = 413 SERVER_ERROR = 500 + + +def validate_positive_int(value: object) -> int: + try: + ivalue = int(value) + except (TypeError, ValueError): + raise ValueError(f"value is not an integer, got {value!r}") + if ivalue <= 0: + raise ValueError(f"value must be greater than 0, got {value!r}") + return ivalue diff --git a/python/openai/tests/test_request_size.py b/python/openai/tests/test_request_size.py new file mode 100644 index 0000000000..5670d00a7a --- /dev/null +++ b/python/openai/tests/test_request_size.py @@ -0,0 +1,165 @@ +# Copyright 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions +# are met: +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of NVIDIA CORPORATION nor the names of its +# contributors may be used to endorse or promote products derived +# from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY +# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY +# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +import asyncio +import json +import os +import sys +from pathlib import Path + +import pytest +import tritonserver +from fastapi.testclient import TestClient + +sys.path.append( + os.path.join(str(Path(__file__).resolve().parent.parent), "openai_frontend") +) + +from frontend.fastapi.middleware.request_size import RequestSizeLimitMiddleware +from tests.utils import setup_fastapi_app, setup_server +from utils.utils import HTTP_DEFAULT_MAX_INPUT_SIZE + +_MODEL = "mock_llm" + +# All POST endpoints. +_ENDPOINTS = ( + "/v1/chat/completions", + "/v1/completions", + "/v1/embeddings", + f"/v1/models/{_MODEL}/load", + f"/v1/models/{_MODEL}/unload", +) + + +@pytest.fixture(scope="module") +def client(): + """FastApiFrontend backed by a real Triton server with mock_llm loaded.""" + model_repository = str(Path(__file__).parent / "test_models") + server = setup_server( + model_repository, + model_control_mode=tritonserver.ModelControlMode.EXPLICIT, + load_models=[_MODEL], + ) + try: + app = setup_fastapi_app(tokenizer="", server=server, backend=None) + with TestClient(app) as test_client: + yield test_client + finally: + server.stop() + + +def _assert_content_too_large(response, actual_bytes: int) -> None: + assert response.status_code == 413 + body = response.json() + assert set(body) == {"error"} + error = body["error"] + assert error["type"] == "invalid_request_error" + assert error["code"] == "content_too_large" + assert error["message"] == RequestSizeLimitMiddleware._oversized_request_message( + actual_bytes, HTTP_DEFAULT_MAX_INPUT_SIZE + ) + + +class TestRequestSizeLimitMiddleware: + @pytest.mark.parametrize("endpoint", _ENDPOINTS) + def test_body_at_limit_is_not_rejected(self, client, endpoint): + response = client.post(endpoint, content=b"x" * HTTP_DEFAULT_MAX_INPUT_SIZE) + assert response.status_code != 413 + + @pytest.mark.parametrize("endpoint", _ENDPOINTS) + def test_body_over_limit_is_rejected(self, client, endpoint): + over = HTTP_DEFAULT_MAX_INPUT_SIZE + 1 + response = client.post(endpoint, content=b"x" * over) + _assert_content_too_large(response, over) + + @pytest.mark.parametrize("endpoint", _ENDPOINTS) + def test_chunked_body_over_limit_is_rejected(self, client, endpoint): + over = HTTP_DEFAULT_MAX_INPUT_SIZE + 1 + + # httpx switches to chunked transfer when content is an Iterable[bytes]. + def chunks(): + yield b"x" * HTTP_DEFAULT_MAX_INPUT_SIZE + yield b"x" + + response = client.post(endpoint, content=chunks()) + _assert_content_too_large(response, over) + + def test_get_without_body_is_unaffected(self, client): + response = client.get(_ENDPOINTS[0]) + assert response.status_code == 405 + + +class TestContentLengthValidation: + """Stage 1 rejects malformed Content-Length with 400.""" + + def _run_with_content_length(self, raw_value: bytes) -> tuple[int, dict]: + captured: dict = {"status": None, "body": b""} + + async def app(scope, receive, send): + raise AssertionError("app must not be reached for invalid Content-Length") + + async def receive(): + raise AssertionError("receive() must not be called when Stage 1 rejects") + + async def send(message): + if message["type"] == "http.response.start": + captured["status"] = message["status"] + elif message["type"] == "http.response.body": + captured["body"] += message.get("body", b"") + + middleware = RequestSizeLimitMiddleware( + app=app, http_max_input_size=HTTP_DEFAULT_MAX_INPUT_SIZE + ) + scope = { + "type": "http", + "method": "POST", + "path": "/v1/chat/completions", + "headers": [(b"content-length", raw_value)], + } + asyncio.run(middleware(scope, receive, send)) + + assert captured["status"] is not None, "no response was sent" + return captured["status"], json.loads(captured["body"]) + + def _assert_invalid_content_length(self, status: int, body: dict): + assert status == 400 + assert set(body) == {"error"} + error = body["error"] + assert error["type"] == "invalid_request_error" + assert error["code"] == "invalid_content_length" + + @pytest.mark.parametrize("raw", [b"not-a-number", b""]) + def test_non_integer_content_length_rejected_with_400(self, raw): + status, body = self._run_with_content_length(raw) + self._assert_invalid_content_length(status, body) + assert "not an integer" in body["error"]["message"] + + @pytest.mark.parametrize("raw", [b"-1", b"-1024"]) + def test_negative_content_length_rejected_with_400(self, raw): + status, body = self._run_with_content_length(raw) + self._assert_invalid_content_length(status, body) + assert "non-negative" in body["error"]["message"] diff --git a/qa/L0_openai/test.sh b/qa/L0_openai/test.sh index 1862698fa7..5fd4daf239 100755 --- a/qa/L0_openai/test.sh +++ b/qa/L0_openai/test.sh @@ -123,12 +123,13 @@ function pre_test() { function run_test() { pushd openai/ TEST_LOG="test_openai.log" + TEST_XML="test_openai.xml" + TEST_LOG_MISTRAL="test_openai_mistral.log" + TEST_XML_MISTRAL="test_openai_mistral.xml" - # Capture error code without exiting to allow log collection set +e - pytest -s -v --junitxml=test_openai.xml tests/ 2>&1 > ${TEST_LOG} - if [ $? -ne 0 ]; then - cat ${TEST_LOG} + pytest -s -v --junitxml=${TEST_XML} tests/ 2>&1 | tee ${TEST_LOG} + if [ ${PIPESTATUS[0]} -ne 0 ]; then echo -e "\n***\n*** Test Failed\n***" RET=1 fi @@ -137,9 +138,9 @@ function run_test() { if [ "$RET" == "0" ]; then # rerun the tool calling tests with mistral model to cover the mistral tool call parser set +e - TEST_TOOL_CALL_PARSER="mistral" TEST_TOKENIZER="mistralai/Mistral-Nemo-Instruct-2407" pytest -s -v --junitxml=test_openai.xml tests/test_tool_calling.py 2>&1 > ${TEST_LOG} - if [ $? -ne 0 ]; then - cat ${TEST_LOG} + TEST_TOOL_CALL_PARSER="mistral" TEST_TOKENIZER="mistralai/Mistral-Nemo-Instruct-2407" \ + pytest -s -v --junitxml=${TEST_XML_MISTRAL} tests/test_tool_calling.py 2>&1 | tee ${TEST_LOG_MISTRAL} + if [ ${PIPESTATUS[0]} -ne 0 ]; then echo -e "\n***\n*** Test Failed\n***" RET=1 fi