Merge pull request #359 from ynput/bugfix/download-resume-integrity

Fix corrupted or incomplete files after interrupted download

Author: iLLiCiTiT

ayon_api/server_api.py

@@ -1686,6 +1686,14 @@ def _download_file_to_stream(
                             progress.set_destination_url(url)
                             continue
                     response.raise_for_status()
+                    if offset > 0 and response.status_code != 206:
+                        # Server ignored 'Range' and sends whole file again,
+                        #   already downloaded content must be discarded
+                        stream.seek(0)
+                        stream.truncate()
+                        progress.reset_transferred()
+                        headers.pop("Range", None)
+
                     if progress.get_content_size() is None:
                         progress.set_content_size(
                             response.headers["Content-length"]
@@ -1694,11 +1702,21 @@ def _download_file_to_stream(
                     for chunk in response.iter_content(chunk_size=chunk_size):
                         stream.write(chunk)
                         progress.add_transferred_chunk(len(chunk))
+
+                content_size = progress.get_content_size()
+                transferred = progress.get_transferred_size()
+                if content_size is not None and transferred != content_size:
+                    # Connection was closed before all content was received
+                    raise requests.exceptions.ConnectionError(
+                        f"Downloaded {transferred} out of {content_size}"
+                        f" bytes from '{url}'."
+                    )
                 break
 
             except (
                 requests.exceptions.Timeout,
                 requests.exceptions.ConnectionError,
+                requests.exceptions.ChunkedEncodingError,
             ):
                 if attempt == retries - 1:
                     raise

ayon_api/utils.py

@@ -955,7 +955,7 @@ def get_content_size(self) -> int | None:
         """
         return self._content_size
 
-    def set_content_size(self, content_size: int) -> None:
+    def set_content_size(self, content_size: int | str) -> None:
         """Set content size in bytes.
 
         Args:
@@ -967,7 +967,7 @@ def set_content_size(self, content_size: int) -> None:
         """
         if self._content_size is not None:
             raise ValueError("Content size was set more then once")
-        self._content_size = content_size
+        self._content_size = int(content_size)
 
     def get_started(self) -> bool:
         """Transfer was started.

tests/test_download_resume.py

@@ -0,0 +1,60 @@
+"""Interrupted downloads. Does not require running AYON server."""
+import io
+
+import pytest
+
+from ayon_api.server_api import ServerAPI
+from ayon_api.utils import RequestTypes
+
+from .fake_transfer import FakeResponse
+
+CONTENT = b"0123456789"
+SIZE_HEADER = {"Content-Length": str(len(CONTENT))}
+
+
+@pytest.fixture
+def con(monkeypatch):
+    monkeypatch.setattr("time.sleep", lambda *args, **kwargs: None)
+    return ServerAPI("http://localhost:0", create_session=False, max_retries=3)
+
+
+def _download(con, get_func):
+    con._base_functions_mapping[RequestTypes.get] = get_func
+    stream = io.BytesIO()
+    progress = con.download_file_to_stream("api/file", stream)
+    return stream.getvalue(), progress
+
+
+def test_incomplete_download_is_continued(con):
+    def get_func(url, **kwargs):
+        if "Range" not in kwargs["headers"]:
+            # Connection closed after 4 bytes without an exception
+            return FakeResponse(200, CONTENT[:4], SIZE_HEADER)
+        assert kwargs["headers"]["Range"] == "bytes=4-"
+        return FakeResponse(206, CONTENT[4:], {"Content-Length": "6"})
+
+    content, progress = _download(con, get_func)
+    assert content == CONTENT
+    assert progress.transferred_size == len(CONTENT)
+
+
+def test_resume_without_range_support_does_not_duplicate(con):
+    calls = []
+
+    def get_func(url, **kwargs):
+        calls.append(url)
+        if len(calls) == 1:
+            return FakeResponse(200, CONTENT[:4], SIZE_HEADER)
+        # Server ignores 'Range' header and sends whole file again
+        return FakeResponse(200, CONTENT, SIZE_HEADER)
+
+    content, progress = _download(con, get_func)
+    assert content == CONTENT
+    assert progress.transferred_size == len(CONTENT)
+
+
+def test_download_without_content_length(con):
+    content, _ = _download(
+        con, lambda url, **kwargs: FakeResponse(200, CONTENT)
+    )
+    assert content == CONTENT