fix(downloader): validate Content-Range before placing 206 bytes

A 206 was trusted from its status alone, so a server answering a different range than requested silently corrupted output.

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

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()