Coverage for gco/services/request_size_middleware.py: 82.43%

1"""Shared request-body size enforcement for GCO FastAPI services.""" 

2 

3from __future__ import annotations 

4 

5from collections import deque 

6 

7from starlette.responses import JSONResponse 

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

9 

10DEFAULT_MAX_REQUEST_BODY_BYTES = 1_048_576 

11 

12 

13class RequestSizeLimitMiddleware: 

14 """Reject request bodies larger than the configured byte limit. 

15 

16 ``Content-Length`` is only an early-rejection optimization. Every accepted 

17 request is read from the ASGI receive channel and counted before it reaches 

18 authentication or a route handler, so a missing, malformed, negative, or 

19 deliberately under-reported header cannot bypass the limit. Buffered 

20 messages are replayed unchanged to preserve the exact bytes used by HMAC 

21 validation and downstream parsing. 

22 """ 

23 

24 def __init__(self, app: ASGIApp, max_body_bytes: int = DEFAULT_MAX_REQUEST_BODY_BYTES) -> None: 

25 if max_body_bytes < 0: 25 ↛ 26line 25 didn't jump to line 26 because the condition on line 25 was never true

26 raise ValueError("max_body_bytes must be non-negative") 

27 self.app = app 

28 self.max_body_bytes = max_body_bytes 

29 

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

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

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

33 return 

34 

35 if self._declared_size_exceeds_limit(scope): 

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

37 return 

38 

39 buffered_messages: deque[Message] = deque() 

40 received_bytes = 0 

41 while True: 

42 message = await receive() 

43 message_type = message["type"] 

44 

45 if message_type == "http.request": 45 ↛ 54line 45 didn't jump to line 54 because the condition on line 45 was always true

46 received_bytes += len(message.get("body", b"")) 

47 if received_bytes > self.max_body_bytes: 

48 await self._too_large_response()(scope, receive, send) 

49 return 

50 buffered_messages.append(message) 

51 if not message.get("more_body", False): 51 ↛ 41line 51 didn't jump to line 41 because the condition on line 51 was always true

52 break 

53 else: 

54 buffered_messages.append(message) 

55 if message_type == "http.disconnect": 

56 break 

57 

58 async def replay_receive() -> Message: 

59 if buffered_messages: 59 ↛ 61line 59 didn't jump to line 61 because the condition on line 59 was always true

60 return buffered_messages.popleft() 

61 return await receive() 

62 

63 await self.app(scope, replay_receive, send) 

64 

65 def _declared_size_exceeds_limit(self, scope: Scope) -> bool: 

66 """Return true if any valid Content-Length value is already too large.""" 

67 for name, raw_value in scope.get("headers", []): 

68 if name.lower() != b"content-length": 

69 continue 

70 try: 

71 declared_size = int(raw_value.decode("ascii")) 

72 except UnicodeDecodeError, ValueError: 

73 continue 

74 if declared_size > self.max_body_bytes: 

75 return True 

76 return False 

77 

78 def _too_large_response(self) -> JSONResponse: 

79 return JSONResponse( 

80 status_code=413, 

81 content={"detail": f"Request body exceeds maximum size of {self.max_body_bytes} bytes"}, 

82 )