Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/middleware/request_size_limit_middleware.py: 36%

1import json 

2from collections.abc import Callable 

3from typing import Final 

4 

5from starlette.types import ASGIApp, Message, Receive, Scope, Send 

6 

7MaxRequestSizeGetter = Callable[[], int | float | None] 

8RequestSizeLimitEnabledGetter = Callable[[], bool] 

9 

10 

11class RequestEntityTooLarge(Exception): 

12 pass 

13 

14 

15class RequestSizeLimitMiddleware: 

16 """ 

17 Reject oversized requests before downstream auth/routes parse the body. 

18 

19 Content-Length can be rejected without reading any body bytes. Requests 

20 without Content-Length are counted as the ASGI stream is consumed, limiting 

21 memory exposure to the configured threshold plus the current chunk. 

22 """ 

23 

24 def __init__( 

25 self, 

26 app: ASGIApp, 

27 get_max_request_size_mb: MaxRequestSizeGetter, 

28 is_request_size_limit_enabled: RequestSizeLimitEnabledGetter, 

29 ) -> None: 

30 self.app = app 

31 self.get_max_request_size_mb = get_max_request_size_mb 

32 self.is_request_size_limit_enabled = is_request_size_limit_enabled 

33 

34 async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: 

35 if scope["type"] != "http": 

36 await self.app(scope, receive, send) 

37 return 

38 

39 max_request_size_mb: Final = self.get_max_request_size_mb() 

40 max_request_size_bytes: Final = _mb_to_bytes(max_request_size_mb) 

41 if max_request_size_bytes is None or not self.is_request_size_limit_enabled(): 41 ↛ 45line 41 didn't jump to line 45 because the condition on line 41 was always true

42 await self.app(scope, receive, send) 

43 return 

44 

45 content_length: Final = _get_content_length(scope=scope) 

46 if content_length is not None and content_length > max_request_size_bytes: 

47 await _send_request_too_large(send=send, max_request_size_mb=max_request_size_mb) 

48 return 

49 

50 received_body_bytes = 0 

51 response_started = False 

52 

53 async def limited_receive() -> Message: 

54 nonlocal received_body_bytes 

55 

56 message: Final = await receive() 

57 if message["type"] != "http.request": 

58 return message 

59 

60 received_body_bytes += len(message.get("body", b"")) 

61 if received_body_bytes > max_request_size_bytes: 

62 raise RequestEntityTooLarge 

63 return message 

64 

65 async def tracking_send(message: Message) -> None: 

66 nonlocal response_started 

67 

68 if message["type"] == "http.response.start": 

69 response_started = True 

70 await send(message) 

71 

72 try: 

73 await self.app(scope, limited_receive, tracking_send) 

74 except RequestEntityTooLarge: 

75 if response_started: 

76 raise 

77 await _send_request_too_large(send=send, max_request_size_mb=max_request_size_mb) 

78 

79 

80def _mb_to_bytes(max_request_size_mb: float | None) -> int | None: 

81 if max_request_size_mb is None: 81 ↛ 83line 81 didn't jump to line 83 because the condition on line 81 was always true

82 return None 

83 if max_request_size_mb <= 0: 

84 return None 

85 return int(max_request_size_mb * 1024 * 1024) 

86 

87 

88def _get_content_length(scope: Scope) -> int | None: 

89 headers: Final = dict(scope.get("headers") or []) 

90 raw_content_length: Final = headers.get(b"content-length") 

91 if raw_content_length is None: 

92 return None 

93 

94 try: 

95 return int(raw_content_length) 

96 except ValueError: 

97 return None 

98 

99 

100async def _send_request_too_large( 

101 send: Send, 

102 max_request_size_mb: float | None, 

103) -> None: 

104 body: Final = json.dumps( 

105 {"error": f"Request size is too large. Max size is {max_request_size_mb} MB"}, 

106 separators=(",", ":"), 

107 ).encode("utf-8") 

108 await send( 

109 { 

110 "type": "http.response.start", 

111 "status": 413, 

112 "headers": [ 

113 (b"content-type", b"application/json"), 

114 (b"content-length", str(len(body)).encode("latin-1")), 

115 ], 

116 } 

117 ) 

118 await send({"type": "http.response.body", "body": body, "more_body": False})