fastapi

1
from pathlib import Path
2
from typing import Optional
3

4
from fastapi import APIRouter, FastAPI, File, UploadFile
5
from fastapi.exceptions import HTTPException
6
from fastapi.testclient import TestClient
7

8
app = FastAPI()
9

10
router = APIRouter()
11

12

13
class ContentSizeLimitMiddleware:
14
    """Content size limiting middleware for ASGI applications
15
    Args:
16
      app (ASGI application): ASGI application
17
      max_content_size (optional): the maximum content size allowed in bytes, None for no limit
18
    """
19

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

24
    def receive_wrapper(self, receive):
25
        received = 0
26

27
        async def inner():
28
            nonlocal received
29
            message = await receive()
30
            if message["type"] != "http.request":
31
                return message  # pragma: no cover
32

33
            body_len = len(message.get("body", b""))
34
            received += body_len
35
            if received > self.max_content_size:
36
                raise HTTPException(
37
                    422,
38
                    detail={
39
                        "name": "ContentSizeLimitExceeded",
40
                        "code": 999,
41
                        "message": "File limit exceeded",
42
                    },
43
                )
44
            return message
45

46
        return inner
47

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

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

56

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

61

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

65

66
client = TestClient(app)
67

68

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

74
    with client:
75
        with open(path, "rb") as file:
76
            response = client.post("/middleware", files={"file": file})
77
        assert response.status_code == 422, response.text
78
        assert response.json() == {
79
            "detail": {
80
                "name": "ContentSizeLimitExceeded",
81
                "code": 999,
82
                "message": "File limit exceeded",
83
            }
84
        }
85

86

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

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