diff --git a/starlette/responses.py b/starlette/responses.py index d5df75d28..9751f360d 100644 --- a/starlette/responses.py +++ b/starlette/responses.py @@ -6,7 +6,8 @@ import os import stat import sys -from collections.abc import AsyncIterable, Awaitable, Callable, Iterable, Mapping, Sequence +from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterable, Mapping, Sequence +from contextlib import asynccontextmanager from datetime import datetime from email.utils import format_datetime, formatdate from functools import partial @@ -365,7 +366,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: http_if_range = headers.get("if-range") if http_range is None or (http_if_range is not None and not self._should_use_range(http_if_range)): - await self._handle_simple(send, send_header_only, send_pathsend) + send_file = partial(self._handle_simple, send, send_header_only, send_pathsend) else: try: ranges = self._parse_range_header(http_range, stat_result.st_size) @@ -376,16 +377,44 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: return await response(scope, receive, send) if len(ranges) == 0: - await self._handle_simple(send, send_header_only, send_pathsend) + send_file = partial(self._handle_simple, send, send_header_only, send_pathsend) elif len(ranges) == 1: start, end = ranges[0] - await self._handle_single_range(send, start, end, stat_result.st_size, send_header_only) + send_file = partial(self._handle_single_range, send, start, end, stat_result.st_size, send_header_only) + send_pathsend = False else: - await self._handle_multiple_ranges(send, ranges, stat_result.st_size, send_header_only) + send_file = partial(self._handle_multiple_ranges, send, ranges, stat_result.st_size, send_header_only) + send_pathsend = False + + spec_version = tuple(map(int, scope.get("asgi", {}).get("spec_version", "2.0").split("."))) + if scope_type != "http" or send_header_only or send_pathsend or spec_version >= (2, 4): + await send_file() + else: + async with create_collapsing_task_group() as task_group: + + async def stream_file() -> None: + await send_file() + task_group.cancel_scope.cancel() + + task_group.start_soon(stream_file) + while True: + if (await receive())["type"] == "http.disconnect": + task_group.cancel_scope.cancel() + break if self.background is not None: await self.background() + @asynccontextmanager + async def _open_file(self) -> AsyncIterator[anyio.AsyncFile[bytes]]: + file = await anyio.open_file(self.path, mode="rb") + try: + yield file + finally: + # Closing must finish even when the transfer is cancelled. + with anyio.CancelScope(shield=True): + await file.aclose() + async def _handle_simple(self, send: Send, send_header_only: bool, send_pathsend: bool) -> None: await send({"type": "http.response.start", "status": self.status_code, "headers": self.raw_headers}) if send_header_only: @@ -393,7 +422,7 @@ async def _handle_simple(self, send: Send, send_header_only: bool, send_pathsend elif send_pathsend: await send({"type": "http.response.pathsend", "path": str(self.path)}) else: - async with await anyio.open_file(self.path, mode="rb") as file: + async with self._open_file() as file: more_body = True while more_body: chunk = await file.read(self.chunk_size) @@ -410,7 +439,7 @@ async def _handle_single_range( if send_header_only: await send({"type": "http.response.body", "body": b"", "more_body": False}) else: - async with await anyio.open_file(self.path, mode="rb") as file: + async with self._open_file() as file: await file.seek(start) more_body = True while more_body: @@ -438,7 +467,7 @@ async def _handle_multiple_ranges( if send_header_only: await send({"type": "http.response.body", "body": b"", "more_body": False}) else: - async with await anyio.open_file(self.path, mode="rb") as file: + async with self._open_file() as file: for start, end in ranges: await send({"type": "http.response.body", "body": header_generator(start, end), "more_body": True}) await file.seek(start) diff --git a/tests/test_responses.py b/tests/test_responses.py index 2733c9f22..f62c2f37a 100644 --- a/tests/test_responses.py +++ b/tests/test_responses.py @@ -1,13 +1,14 @@ from __future__ import annotations import datetime as dt +import os import sys import time from collections.abc import AsyncGenerator, AsyncIterator, Iterator from dataclasses import dataclass from http.cookies import SimpleCookie from pathlib import Path -from typing import Any +from typing import Any, Literal import anyio import pytest @@ -395,6 +396,130 @@ async def send(message: Message) -> None: ) +@pytest.fixture +def file_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> tuple[Path, list[anyio.AsyncFile[bytes]]]: + path = tmp_path / "file.bin" + path.write_bytes(b"x" * (8 * FileResponse.chunk_size)) + files: list[anyio.AsyncFile[bytes]] = [] + open_file = anyio.open_file + + async def track_open_file(path: str | os.PathLike[str], mode: Literal["rb"]) -> anyio.AsyncFile[bytes]: + file = await open_file(path, mode=mode) + files.append(file) + return file + + monkeypatch.setattr(anyio, "open_file", track_open_file) + return path, files + + +@pytest.fixture(params=[None, b"bytes=0-", b"bytes=0-131071,196608-524287"]) +def scope(request: pytest.FixtureRequest) -> Scope: + return { + "type": "http", + "method": "GET", + "headers": [] if request.param is None else [(b"range", request.param)], + "extensions": {"http.response.pathsend": {}} if request.param is not None else {}, + } + + +@pytest.mark.anyio +@pytest.mark.parametrize("spec_version", [None, "2.0", "2.3"]) +async def test_file_response_stops_on_disconnect( + file_path: tuple[Path, list[anyio.AsyncFile[bytes]]], scope: Scope, spec_version: str | None +) -> None: + path, files = file_path + if spec_version is not None: + scope["asgi"] = {"spec_version": spec_version} + disconnected = anyio.Event() + received_request = False + submitted = 0 + background_ran = False + + async def receive() -> Message: + nonlocal received_request + if not received_request: + received_request = True + return {"type": "http.request", "body": b"", "more_body": False} + await disconnected.wait() + return {"type": "http.disconnect"} + + async def send(message: Message) -> None: + nonlocal submitted + if message["type"] == "http.response.body": + submitted += len(message["body"]) + if submitted >= FileResponse.chunk_size: + disconnected.set() + + async def cleanup() -> None: + nonlocal background_ran + assert len(files) == 1 + assert files[0].closed + await anyio.sleep(0) + background_ran = True + + with anyio.fail_after(5): + await FileResponse(path, background=BackgroundTask(cleanup))(scope, receive, send) + + assert FileResponse.chunk_size <= submitted < 3 * FileResponse.chunk_size + assert background_ran + + +@pytest.mark.anyio +@pytest.mark.parametrize("spec_version", ["2.0", "2.3", "2.4"]) +async def test_file_response_closes_on_cancellation( + file_path: tuple[Path, list[anyio.AsyncFile[bytes]]], scope: Scope, spec_version: str +) -> None: + path, files = file_path + scope["asgi"] = {"spec_version": spec_version} + + async def receive() -> Message: + await anyio.sleep_forever() + pytest.fail("The disconnect listener should be cancelled") # pragma: no cover - sleep never returns + + async def send(message: Message) -> None: + if message["type"] == "http.response.body": + cancel_scope.cancel() + await anyio.sleep_forever() + + async def cleanup() -> None: + pytest.fail( + "Background tasks should not run after external cancellation" + ) # pragma: no cover - failure sentinel + + with anyio.fail_after(5), anyio.CancelScope() as cancel_scope: + await FileResponse(path, background=BackgroundTask(cleanup))(scope, receive, send) + + assert cancel_scope.cancelled_caught + assert len(files) == 1 + assert files[0].closed + + +@pytest.mark.anyio +@pytest.mark.parametrize("spec_version", ["2.0", "2.3", "2.4"]) +async def test_file_response_closes_on_send_error( + file_path: tuple[Path, list[anyio.AsyncFile[bytes]]], scope: Scope, spec_version: str +) -> None: + path, files = file_path + scope["asgi"] = {"spec_version": spec_version} + error = OSError("Disconnected") + + async def receive() -> Message: + assert spec_version != "2.4" + await anyio.sleep_forever() + pytest.fail("The disconnect listener should be cancelled") # pragma: no cover - sleep never returns + + async def send(message: Message) -> None: + if message["type"] == "http.response.body": + raise error + + with anyio.fail_after(5), pytest.raises(OSError) as exc: + await FileResponse(path)(scope, receive, send) + + assert exc.value is error + assert len(files) == 1 + assert files[0].closed + + def test_set_cookie(test_client_factory: TestClientFactory, monkeypatch: pytest.MonkeyPatch) -> None: # Mock time used as a reference for `Expires` by stdlib `SimpleCookie`. mocked_now = dt.datetime(2037, 1, 22, 12, 0, 0, tzinfo=dt.timezone.utc) @@ -1018,7 +1143,8 @@ class SmallChunkSizeFileResponse(FileResponse): start_message: dict[str, Any] = {} async def receive() -> Message: - raise NotImplementedError("Should not be called!") + await anyio.sleep_forever() + pytest.fail("The disconnect listener should be cancelled") # pragma: no cover - sleep never returns async def send(message: Message) -> None: if message["type"] == "http.response.start":