diff --git a/tests/downloaders/test_failure_surfacing.py b/tests/downloaders/test_failure_surfacing.py index c43cfc54..ea8f5861 100644 --- a/tests/downloaders/test_failure_surfacing.py +++ b/tests/downloaders/test_failure_surfacing.py @@ -6,17 +6,18 @@ import pytest import requests as rq from unshackle.core.constants import DOWNLOAD_CANCELLED -from unshackle.core.downloaders.requests import download, has_range_header +from unshackle.core.downloaders.requests import download, has_range_header, parse_content_range, request_range_start from unshackle.core.downloaders.requests import requests as requests_downloader # the package re-exports the `requests` function, shadowing the module of the same name downloader = sys.modules[download.__module__] BODY = b"".join(f"{i:05d} payload\n".encode() for i in range(500)) +CUT = len(BODY) // 2 class _Server: - """Serves BODY, with a switch for the two failure modes retries have to handle.""" + """Serves BODY, with a switch for the failure modes retries have to handle.""" def __init__(self, mode): self.mode = mode @@ -26,8 +27,16 @@ class _Server: self.sock.listen(8) self.url = f"http://127.0.0.1:{self.sock.getsockname()[1]}/seg.mp4" self.requests = [] + self.plain_hits = 0 threading.Thread(target=self.serve, daemon=True).start() + def send_206(self, conn, body, content_range=None): + head = "HTTP/1.1 206 Partial Content\r\n" + if content_range: + head += f"Content-Range: {content_range}\r\n" + head += f"Content-Length: {len(body)}\r\nConnection: close\r\n\r\n" + conn.sendall(head.encode() + body) + def serve(self): while True: try: @@ -45,13 +54,31 @@ class _Server: ranged = "range" in headers if self.mode == "unavailable": conn.sendall(b"HTTP/1.1 503 Service Unavailable\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") - elif ranged: - conn.sendall( - b"HTTP/1.1 416 Range Not Satisfiable\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" - ) + elif not ranged: + self.plain_hits += 1 + if self.mode.startswith("resume_") and self.plain_hits == 1: + # promise the full body but hang up early: leaves a CUT-byte + # partial behind so the next attempt sends a resume Range + conn.sendall( + f"HTTP/1.1 200 OK\r\nContent-Length: {len(BODY)}\r\nConnection: close\r\n\r\n".encode() + + BODY[:CUT] + ) + else: + conn.sendall( + f"HTTP/1.1 200 OK\r\nContent-Length: {len(BODY)}\r\nConnection: close\r\n\r\n".encode() + + BODY + ) + elif self.mode == "resume_whole_body_206": + # ignores the requested start: a 206 carrying the entire resource + self.send_206(conn, BODY, f"bytes 0-{len(BODY) - 1}/{len(BODY)}") + elif self.mode == "resume_alien_206": + # a 206 starting neither at the requested offset nor at zero + self.send_206(conn, BODY[5:], f"bytes 5-{len(BODY) - 1}/{len(BODY)}") + elif self.mode == "resume_missing_content_range": + self.send_206(conn, BODY) else: conn.sendall( - f"HTTP/1.1 200 OK\r\nContent-Length: {len(BODY)}\r\nConnection: close\r\n\r\n".encode() + BODY + b"HTTP/1.1 416 Range Not Satisfiable\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" ) except Exception: pass @@ -111,6 +138,65 @@ def test_stale_partial_restarts_clean_on_416(tmp_path, server): assert [r.get("range") for r in server.requests] == [f"bytes={len(BODY)}-", None] +@pytest.mark.parametrize("server", ["resume_whole_body_206"], indirect=True) +def test_resume_206_carrying_the_whole_body_is_rewritten_not_appended(tmp_path, server): + save_path = tmp_path / "0.mp4" + events = list(download(url=server.url, save_path=save_path, session=rq.Session())) + assert save_path.read_bytes() == BODY + assert events[-1]["written"] == len(BODY) + assert [r.get("range") for r in server.requests] == [None, f"bytes={CUT}-"] + + +@pytest.mark.parametrize("server", ["resume_alien_206"], indirect=True) +def test_resume_206_with_wrong_start_restarts_clean(tmp_path, server): + save_path = tmp_path / "0.mp4" + list(download(url=server.url, save_path=save_path, session=rq.Session())) + assert save_path.read_bytes() == BODY + assert [r.get("range") for r in server.requests] == [None, f"bytes={CUT}-", None] + + +@pytest.mark.parametrize("server", ["resume_missing_content_range"], indirect=True) +def test_resume_206_without_content_range_restarts_clean(tmp_path, server): + save_path = tmp_path / "0.mp4" + list(download(url=server.url, save_path=save_path, session=rq.Session())) + assert save_path.read_bytes() == BODY + assert [r.get("range") for r in server.requests] == [None, f"bytes={CUT}-", None] + + +@pytest.mark.parametrize("server", ["resume_whole_body_206"], indirect=True) +def test_part_with_wrong_content_range_start_never_writes(tmp_path, server): + save_path = tmp_path / "0.mp4" + save_path.write_bytes(b"\x00" * len(BODY)) + with pytest.raises(OSError): + list( + download( + url=server.url, + save_path=save_path, + session=rq.Session(), + part_offset=CUT, + part_end=len(BODY) - 1, + ) + ) + assert save_path.read_bytes() == b"\x00" * len(BODY) + assert len(server.requests) == downloader.MAX_ATTEMPTS + + +@pytest.mark.parametrize("server", ["resume_whole_body_206"], indirect=True) +def test_slice_segment_rejects_a_206_for_the_wrong_range(tmp_path, server): + save_path = tmp_path / "0.mp4" + with pytest.raises(OSError): + list( + download( + url=server.url, + save_path=save_path, + session=rq.Session(), + headers={"Range": f"bytes={CUT}-{len(BODY) - 1}"}, + ) + ) + assert not save_path.exists() + assert len(server.requests) == downloader.MAX_ATTEMPTS + + @pytest.mark.parametrize("server", ["range_unsatisfiable"], indirect=True) def test_slice_segment_never_gets_a_resume_range(tmp_path, server): save_path = tmp_path / "0.mp4" @@ -185,3 +271,33 @@ def test_batch_larger_than_the_queue_window_downloads_every_segment(tmp_path, se ) def test_has_range_header(item, expected): assert has_range_header(item) is expected + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + (None, None), + ("", None), + ("garbage", None), + ("bytes */400000", None), + ("bytes 0-399999/400000", (0, 399999, 400000)), + ("BYTES 150000-399999/*", (150000, 399999, None)), + (b"bytes 5-9/10", (5, 9, 10)), + ], +) +def test_parse_content_range(value, expected): + assert parse_content_range(value) == expected + + +@pytest.mark.parametrize( + ("headers", "expected"), + [ + ({}, None), + ({"Accept": "*/*"}, None), + ({"Range": "bytes=-500"}, None), + ({"Range": "bytes=100-"}, 100), + ({"range": "bytes=100-227"}, 100), + ], +) +def test_request_range_start(headers, expected): + assert request_range_start(headers) == expected diff --git a/unshackle/core/downloaders/requests.py b/unshackle/core/downloaders/requests.py index 54d3d7f1..50ecd700 100644 --- a/unshackle/core/downloaders/requests.py +++ b/unshackle/core/downloaders/requests.py @@ -422,6 +422,27 @@ def range_covers_full_body(headers: MutableMapping[str, Any], content_length: in return start == 0 and content_length == end + 1 +def parse_content_range(value: Any) -> Optional[tuple[int, int, Optional[int]]]: + """Parse a 206's ``Content-Range`` into ``(start, end, total)``; None when absent or unparseable. + + ``total`` is None for an unknown complete length (``*``). None fails closed: a caller + must not place bytes at an offset it cannot prove the server agreed to. + """ + if isinstance(value, bytes): + value = value.decode("latin1", "replace") + match = re.fullmatch(r"bytes\s+(\d+)-(\d+)/(\d+|\*)", str(value or "").strip(), re.IGNORECASE) + if not match: + return None + return int(match.group(1)), int(match.group(2)), None if match.group(3) == "*" else int(match.group(3)) + + +def request_range_start(headers: MutableMapping[str, Any]) -> Optional[int]: + """Start byte of the request's ``Range`` header; None when absent or not ``bytes=start-...``.""" + value = str(next((v for k, v in headers.items() if str(k).lower() == "range"), "")) + match = re.match(r"bytes=(\d+)-", value.strip()) + return int(match.group(1)) if match else None + + def force_unblock_stream(stream: Any) -> None: """Best-effort: shut the stream's socket so a thread parked in a blocking read returns now. @@ -723,10 +744,42 @@ def download( # item_range 206s reflect the slice's own Range, not a resume; those retries rewrite in "wb" resumed = (not part_mode) and (not item_range) and resume_offset > 0 and stream.status_code == 206 + if resumed: + content_range = parse_content_range( + stream.headers.get("Content-Range") or stream.headers.get("content-range") + ) + if content_range is None or content_range[0] != resume_offset: + # some servers answer `bytes=X-` with a 206 for the whole resource; + # appending that body would misplace every byte after the partial. A + # proven full-resource-from-zero body is byte-identical to a 200, so + # keep the response and rewrite from scratch; anything else (alien + # start, prefix slice, no provable placement) restarts clean like a 416 + start, end, total = content_range or (None, None, None) + if start == 0 and total is not None and end == total - 1: + resumed = False + resume_offset = 0 + else: + try: + stream.close() + except Exception: + pass + tmp_file.unlink(missing_ok=True) + resume_offset = 0 + written = 0 + continue if (not part_mode) and resume_offset > 0 and not resumed: resume_offset = 0 - if part_mode and stream.status_code != 206: - raise IOError(f"expected 206 for ranged part, got {stream.status_code}") + if part_mode and part_offset is not None and part_end is not None: # spelled so mypy narrows + if stream.status_code != 206: + raise IOError(f"expected 206 for ranged part, got {stream.status_code}") + content_range = parse_content_range( + stream.headers.get("Content-Range") or stream.headers.get("content-range") + ) + if content_range is None or content_range[0] != part_offset + written or content_range[1] > part_end: + # a wrong start would land bytes at the wrong offset in the shared file, + # and an end past part_end would spill into a sibling part's region; fail + # the attempt before any byte is written + raise IOError(f"part {part_offset}-{part_end} got Content-Range {content_range!r}") if use_rnet: content_encoded = is_content_encoded( stream.headers.get("Content-Encoding") or stream.headers.get("content-encoding") @@ -749,6 +802,19 @@ def download( if not (stream.status_code == 200 and range_covers_full_body(req_headers, content_length)): raise IOError(f"expected 206 for byte-range segment, got {stream.status_code}") + if item_range and not part_mode and stream.status_code == 206: + # a 206 for a different range than the slice asked for (e.g. the whole + # parent) would be written as the segment and silently corrupt the merge + slice_start = request_range_start(req_headers) + if slice_start is not None: + content_range = parse_content_range( + stream.headers.get("Content-Range") or stream.headers.get("content-range") + ) + if content_range is None or content_range[0] != slice_start: + raise IOError( + f"byte-range segment got Content-Range {content_range!r}, expected start {slice_start}" + ) + if resumed and content_encoded: try: stream.close()