fix(downloader): validate partial download responses and sizes

src/downloader.py

@@ -1,5 +1,6 @@
 import json
 import os
+import re
 import shutil
 import sys
 import time
@@ -187,17 +188,27 @@ def _get_download_headers(resume_byte_pos, url):
 
 
 def _calculate_total_size(response, headers, resume_byte_pos):
-    if headers and response.status_code == 206:
-        content_range = response.headers.get("content-range", "")
-        if content_range:
-            return int(content_range.split("/")[-1])
-        return resume_byte_pos + int(response.headers.get("content-length", 0))
+    if response.status_code == 206:
+        content_range = re.fullmatch(
+            r"bytes (\d+)-(\d+)/(\d+)", response.headers.get("content-range", ""), flags=re.IGNORECASE
+        )
+        if not headers or content_range is None:
+            raise ValueError("Unexpected or invalid partial download response")
+        start, end, total = map(int, content_range.groups())
+        if start != resume_byte_pos or not start <= end < total or end != total - 1:
+            raise ValueError("Content-Range does not match the requested download range")
+        content_length = response.headers.get("content-length")
+        if content_length is not None and int(content_length) != end - start + 1:
+            raise ValueError("Content-Length does not match Content-Range")
+        return total
+    if int(response.headers.get("content-length", 0)) < 0:
+        raise ValueError("Invalid Content-Length")
     return int(response.headers.get("content-length", 0))
 
 
 def _handle_no_range_support(resume_byte_pos, destination):
     if resume_byte_pos > 0:
-        logging.warning("Server doesn't support range requests, restarting download")
+        logging.warning("Server returned a full response, restarting download from byte zero")
         if os.path.exists(destination):
             os.remove(destination)
         return 0, "wb"
@@ -238,6 +249,9 @@ def _download_content(response, destination, mode, url, total_size, resume_byte_
                 if not chunk:
                     continue
 
+                if (total_size > 0 or "content-length" in response.headers) and downloaded + len(chunk) > total_size:
+                    raise ValueError(f"Download exceeds expected size: {downloaded + len(chunk)}/{total_size} bytes")
+
                 size = f.write(chunk)
                 downloaded += size
 
@@ -295,14 +309,19 @@ def _perform_download(url, destination, resume_byte_pos, mode, start_time):
     headers = _get_download_headers(resume_byte_pos, url)
 
     with requests.get(url, stream=True, headers=headers, timeout=(30, 60)) as response:
+        if response.status_code == 416 and resume_byte_pos > 0:
+            cleanup_download_state(destination)
+            raise RequestException("Resume range rejected; retrying download from byte zero")
         response.raise_for_status()
+        if response.status_code not in (200, 206):
+            raise ValueError(f"Unexpected download response status: {response.status_code}")
 
         total_size = _calculate_total_size(response, headers, resume_byte_pos)
 
         if total_size > 0:
             logging.info(f"Starting download of {total_size / (1024**3):.2f}GB to {os.path.basename(destination)}")
 
-        if not headers and response.status_code != 206:
+        if response.status_code == 200:
             new_pos, new_mode = _handle_no_range_support(resume_byte_pos, destination)
             if new_mode:
                 resume_byte_pos = new_pos
@@ -318,7 +337,7 @@ def _perform_download(url, destination, resume_byte_pos, mode, start_time):
 
             save_download_state(destination, url, downloaded, total_size)
 
-            if total_size > 0 and downloaded < total_size:
+            if (total_size > 0 or "content-length" in response.headers) and downloaded != total_size:
                 raise Exception(f"Download incomplete: {downloaded}/{total_size} bytes")
 
             cleanup_download_state(destination)

tests/test_downloader.py

@@ -241,15 +241,10 @@ def test_get_download_headers_empty_when_no_range_support(monkeypatch: pytest.Mo
 
 
 def test_calculate_total_size_with_range_response_using_content_range():
-    resp = _mock_response(status_code=206, headers={"content-range": "bytes 0-99/12345"})
+    resp = _mock_response(status_code=206, headers={"content-range": "bytes 0-12344/12345"})
     assert downloader._calculate_total_size(resp, {"Range": "bytes=0-"}, 0) == 12345
 
 
-def test_calculate_total_size_with_range_response_no_content_range():
-    resp = _mock_response(status_code=206, headers={"content-length": "100"})
-    assert downloader._calculate_total_size(resp, {"Range": "bytes=50-"}, 50) == 150
-
-
 def test_calculate_total_size_without_range_uses_content_length():
     resp = _mock_response(status_code=200, headers={"content-length": "9999"})
     assert downloader._calculate_total_size(resp, {}, 0) == 9999
@@ -375,3 +370,82 @@ def test_download_file_propagates_disk_write_oserror(tmp_path: Path, monkeypatch
         pytest.raises(OSError, match="disk full"),
     ):
         downloader.download_file("https://example.com/x", str(dest))
+
+
+@pytest.mark.parametrize(("status", "range_unit"), [(200, "bytes"), (206, "bytes"), (206, "Bytes")])
+def test_resume_uses_response_status(tmp_path: Path, status: int, range_unit: str):
+    dest = tmp_path / "archive"
+    dest.write_bytes(b"abcd")
+    url = "https://example.com/archive"
+    downloader.save_download_state(str(dest), url, 4, 10)
+    headers = {"content-length": "10" if status == 200 else "6"}
+    if status == 206:
+        headers["content-range"] = f"{range_unit} 4-9/10"
+    response = _mock_response(status, headers, [b"abcdefghij" if status == 200 else b"efghij"])
+    with (
+        patch.object(downloader, "supports_range_requests", return_value=True),
+        patch.object(downloader.requests, "get", return_value=response) as get,
+    ):
+        assert downloader.download_file(url, str(dest)) is True
+    assert get.call_count == 1
+    assert get.call_args.kwargs["headers"] == {"Range": "bytes=4-"}
+    assert dest.read_bytes() == b"abcdefghij"
+    assert not Path(downloader.get_download_state_file(str(dest))).exists()
+
+
+@pytest.mark.parametrize(
+    "headers",
+    [
+        {},
+        {"content-range": "bytes 0-9/10"},
+        {"content-range": "bytes 4-10/10"},
+        {"content-range": "bytes 4-9/10", "content-length": "10"},
+    ],
+)
+def test_invalid_resume_range_preserves_partial_file(tmp_path: Path, headers: dict):
+    dest = tmp_path / "archive"
+    dest.write_bytes(b"abcd")
+    url = "https://example.com/archive"
+    downloader.save_download_state(str(dest), url, 4, 10)
+    response = _mock_response(206, headers, [b"efghij"])
+    with (
+        patch.object(downloader, "supports_range_requests", return_value=True),
+        patch.object(downloader.requests, "get", return_value=response),
+    ):
+        assert downloader.download_file(url, str(dest)) is False
+    assert dest.read_bytes() == b"abcd"
+    response.iter_content.assert_not_called()
+
+
+def test_oversized_response_rejected_before_writing_chunk(tmp_path: Path):
+    dest = tmp_path / "archive"
+    response = _mock_response(200, {"content-length": "3"}, [b"abcd"])
+    with patch.object(downloader.requests, "get", return_value=response):
+        assert downloader.download_file("https://example.com/archive", str(dest)) is False
+    assert dest.read_bytes() == b""
+
+
+@pytest.mark.parametrize("retries", [1, 2])
+def test_range_rejection_restarts_within_retry_limit(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, retries: int):
+    dest = tmp_path / "archive"
+    dest.write_bytes(b"oversized partial file")
+    url = "https://example.com/archive"
+    downloader.save_download_state(str(dest), url, dest.stat().st_size, 4)
+    monkeypatch.setattr(config, "DOWNLOAD_MAX_RETRIES", str(retries))
+    monkeypatch.setattr(downloader.time, "sleep", lambda _: None)
+    with (
+        patch.object(downloader, "supports_range_requests", return_value=True),
+        patch.object(
+            downloader.requests,
+            "get",
+            side_effect=[_mock_response(416), _mock_response(200, {"content-length": "4"}, [b"abcd"])],
+        ) as get,
+    ):
+        assert downloader.download_file(url, str(dest)) is (retries == 2)
+    assert get.call_count == retries
+    assert not Path(downloader.get_download_state_file(str(dest))).exists()
+    if retries == 2:
+        assert get.call_args.kwargs["headers"] == {}
+        assert dest.read_bytes() == b"abcd"
+    else:
+        assert dest.read_bytes() == b"oversized partial file"