content size limiting middleware for asgi applications args app asgi application asgi application maxcontentsize optional the maximum content size allowed in bytes none for no limit content size limiting middleware for asgi applications args app asgi application asgi application max_content_size optional the maximum content size allowed in bytes none for no limit pragma no cover
from pathlib import Path
from typing import Optional

from fastapi import APIRouter, FastAPI, File, UploadFile
from fastapi.exceptions import HTTPException
from fastapi.testclient import TestClient

app = FastAPI()

router = APIRouter()


class ContentSizeLimitMiddleware:
    

    def __init__(self, app: APIRouter, max_content_size: Optional[int] = None):
        self.app = app
        self.max_content_size = max_content_size

    def receive_wrapper(self, receive):
        received = 0

        async def inner():
            nonlocal received
            message = await receive()
            if message["type"] != "http.request":
                return message  

            body_len = len(message.get("body", b""))
            received += body_len
            if received > self.max_content_size:
                raise HTTPException(
                    422,
                    detail={
                        "name": "ContentSizeLimitExceeded",
                        "code": 999,
                        "message": "File limit exceeded",
                    },
                )
            return message

        return inner

    async def __call__(self, scope, receive, send):
        if scope["type"] != "http" or self.max_content_size is None:
            await self.app(scope, receive, send)
            return

        wrapper = self.receive_wrapper(receive)
        await self.app(scope, wrapper, send)


@router.post("/middleware")
def run_middleware(file: UploadFile = File(..., description="Big File")):
    return {"message": "OK"}


app.include_router(router)
app.add_middleware(ContentSizeLimitMiddleware, max_content_size=2**8)


client = TestClient(app)


def test_custom_middleware_exception(tmp_path: Path):
    default_pydantic_max_size = 2**16
    path = tmp_path / "test.txt"
    path.write_bytes(b"x" * (default_pydantic_max_size + 1))

    with client:
        with open(path, "rb") as file:
            response = client.post("/middleware", files={"file": file})
        assert response.status_code == 422, response.text
        assert response.json() == {
            "detail": {
                "name": "ContentSizeLimitExceeded",
                "code": 999,
                "message": "File limit exceeded",
            }
        }


def test_custom_middleware_exception_not_raised(tmp_path: Path):
    path = tmp_path / "test.txt"
    path.write_bytes(b"<file content>")

    with client:
        with open(path, "rb") as file:
            response = client.post("/middleware", files={"file": file})
        assert response.status_code == 200, response.text
        assert response.json() == {"message": "OK"}