responses.py 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568
  1. from __future__ import annotations
  2. import hashlib
  3. import http.cookies
  4. import json
  5. import os
  6. import stat
  7. import sys
  8. from collections.abc import AsyncIterable, Awaitable, Callable, Iterable, Mapping, Sequence
  9. from datetime import datetime
  10. from email.utils import format_datetime, formatdate
  11. from functools import partial
  12. from mimetypes import guess_type
  13. from secrets import token_hex
  14. from typing import Any, Literal
  15. from urllib.parse import quote
  16. import anyio
  17. import anyio.to_thread
  18. from starlette._utils import collapse_excgroups
  19. from starlette.background import BackgroundTask
  20. from starlette.concurrency import iterate_in_threadpool
  21. from starlette.datastructures import URL, Headers, MutableHeaders
  22. from starlette.requests import ClientDisconnect
  23. from starlette.types import Message, Receive, Scope, Send
  24. class Response:
  25. media_type = None
  26. charset = "utf-8"
  27. def __init__(
  28. self,
  29. content: Any = None,
  30. status_code: int = 200,
  31. headers: Mapping[str, str] | None = None,
  32. media_type: str | None = None,
  33. background: BackgroundTask | None = None,
  34. ) -> None:
  35. self.status_code = status_code
  36. if media_type is not None:
  37. self.media_type = media_type
  38. self.background = background
  39. self.body = self.render(content)
  40. self.init_headers(headers)
  41. def render(self, content: Any) -> bytes | memoryview:
  42. if content is None:
  43. return b""
  44. if isinstance(content, bytes | memoryview):
  45. return content
  46. return content.encode(self.charset) # type: ignore
  47. def init_headers(self, headers: Mapping[str, str] | None = None) -> None:
  48. if headers is None:
  49. raw_headers: list[tuple[bytes, bytes]] = []
  50. populate_content_length = True
  51. populate_content_type = True
  52. else:
  53. raw_headers = [(k.lower().encode("latin-1"), v.encode("latin-1")) for k, v in headers.items()]
  54. keys = [h[0] for h in raw_headers]
  55. populate_content_length = b"content-length" not in keys
  56. populate_content_type = b"content-type" not in keys
  57. body = getattr(self, "body", None)
  58. if (
  59. body is not None
  60. and populate_content_length
  61. and not (self.status_code < 200 or self.status_code in (204, 304))
  62. ):
  63. content_length = str(len(body))
  64. raw_headers.append((b"content-length", content_length.encode("latin-1")))
  65. content_type = self.media_type
  66. if content_type is not None and populate_content_type:
  67. if content_type.startswith("text/") and "charset=" not in content_type.lower():
  68. content_type += "; charset=" + self.charset
  69. raw_headers.append((b"content-type", content_type.encode("latin-1")))
  70. self.raw_headers = raw_headers
  71. @property
  72. def headers(self) -> MutableHeaders:
  73. if not hasattr(self, "_headers"):
  74. self._headers = MutableHeaders(raw=self.raw_headers)
  75. return self._headers
  76. def set_cookie(
  77. self,
  78. key: str,
  79. value: str = "",
  80. max_age: int | None = None,
  81. expires: datetime | str | int | None = None,
  82. path: str | None = "/",
  83. domain: str | None = None,
  84. secure: bool = False,
  85. httponly: bool = False,
  86. samesite: Literal["lax", "strict", "none"] | None = "lax",
  87. partitioned: bool = False,
  88. ) -> None:
  89. cookie: http.cookies.BaseCookie[str] = http.cookies.SimpleCookie()
  90. cookie[key] = value
  91. if max_age is not None:
  92. cookie[key]["max-age"] = max_age
  93. if expires is not None:
  94. if isinstance(expires, datetime):
  95. cookie[key]["expires"] = format_datetime(expires, usegmt=True)
  96. else:
  97. cookie[key]["expires"] = expires
  98. if path is not None:
  99. cookie[key]["path"] = path
  100. if domain is not None:
  101. cookie[key]["domain"] = domain
  102. if secure:
  103. cookie[key]["secure"] = True
  104. if httponly:
  105. cookie[key]["httponly"] = True
  106. if samesite is not None:
  107. assert samesite.lower() in [
  108. "strict",
  109. "lax",
  110. "none",
  111. ], "samesite must be either 'strict', 'lax' or 'none'"
  112. cookie[key]["samesite"] = samesite
  113. if partitioned:
  114. if sys.version_info < (3, 14):
  115. raise ValueError("Partitioned cookies are only supported in Python 3.14 and above.") # pragma: no cover
  116. cookie[key]["partitioned"] = True # pragma: no cover
  117. cookie_val = cookie.output(header="").strip()
  118. self.raw_headers.append((b"set-cookie", cookie_val.encode("latin-1")))
  119. def delete_cookie(
  120. self,
  121. key: str,
  122. path: str = "/",
  123. domain: str | None = None,
  124. secure: bool = False,
  125. httponly: bool = False,
  126. samesite: Literal["lax", "strict", "none"] | None = "lax",
  127. ) -> None:
  128. self.set_cookie(
  129. key,
  130. max_age=0,
  131. expires=0,
  132. path=path,
  133. domain=domain,
  134. secure=secure,
  135. httponly=httponly,
  136. samesite=samesite,
  137. )
  138. def _wrap_websocket_denial_send(self, send: Send) -> Send:
  139. async def wrapped(message: Message) -> None:
  140. message_type = message["type"]
  141. if message_type in {"http.response.start", "http.response.body"}: # pragma: no branch
  142. message = {**message, "type": "websocket." + message_type}
  143. await send(message)
  144. return wrapped
  145. async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
  146. if scope["type"] == "websocket":
  147. send = self._wrap_websocket_denial_send(send)
  148. await send({"type": "http.response.start", "status": self.status_code, "headers": self.raw_headers})
  149. await send({"type": "http.response.body", "body": self.body})
  150. if self.background is not None:
  151. await self.background()
  152. class HTMLResponse(Response):
  153. media_type = "text/html"
  154. class PlainTextResponse(Response):
  155. media_type = "text/plain"
  156. class JSONResponse(Response):
  157. media_type = "application/json"
  158. def __init__(
  159. self,
  160. content: Any,
  161. status_code: int = 200,
  162. headers: Mapping[str, str] | None = None,
  163. media_type: str | None = None,
  164. background: BackgroundTask | None = None,
  165. ) -> None:
  166. super().__init__(content, status_code, headers, media_type, background)
  167. def render(self, content: Any) -> bytes:
  168. return json.dumps(
  169. content,
  170. ensure_ascii=False,
  171. allow_nan=False,
  172. indent=None,
  173. separators=(",", ":"),
  174. ).encode("utf-8")
  175. class RedirectResponse(Response):
  176. def __init__(
  177. self,
  178. url: str | URL,
  179. status_code: int = 307,
  180. headers: Mapping[str, str] | None = None,
  181. background: BackgroundTask | None = None,
  182. ) -> None:
  183. super().__init__(content=b"", status_code=status_code, headers=headers, background=background)
  184. self.headers["location"] = quote(str(url), safe=":/%#?=@[]!$&'()*+,;")
  185. Content = str | bytes | memoryview
  186. SyncContentStream = Iterable[Content]
  187. AsyncContentStream = AsyncIterable[Content]
  188. ContentStream = AsyncContentStream | SyncContentStream
  189. class StreamingResponse(Response):
  190. body_iterator: AsyncContentStream
  191. def __init__(
  192. self,
  193. content: ContentStream,
  194. status_code: int = 200,
  195. headers: Mapping[str, str] | None = None,
  196. media_type: str | None = None,
  197. background: BackgroundTask | None = None,
  198. ) -> None:
  199. if isinstance(content, AsyncIterable):
  200. self.body_iterator = content
  201. else:
  202. self.body_iterator = iterate_in_threadpool(content)
  203. self.status_code = status_code
  204. self.media_type = self.media_type if media_type is None else media_type
  205. self.background = background
  206. self.init_headers(headers)
  207. async def listen_for_disconnect(self, receive: Receive) -> None:
  208. while True:
  209. message = await receive()
  210. if message["type"] == "http.disconnect":
  211. break
  212. async def stream_response(self, send: Send) -> None:
  213. await send({"type": "http.response.start", "status": self.status_code, "headers": self.raw_headers})
  214. async for chunk in self.body_iterator:
  215. if not isinstance(chunk, bytes | memoryview):
  216. chunk = chunk.encode(self.charset)
  217. await send({"type": "http.response.body", "body": chunk, "more_body": True})
  218. await send({"type": "http.response.body", "body": b"", "more_body": False})
  219. async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
  220. if scope["type"] == "websocket":
  221. send = self._wrap_websocket_denial_send(send)
  222. await self.stream_response(send)
  223. if self.background is not None:
  224. await self.background()
  225. return
  226. spec_version = tuple(map(int, scope.get("asgi", {}).get("spec_version", "2.0").split(".")))
  227. if spec_version >= (2, 4):
  228. try:
  229. await self.stream_response(send)
  230. except OSError:
  231. raise ClientDisconnect()
  232. else:
  233. with collapse_excgroups():
  234. async with anyio.create_task_group() as task_group:
  235. async def wrap(func: Callable[[], Awaitable[None]]) -> None:
  236. await func()
  237. task_group.cancel_scope.cancel()
  238. task_group.start_soon(wrap, partial(self.stream_response, send))
  239. await wrap(partial(self.listen_for_disconnect, receive))
  240. if self.background is not None:
  241. await self.background()
  242. class MalformedRangeHeader(Exception):
  243. def __init__(self, content: str = "Malformed range header.") -> None:
  244. self.content = content
  245. class RangeNotSatisfiable(Exception):
  246. def __init__(self, max_size: int) -> None:
  247. self.max_size = max_size
  248. class FileResponse(Response):
  249. chunk_size = 64 * 1024
  250. def __init__(
  251. self,
  252. path: str | os.PathLike[str],
  253. status_code: int = 200,
  254. headers: Mapping[str, str] | None = None,
  255. media_type: str | None = None,
  256. background: BackgroundTask | None = None,
  257. filename: str | None = None,
  258. stat_result: os.stat_result | None = None,
  259. content_disposition_type: str = "attachment",
  260. ) -> None:
  261. self.path = path
  262. self.status_code = status_code
  263. self.filename = filename
  264. if media_type is None:
  265. media_type = guess_type(filename or path)[0] or "text/plain"
  266. self.media_type = media_type
  267. self.background = background
  268. self.init_headers(headers)
  269. self.headers.setdefault("accept-ranges", "bytes")
  270. if self.filename is not None:
  271. content_disposition_filename = quote(self.filename)
  272. if content_disposition_filename != self.filename:
  273. content_disposition = f"{content_disposition_type}; filename*=utf-8''{content_disposition_filename}"
  274. else:
  275. content_disposition = f'{content_disposition_type}; filename="{self.filename}"'
  276. self.headers.setdefault("content-disposition", content_disposition)
  277. self.stat_result = stat_result
  278. if stat_result is not None:
  279. self.set_stat_headers(stat_result)
  280. def set_stat_headers(self, stat_result: os.stat_result) -> None:
  281. content_length = str(stat_result.st_size)
  282. last_modified = formatdate(stat_result.st_mtime, usegmt=True)
  283. etag_base = str(stat_result.st_mtime) + "-" + str(stat_result.st_size)
  284. etag = f'"{hashlib.md5(etag_base.encode(), usedforsecurity=False).hexdigest()}"'
  285. self.headers.setdefault("content-length", content_length)
  286. self.headers.setdefault("last-modified", last_modified)
  287. self.headers.setdefault("etag", etag)
  288. async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
  289. scope_type = scope["type"]
  290. send_header_only = scope_type == "http" and scope["method"].upper() == "HEAD"
  291. send_pathsend = scope_type == "http" and "http.response.pathsend" in scope.get("extensions", {})
  292. if scope_type == "websocket":
  293. send = self._wrap_websocket_denial_send(send)
  294. if self.stat_result is None:
  295. try:
  296. stat_result = await anyio.to_thread.run_sync(os.stat, self.path)
  297. self.set_stat_headers(stat_result)
  298. except FileNotFoundError:
  299. raise RuntimeError(f"File at path {self.path} does not exist.")
  300. else:
  301. mode = stat_result.st_mode
  302. if not stat.S_ISREG(mode):
  303. raise RuntimeError(f"File at path {self.path} is not a file.")
  304. else:
  305. stat_result = self.stat_result
  306. headers = Headers(scope=scope)
  307. http_range = headers.get("range")
  308. http_if_range = headers.get("if-range")
  309. if http_range is None or (http_if_range is not None and not self._should_use_range(http_if_range)):
  310. await self._handle_simple(send, send_header_only, send_pathsend)
  311. else:
  312. try:
  313. ranges = self._parse_range_header(http_range, stat_result.st_size)
  314. except MalformedRangeHeader as exc:
  315. return await PlainTextResponse(exc.content, status_code=400)(scope, receive, send)
  316. except RangeNotSatisfiable as exc:
  317. response = PlainTextResponse(status_code=416, headers={"Content-Range": f"bytes */{exc.max_size}"})
  318. return await response(scope, receive, send)
  319. if len(ranges) == 1:
  320. start, end = ranges[0]
  321. await self._handle_single_range(send, start, end, stat_result.st_size, send_header_only)
  322. else:
  323. await self._handle_multiple_ranges(send, ranges, stat_result.st_size, send_header_only)
  324. if self.background is not None:
  325. await self.background()
  326. async def _handle_simple(self, send: Send, send_header_only: bool, send_pathsend: bool) -> None:
  327. await send({"type": "http.response.start", "status": self.status_code, "headers": self.raw_headers})
  328. if send_header_only:
  329. await send({"type": "http.response.body", "body": b"", "more_body": False})
  330. elif send_pathsend:
  331. await send({"type": "http.response.pathsend", "path": str(self.path)})
  332. else:
  333. async with await anyio.open_file(self.path, mode="rb") as file:
  334. more_body = True
  335. while more_body:
  336. chunk = await file.read(self.chunk_size)
  337. more_body = len(chunk) == self.chunk_size
  338. await send({"type": "http.response.body", "body": chunk, "more_body": more_body})
  339. async def _handle_single_range(
  340. self, send: Send, start: int, end: int, file_size: int, send_header_only: bool
  341. ) -> None:
  342. headers = MutableHeaders(raw=list(self.raw_headers))
  343. headers["content-range"] = f"bytes {start}-{end - 1}/{file_size}"
  344. headers["content-length"] = str(end - start)
  345. await send({"type": "http.response.start", "status": 206, "headers": headers.raw})
  346. if send_header_only:
  347. await send({"type": "http.response.body", "body": b"", "more_body": False})
  348. else:
  349. async with await anyio.open_file(self.path, mode="rb") as file:
  350. await file.seek(start)
  351. more_body = True
  352. while more_body:
  353. chunk = await file.read(min(self.chunk_size, end - start))
  354. start += len(chunk)
  355. more_body = len(chunk) == self.chunk_size and start < end
  356. await send({"type": "http.response.body", "body": chunk, "more_body": more_body})
  357. async def _handle_multiple_ranges(
  358. self,
  359. send: Send,
  360. ranges: list[tuple[int, int]],
  361. file_size: int,
  362. send_header_only: bool,
  363. ) -> None:
  364. # In firefox and chrome, they use boundary with 95-96 bits entropy (that's roughly 13 bytes).
  365. boundary = token_hex(13)
  366. content_length, header_generator = self.generate_multipart(
  367. ranges, boundary, file_size, self.headers["content-type"]
  368. )
  369. headers = MutableHeaders(raw=list(self.raw_headers))
  370. headers["content-type"] = f"multipart/byteranges; boundary={boundary}"
  371. headers["content-length"] = str(content_length)
  372. await send({"type": "http.response.start", "status": 206, "headers": headers.raw})
  373. if send_header_only:
  374. await send({"type": "http.response.body", "body": b"", "more_body": False})
  375. else:
  376. async with await anyio.open_file(self.path, mode="rb") as file:
  377. for start, end in ranges:
  378. await send({"type": "http.response.body", "body": header_generator(start, end), "more_body": True})
  379. await file.seek(start)
  380. while start < end:
  381. chunk = await file.read(min(self.chunk_size, end - start))
  382. start += len(chunk)
  383. await send({"type": "http.response.body", "body": chunk, "more_body": True})
  384. await send({"type": "http.response.body", "body": b"\r\n", "more_body": True})
  385. await send(
  386. {
  387. "type": "http.response.body",
  388. "body": f"--{boundary}--".encode("latin-1"),
  389. "more_body": False,
  390. }
  391. )
  392. def _should_use_range(self, http_if_range: str) -> bool:
  393. return http_if_range == self.headers["last-modified"] or http_if_range == self.headers["etag"]
  394. @classmethod
  395. def _parse_range_header(cls, http_range: str, file_size: int) -> list[tuple[int, int]]:
  396. ranges: list[tuple[int, int]] = []
  397. try:
  398. units, range_ = http_range.split("=", 1)
  399. except ValueError:
  400. raise MalformedRangeHeader()
  401. units = units.strip().lower()
  402. if units != "bytes":
  403. raise MalformedRangeHeader("Only support bytes range")
  404. ranges = cls._parse_ranges(range_, file_size)
  405. if len(ranges) == 0:
  406. raise MalformedRangeHeader("Range header: range must be requested")
  407. if any(not (0 <= start < file_size) for start, _ in ranges):
  408. raise RangeNotSatisfiable(file_size)
  409. if any(start > end for start, end in ranges):
  410. raise MalformedRangeHeader("Range header: start must be less than end")
  411. if len(ranges) == 1:
  412. return ranges
  413. # Merge overlapping ranges
  414. ranges.sort()
  415. result: list[tuple[int, int]] = [ranges[0]]
  416. for start, end in ranges[1:]:
  417. last_start, last_end = result[-1]
  418. if start <= last_end:
  419. result[-1] = (last_start, max(last_end, end))
  420. else:
  421. result.append((start, end))
  422. return result
  423. @classmethod
  424. def _parse_ranges(cls, range_: str, file_size: int) -> list[tuple[int, int]]:
  425. ranges: list[tuple[int, int]] = []
  426. for part in range_.split(","):
  427. part = part.strip()
  428. # If the range is empty or a single dash, we ignore it.
  429. if not part or part == "-":
  430. continue
  431. # If the range is not in the format "start-end", we ignore it.
  432. if "-" not in part:
  433. continue
  434. start_str, end_str = part.split("-", 1)
  435. start_str = start_str.strip()
  436. end_str = end_str.strip()
  437. try:
  438. start = int(start_str) if start_str else file_size - int(end_str)
  439. end = int(end_str) + 1 if start_str and end_str and int(end_str) < file_size else file_size
  440. ranges.append((start, end))
  441. except ValueError:
  442. # If the range is not numeric, we ignore it.
  443. continue
  444. return ranges
  445. def generate_multipart(
  446. self,
  447. ranges: Sequence[tuple[int, int]],
  448. boundary: str,
  449. max_size: int,
  450. content_type: str,
  451. ) -> tuple[int, Callable[[int, int], bytes]]:
  452. r"""
  453. Multipart response headers generator.
  454. ```
  455. --{boundary}\r\n
  456. Content-Type: {content_type}\r\n
  457. Content-Range: bytes {start}-{end-1}/{max_size}\r\n
  458. \r\n
  459. ..........content...........\r\n
  460. --{boundary}\r\n
  461. Content-Type: {content_type}\r\n
  462. Content-Range: bytes {start}-{end-1}/{max_size}\r\n
  463. \r\n
  464. ..........content...........\r\n
  465. --{boundary}--
  466. ```
  467. """
  468. boundary_len = len(boundary)
  469. static_header_part_len = 49 + boundary_len + len(content_type) + len(str(max_size))
  470. content_length = sum(
  471. (len(str(start)) + len(str(end - 1)) + static_header_part_len) # Headers
  472. + (end - start) # Content
  473. for start, end in ranges
  474. ) + (
  475. 4 + boundary_len # --boundary--
  476. )
  477. return (
  478. content_length,
  479. lambda start, end: (
  480. f"""\
  481. --{boundary}\r
  482. Content-Type: {content_type}\r
  483. Content-Range: bytes {start}-{end - 1}/{max_size}\r
  484. \r
  485. """
  486. ).encode("latin-1"),
  487. )