base.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244
  1. from __future__ import annotations
  2. from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Mapping, MutableMapping
  3. from typing import Any, TypeVar
  4. import anyio
  5. from starlette._utils import collapse_excgroups
  6. from starlette.requests import ClientDisconnect, Request
  7. from starlette.responses import Response
  8. from starlette.types import ASGIApp, Message, Receive, Scope, Send
  9. RequestResponseEndpoint = Callable[[Request], Awaitable[Response]]
  10. DispatchFunction = Callable[[Request, RequestResponseEndpoint], Awaitable[Response]]
  11. BodyStreamGenerator = AsyncGenerator[bytes | MutableMapping[str, Any], None]
  12. AsyncContentStream = AsyncIterable[str | bytes | memoryview | MutableMapping[str, Any]]
  13. T = TypeVar("T")
  14. class _CachedRequest(Request):
  15. """
  16. If the user calls Request.body() from their dispatch function
  17. we cache the entire request body in memory and pass that to downstream middlewares,
  18. but if they call Request.stream() then all we do is send an
  19. empty body so that downstream things don't hang forever.
  20. """
  21. def __init__(self, scope: Scope, receive: Receive):
  22. super().__init__(scope, receive)
  23. self._wrapped_rcv_disconnected = False
  24. self._wrapped_rcv_consumed = False
  25. self._wrapped_rc_stream = self.stream()
  26. async def wrapped_receive(self) -> Message:
  27. # wrapped_rcv state 1: disconnected
  28. if self._wrapped_rcv_disconnected:
  29. # we've already sent a disconnect to the downstream app
  30. # we don't need to wait to get another one
  31. # (although most ASGI servers will just keep sending it)
  32. return {"type": "http.disconnect"}
  33. # wrapped_rcv state 1: consumed but not yet disconnected
  34. if self._wrapped_rcv_consumed:
  35. # since the downstream app has consumed us all that is left
  36. # is to send it a disconnect
  37. if self._is_disconnected:
  38. # the middleware has already seen the disconnect
  39. # since we know the client is disconnected no need to wait
  40. # for the message
  41. self._wrapped_rcv_disconnected = True
  42. return {"type": "http.disconnect"}
  43. # we don't know yet if the client is disconnected or not
  44. # so we'll wait until we get that message
  45. msg = await self.receive()
  46. if msg["type"] != "http.disconnect": # pragma: no cover
  47. # at this point a disconnect is all that we should be receiving
  48. # if we get something else, things went wrong somewhere
  49. raise RuntimeError(f"Unexpected message received: {msg['type']}")
  50. self._wrapped_rcv_disconnected = True
  51. return msg
  52. # wrapped_rcv state 3: not yet consumed
  53. if getattr(self, "_body", None) is not None:
  54. # body() was called, we return it even if the client disconnected
  55. self._wrapped_rcv_consumed = True
  56. return {
  57. "type": "http.request",
  58. "body": self._body,
  59. "more_body": False,
  60. }
  61. elif self._stream_consumed:
  62. # stream() was called to completion
  63. # return an empty body so that downstream apps don't hang
  64. # waiting for a disconnect
  65. self._wrapped_rcv_consumed = True
  66. return {
  67. "type": "http.request",
  68. "body": b"",
  69. "more_body": False,
  70. }
  71. else:
  72. # body() was never called and stream() wasn't consumed
  73. try:
  74. stream = self.stream()
  75. chunk = await stream.__anext__()
  76. self._wrapped_rcv_consumed = self._stream_consumed
  77. return {
  78. "type": "http.request",
  79. "body": chunk,
  80. "more_body": not self._stream_consumed,
  81. }
  82. except ClientDisconnect:
  83. self._wrapped_rcv_disconnected = True
  84. return {"type": "http.disconnect"}
  85. class BaseHTTPMiddleware:
  86. def __init__(self, app: ASGIApp, dispatch: DispatchFunction | None = None) -> None:
  87. self.app = app
  88. self.dispatch_func = self.dispatch if dispatch is None else dispatch
  89. async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
  90. if scope["type"] != "http":
  91. await self.app(scope, receive, send)
  92. return
  93. request = _CachedRequest(scope, receive)
  94. wrapped_receive = request.wrapped_receive
  95. response_sent = anyio.Event()
  96. app_exc: Exception | None = None
  97. exception_already_raised = False
  98. async def call_next(request: Request) -> Response:
  99. async def receive_or_disconnect() -> Message:
  100. if response_sent.is_set():
  101. return {"type": "http.disconnect"}
  102. async with anyio.create_task_group() as task_group:
  103. async def wrap(func: Callable[[], Awaitable[T]]) -> T:
  104. result = await func()
  105. task_group.cancel_scope.cancel()
  106. return result
  107. task_group.start_soon(wrap, response_sent.wait)
  108. message = await wrap(wrapped_receive)
  109. if response_sent.is_set():
  110. return {"type": "http.disconnect"}
  111. return message
  112. async def send_no_error(message: Message) -> None:
  113. try:
  114. await send_stream.send(message)
  115. except anyio.BrokenResourceError:
  116. # recv_stream has been closed, i.e. response_sent has been set.
  117. return
  118. async def coro() -> None:
  119. nonlocal app_exc
  120. with send_stream:
  121. try:
  122. await self.app(scope, receive_or_disconnect, send_no_error)
  123. except Exception as exc:
  124. app_exc = exc
  125. task_group.start_soon(coro)
  126. try:
  127. message = await recv_stream.receive()
  128. info = message.get("info", None)
  129. if message["type"] == "http.response.debug" and info is not None:
  130. message = await recv_stream.receive()
  131. except anyio.EndOfStream:
  132. if app_exc is not None:
  133. nonlocal exception_already_raised
  134. exception_already_raised = True
  135. # Prevent `anyio.EndOfStream` from polluting app exception context.
  136. # If both cause and context are None then the context is suppressed
  137. # and `anyio.EndOfStream` is not present in the exception traceback.
  138. # If exception cause is not None then it is propagated with
  139. # reraising here.
  140. # If exception has no cause but has context set then the context is
  141. # propagated as a cause with the reraise. This is necessary in order
  142. # to prevent `anyio.EndOfStream` from polluting the exception
  143. # context.
  144. raise app_exc from app_exc.__cause__ or app_exc.__context__
  145. raise RuntimeError("No response returned.")
  146. assert message["type"] == "http.response.start"
  147. async def body_stream() -> BodyStreamGenerator:
  148. async for message in recv_stream:
  149. if message["type"] == "http.response.pathsend":
  150. yield message
  151. break
  152. assert message["type"] == "http.response.body", f"Unexpected message: {message}"
  153. body = message.get("body", b"")
  154. if body:
  155. yield body
  156. if not message.get("more_body", False):
  157. break
  158. response = _StreamingResponse(status_code=message["status"], content=body_stream(), info=info)
  159. response.raw_headers = message["headers"]
  160. return response
  161. streams: anyio.create_memory_object_stream[Message] = anyio.create_memory_object_stream()
  162. send_stream, recv_stream = streams
  163. with recv_stream, send_stream, collapse_excgroups():
  164. async with anyio.create_task_group() as task_group:
  165. response = await self.dispatch_func(request, call_next)
  166. await response(scope, wrapped_receive, send)
  167. response_sent.set()
  168. recv_stream.close()
  169. if app_exc is not None and not exception_already_raised:
  170. raise app_exc
  171. async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
  172. raise NotImplementedError() # pragma: no cover
  173. class _StreamingResponse(Response):
  174. def __init__(
  175. self,
  176. content: AsyncContentStream,
  177. status_code: int = 200,
  178. headers: Mapping[str, str] | None = None,
  179. media_type: str | None = None,
  180. info: Mapping[str, Any] | None = None,
  181. ) -> None:
  182. self.info = info
  183. self.body_iterator = content
  184. self.status_code = status_code
  185. self.media_type = media_type
  186. self.init_headers(headers)
  187. self.background = None
  188. async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
  189. if self.info is not None:
  190. await send({"type": "http.response.debug", "info": self.info})
  191. await send(
  192. {
  193. "type": "http.response.start",
  194. "status": self.status_code,
  195. "headers": self.raw_headers,
  196. }
  197. )
  198. should_close_body = True
  199. async for chunk in self.body_iterator:
  200. if isinstance(chunk, dict):
  201. # We got an ASGI message which is not response body (eg: pathsend)
  202. should_close_body = False
  203. await send(chunk)
  204. continue
  205. await send({"type": "http.response.body", "body": chunk, "more_body": True})
  206. if should_close_body:
  207. await send({"type": "http.response.body", "body": b"", "more_body": False})
  208. if self.background:
  209. await self.background()