Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/starlette/middleware/body_limit.py: 30%

1from __future__ import annotations 

2 

3from typing import cast 

4 

5from starlette.datastructures import Headers 

6from starlette.exceptions import HTTPException 

7from starlette.responses import PlainTextResponse 

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

9 

10MAX_BODY_SIZE_SCOPE_KEY = "starlette.max_body_size" 

11_BODY_LIMIT_RESPONDER_SCOPE_KEY = "starlette._body_limit_responder" 

12 

13 

14class _Missing: 

15 __slots__ = () 

16 

17 

18_MISSING = _Missing() 

19 

20 

21class _RequestBodyTooLarge(HTTPException): 

22 def __init__(self) -> None: 

23 super().__init__(status_code=413, detail="Content Too Large") 

24 

25 

26class _RequestBodyLimitResponseSent(Exception): 

27 pass 

28 

29 

30class RequestBodyLimitMiddleware: 

31 """Limit the total size of an HTTP request body.""" 

32 

33 def __init__(self, app: ASGIApp, max_body_size: int) -> None: 

34 self.app = app 

35 self.max_body_size = max_body_size 

36 

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

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

39 return await self.app(scope, receive, send) 

40 

41 responder = RequestBodyLimitResponder(self.app, self.max_body_size) 

42 await responder(scope, receive, send) 

43 

44 

45class RequestBodyLimitResponder: 

46 def __init__(self, app: ASGIApp, max_body_size: int) -> None: 

47 self.app = app 

48 self.max_body_size = max_body_size 

49 self._scope: Scope | None = None 

50 self._receive: Receive | None = None 

51 self._send: Send | None = None 

52 self.content_length: int | None = None 

53 self.total_size = 0 

54 self.response_started = False 

55 

56 @property 

57 def scope(self) -> Scope: 

58 assert self._scope is not None 

59 return self._scope 

60 

61 @property 

62 def receive(self) -> Receive: 

63 assert self._receive is not None 

64 return self._receive 

65 

66 @property 

67 def send(self) -> Send: 

68 assert self._send is not None 

69 return self._send 

70 

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

72 previous_scope_limit = cast(int | _Missing, scope.get(MAX_BODY_SIZE_SCOPE_KEY, _MISSING)) 

73 scope[MAX_BODY_SIZE_SCOPE_KEY] = self.max_body_size 

74 

75 active_responder = cast(RequestBodyLimitResponder | None, scope.get(_BODY_LIMIT_RESPONDER_SCOPE_KEY)) 

76 if active_responder is not None: 

77 active_responder.max_body_size = self.max_body_size 

78 if active_responder.total_size > active_responder.max_body_size: 

79 raise _RequestBodyTooLarge 

80 return await self.app(scope, receive, send) 

81 

82 self._scope = scope 

83 self._receive = receive 

84 self._send = send 

85 self.content_length = _get_content_length(scope) 

86 scope[_BODY_LIMIT_RESPONDER_SCOPE_KEY] = self 

87 

88 try: 

89 await self.app(scope, self.receive_with_limit, self.send_with_limit) 

90 except _RequestBodyTooLarge: 

91 if self.response_started: 

92 raise 

93 response = PlainTextResponse("Content Too Large", status_code=413) 

94 await response(scope, receive, send) 

95 except _RequestBodyLimitResponseSent: 

96 pass 

97 finally: 

98 scope.pop(_BODY_LIMIT_RESPONDER_SCOPE_KEY, None) 

99 if isinstance(previous_scope_limit, _Missing): 

100 scope.pop(MAX_BODY_SIZE_SCOPE_KEY, None) 

101 else: 

102 scope[MAX_BODY_SIZE_SCOPE_KEY] = previous_scope_limit 

103 

104 async def receive_with_limit(self) -> Message: 

105 if self.content_length is not None and self.content_length > self.max_body_size: 

106 raise _RequestBodyTooLarge 

107 

108 message = await self.receive() 

109 if message["type"] == "http.request": 

110 self.total_size += len(message.get("body", b"")) 

111 if self.total_size > self.max_body_size: 

112 raise _RequestBodyTooLarge 

113 return message 

114 

115 async def send_with_limit(self, message: Message) -> None: 

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

117 self.response_started = True 

118 if self.content_length is not None and self.content_length > self.max_body_size: 

119 response = PlainTextResponse("Content Too Large", status_code=413) 

120 await response(self.scope, self.receive, self.send) 

121 raise _RequestBodyLimitResponseSent 

122 await self.send(message) 

123 

124 

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

126 content_length = Headers(scope=scope).get("content-length") 

127 if content_length is None: 

128 return None 

129 try: 

130 return int(content_length) 

131 except ValueError: 

132 return None