cors.py 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172
  1. from __future__ import annotations
  2. import functools
  3. import re
  4. import typing
  5. from starlette.datastructures import Headers, MutableHeaders
  6. from starlette.responses import PlainTextResponse, Response
  7. from starlette.types import ASGIApp, Message, Receive, Scope, Send
  8. ALL_METHODS = ("DELETE", "GET", "HEAD", "OPTIONS", "PATCH", "POST", "PUT")
  9. SAFELISTED_HEADERS = {"Accept", "Accept-Language", "Content-Language", "Content-Type"}
  10. class CORSMiddleware:
  11. def __init__(
  12. self,
  13. app: ASGIApp,
  14. allow_origins: typing.Sequence[str] = (),
  15. allow_methods: typing.Sequence[str] = ("GET",),
  16. allow_headers: typing.Sequence[str] = (),
  17. allow_credentials: bool = False,
  18. allow_origin_regex: str | None = None,
  19. expose_headers: typing.Sequence[str] = (),
  20. max_age: int = 600,
  21. ) -> None:
  22. if "*" in allow_methods:
  23. allow_methods = ALL_METHODS
  24. compiled_allow_origin_regex = None
  25. if allow_origin_regex is not None:
  26. compiled_allow_origin_regex = re.compile(allow_origin_regex)
  27. allow_all_origins = "*" in allow_origins
  28. allow_all_headers = "*" in allow_headers
  29. preflight_explicit_allow_origin = not allow_all_origins or allow_credentials
  30. simple_headers = {}
  31. if allow_all_origins:
  32. simple_headers["Access-Control-Allow-Origin"] = "*"
  33. if allow_credentials:
  34. simple_headers["Access-Control-Allow-Credentials"] = "true"
  35. if expose_headers:
  36. simple_headers["Access-Control-Expose-Headers"] = ", ".join(expose_headers)
  37. preflight_headers = {}
  38. if preflight_explicit_allow_origin:
  39. # The origin value will be set in preflight_response() if it is allowed.
  40. preflight_headers["Vary"] = "Origin"
  41. else:
  42. preflight_headers["Access-Control-Allow-Origin"] = "*"
  43. preflight_headers.update(
  44. {
  45. "Access-Control-Allow-Methods": ", ".join(allow_methods),
  46. "Access-Control-Max-Age": str(max_age),
  47. }
  48. )
  49. allow_headers = sorted(SAFELISTED_HEADERS | set(allow_headers))
  50. if allow_headers and not allow_all_headers:
  51. preflight_headers["Access-Control-Allow-Headers"] = ", ".join(allow_headers)
  52. if allow_credentials:
  53. preflight_headers["Access-Control-Allow-Credentials"] = "true"
  54. self.app = app
  55. self.allow_origins = allow_origins
  56. self.allow_methods = allow_methods
  57. self.allow_headers = [h.lower() for h in allow_headers]
  58. self.allow_all_origins = allow_all_origins
  59. self.allow_all_headers = allow_all_headers
  60. self.preflight_explicit_allow_origin = preflight_explicit_allow_origin
  61. self.allow_origin_regex = compiled_allow_origin_regex
  62. self.simple_headers = simple_headers
  63. self.preflight_headers = preflight_headers
  64. async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
  65. if scope["type"] != "http": # pragma: no cover
  66. await self.app(scope, receive, send)
  67. return
  68. method = scope["method"]
  69. headers = Headers(scope=scope)
  70. origin = headers.get("origin")
  71. if origin is None:
  72. await self.app(scope, receive, send)
  73. return
  74. if method == "OPTIONS" and "access-control-request-method" in headers:
  75. response = self.preflight_response(request_headers=headers)
  76. await response(scope, receive, send)
  77. return
  78. await self.simple_response(scope, receive, send, request_headers=headers)
  79. def is_allowed_origin(self, origin: str) -> bool:
  80. if self.allow_all_origins:
  81. return True
  82. if self.allow_origin_regex is not None and self.allow_origin_regex.fullmatch(origin):
  83. return True
  84. return origin in self.allow_origins
  85. def preflight_response(self, request_headers: Headers) -> Response:
  86. requested_origin = request_headers["origin"]
  87. requested_method = request_headers["access-control-request-method"]
  88. requested_headers = request_headers.get("access-control-request-headers")
  89. headers = dict(self.preflight_headers)
  90. failures = []
  91. if self.is_allowed_origin(origin=requested_origin):
  92. if self.preflight_explicit_allow_origin:
  93. # The "else" case is already accounted for in self.preflight_headers
  94. # and the value would be "*".
  95. headers["Access-Control-Allow-Origin"] = requested_origin
  96. else:
  97. failures.append("origin")
  98. if requested_method not in self.allow_methods:
  99. failures.append("method")
  100. # If we allow all headers, then we have to mirror back any requested
  101. # headers in the response.
  102. if self.allow_all_headers and requested_headers is not None:
  103. headers["Access-Control-Allow-Headers"] = requested_headers
  104. elif requested_headers is not None:
  105. for header in [h.lower() for h in requested_headers.split(",")]:
  106. if header.strip() not in self.allow_headers:
  107. failures.append("headers")
  108. break
  109. # We don't strictly need to use 400 responses here, since its up to
  110. # the browser to enforce the CORS policy, but its more informative
  111. # if we do.
  112. if failures:
  113. failure_text = "Disallowed CORS " + ", ".join(failures)
  114. return PlainTextResponse(failure_text, status_code=400, headers=headers)
  115. return PlainTextResponse("OK", status_code=200, headers=headers)
  116. async def simple_response(self, scope: Scope, receive: Receive, send: Send, request_headers: Headers) -> None:
  117. send = functools.partial(self.send, send=send, request_headers=request_headers)
  118. await self.app(scope, receive, send)
  119. async def send(self, message: Message, send: Send, request_headers: Headers) -> None:
  120. if message["type"] != "http.response.start":
  121. await send(message)
  122. return
  123. message.setdefault("headers", [])
  124. headers = MutableHeaders(scope=message)
  125. headers.update(self.simple_headers)
  126. origin = request_headers["Origin"]
  127. has_cookie = "cookie" in request_headers
  128. # If request includes any cookie headers, then we must respond
  129. # with the specific origin instead of '*'.
  130. if self.allow_all_origins and has_cookie:
  131. self.allow_explicit_origin(headers, origin)
  132. # If we only allow specific origins, then we have to mirror back
  133. # the Origin header in the response.
  134. elif not self.allow_all_origins and self.is_allowed_origin(origin=origin):
  135. self.allow_explicit_origin(headers, origin)
  136. await send(message)
  137. @staticmethod
  138. def allow_explicit_origin(headers: MutableHeaders, origin: str) -> None:
  139. headers["Access-Control-Allow-Origin"] = origin
  140. headers.add_vary_header("Origin")