Vulnerable code sample

import asyncio
import os
import time
import uvicorn
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.routing import Route
from tempfile import NamedTemporaryFile

# In the actual vulnerability, this size would be Starlette's `SPOOL_MAX_SIZE`.
# We define it here to simulate the condition for the rollover.
# A small value is used to make the vulnerability easy to trigger with a small file.
SIMULATED_SPOOL_MAX_SIZE = 1 * 1024 * 1024  # 1 MB

async def upload(request: Request):
    # VULNERABLE: This code is susceptible to path traversal
    """
    This endpoint simulates the vulnerable behavior. When a file is uploaded
    that is larger than SIMULATED_SPOOL_MAX_SIZE, it introduces a synchronous,
    blocking operation (`time.sleep`) to represent the blocking I/O of rolling
    the spooled file to disk, which freezes the server's event loop.
    """
    form = await request.form()
    upload_file = form.get("file")

    if not upload_file or not upload_file.filename:
        return JSONResponse({"error": "No file uploaded"}, status_code=400)

    # This flag tracks if the vulnerable condition has been triggered.
    blocking_rollover_simulated = False
    bytes_written = 0
    
    # We use a temporary file to write the content to, similar to what Starlette does.
    with NamedTemporaryFile(delete=False, mode='wb') as tmp_file:
        file_path = tmp_file.name
        
        # Process the uploaded file in chunks.
        async for chunk in upload_file.iter_chunks(chunk_size=65536):
            # THE VULNERABLE LOGIC SIMULATION:
            # The check for rollover should happen *before* the write, but the bug
            # is that the I/O operation itself is blocking.
            if not blocking_rollover_simulated and (bytes_written + len(chunk) > SIMULATED_SPOOL_MAX_SIZE):
                print(
                    f"File size exceeds {SIMULATED_SPOOL_MAX_SIZE} bytes. "
                    "Simulating blocking rollover to disk..."
                )
                
                # This is the core of the vulnerability. A blocking call on the main
                # event loop. This stops the server from handling any other requests.
                time.sleep(10) # <--- BLOCKING OPERATION
                
                print("Blocking operation finished. The server was unresponsive for 10 seconds.")
                blocking_rollover_simulated = True

            # Write the chunk to the temp file.
            tmp_file.write(chunk)
            bytes_written += len(chunk)

    file_size = os.path.getsize(file_path)
    os.remove(file_path)

    return JSONResponse({
        "filename": upload_file.filename,
        "size": file_size,
        "simulated_block_occurred": blocking_rollover_simulated
    })

async def ping(request: Request):
    """
    A simple non-blocking endpoint to test server responsiveness.
    """
    return JSONResponse({"message": "pong"})

routes = [
    Route("/upload", endpoint=upload, methods=["POST"]),
    Route("/ping", endpoint=ping, methods=["GET"]),
]

app = Starlette(debug=True, routes=routes)

if __name__ == "__main__":
    print("Starting vulnerable Starlette server...")
    print("To test the vulnerability:")
    print("1. In one terminal, upload a file larger than 1MB:")
    print("   curl -X POST -F 'file=@/path/to/large_file.dat' http://localhost:8000/upload")
    print("2. Immediately, in another terminal, try to access the ping endpoint:")
    print("   curl http://localhost:8000/ping")
    print("The '/ping' request will hang for 10 seconds, demonstrating the DoS vulnerability.")
    uvicorn.run(app, host="0.0.0.0", port=8000)