diff --git a/shelfmark/download/clients/alldebrid.py b/shelfmark/download/clients/alldebrid.py index 0555bca1..8e2d404c 100644 --- a/shelfmark/download/clients/alldebrid.py +++ b/shelfmark/download/clients/alldebrid.py @@ -662,13 +662,15 @@ class AllDebridClient(DownloadClient): buf = download_url( direct_link, referer="https://alldebrid.com/", + spool_to_disk=True, ) if not buf: msg = f"Failed to download from {safe_url(direct_link)}" _raise_runtime_error(msg) - with dest.open("wb") as fh: - fh.write(buf.getvalue()) + with buf, dest.open("wb") as fh: + buf.seek(0) + shutil.copyfileobj(buf, fh) with state.lock: state.progress = 50.0 + (idx + 1) / total * 50.0 diff --git a/shelfmark/download/clients/debridlink.py b/shelfmark/download/clients/debridlink.py index 060eba90..775e60d2 100644 --- a/shelfmark/download/clients/debridlink.py +++ b/shelfmark/download/clients/debridlink.py @@ -740,8 +740,11 @@ class DebridLinkClient(DownloadClient): str(file_info.get("downloadUrl", "")), referer="https://debrid-link.com/", cancel_flag=state.cancel_event, + spool_to_disk=True, ) if state.cancel_event.is_set(): + if buf is not None: + buf.close() return if not buf: # The URL is signed and stays out of the message, which ends up @@ -749,7 +752,7 @@ class DebridLinkClient(DownloadClient): msg = f"Failed to download {rel_path} from Debrid-Link" _raise_runtime_error(msg) - with dest.open("wb") as fh: + with buf, dest.open("wb") as fh: buf.seek(0) shutil.copyfileobj(buf, fh) diff --git a/shelfmark/download/clients/realdebrid.py b/shelfmark/download/clients/realdebrid.py index 0a2bd31e..c8ebc96f 100644 --- a/shelfmark/download/clients/realdebrid.py +++ b/shelfmark/download/clients/realdebrid.py @@ -517,13 +517,15 @@ class RealDebridClient(DownloadClient): buf = download_url( direct_url, referer="https://real-debrid.com/", + spool_to_disk=True, ) if not buf: msg = f"Failed to download from {direct_url}" _raise_runtime_error(msg) - with dest.open("wb") as fh: - fh.write(buf.getvalue()) + with buf, dest.open("wb") as fh: + buf.seek(0) + shutil.copyfileobj(buf, fh) with state.lock: state.progress = 50.0 + (idx + 1) / total * 50.0 diff --git a/shelfmark/download/clients/torbox.py b/shelfmark/download/clients/torbox.py index 86a9e0b6..3e4a76de 100644 --- a/shelfmark/download/clients/torbox.py +++ b/shelfmark/download/clients/torbox.py @@ -528,8 +528,11 @@ class TorBoxClient(DownloadClient): direct_url, referer="https://torbox.app/", cancel_flag=state.cancel_event, + spool_to_disk=True, ) if state.cancel_event.is_set(): + if buffer is not None: + buffer.close() return if not buffer: _raise_runtime_error( @@ -538,7 +541,7 @@ class TorBoxClient(DownloadClient): destination = state.target_dir / relative_path destination.parent.mkdir(parents=True, exist_ok=True) - with destination.open("wb") as output: + with buffer, destination.open("wb") as output: buffer.seek(0) shutil.copyfileobj(buffer, output) with state.lock: diff --git a/shelfmark/download/http.py b/shelfmark/download/http.py index c9e821da..1c736315 100644 --- a/shelfmark/download/http.py +++ b/shelfmark/download/http.py @@ -4,7 +4,8 @@ import random import time from http import HTTPStatus from io import BytesIO -from typing import TYPE_CHECKING, NoReturn +from tempfile import SpooledTemporaryFile +from typing import TYPE_CHECKING, BinaryIO, NoReturn from urllib.parse import parse_qsl, urlencode, urljoin, urlparse, urlunparse import requests @@ -881,7 +882,8 @@ def download_url( _selector: network.AAMirrorSelector | None = None, status_callback: Callable[[str, str | None], None] | None = None, referer: str | None = None, -) -> BytesIO | None: + spool_to_disk: bool = False, +) -> BinaryIO | None: """Download content from URL with automatic retry and resume support.""" selector = _selector or network.AAMirrorSelector() current_url = selector.rewrite(link) @@ -904,134 +906,146 @@ def download_url( if cancel_flag and cancel_flag.is_set(): return None - buffer = BytesIO() + buffer = ( + SpooledTemporaryFile(max_size=8 * 1024 * 1024, mode="w+b") # noqa: SIM115 - caller owns successful downloads + if spool_to_disk + else BytesIO() + ) bytes_downloaded = 0 + returned = False try: - if attempt > 0 and status_callback: - status_callback( - "resolving", - f"Connecting (Attempt {attempt + 1}/{MAX_DOWNLOAD_RETRIES})", - ) + try: + if attempt > 0 and status_callback: + status_callback( + "resolving", + f"Connecting (Attempt {attempt + 1}/{MAX_DOWNLOAD_RETRIES})", + ) - logger.info("Downloading (attempt %s/%s)", attempt + 1, MAX_DOWNLOAD_RETRIES) - # Try with CF cookies/UA if available - cookies = _apply_cf_bypass(current_url, headers) - response = requests.get( - current_url, - stream=True, - proxies=get_proxies(current_url), - timeout=REQUEST_TIMEOUT, - cookies=cookies, - headers=headers, - verify=get_ssl_verify(current_url), - hooks=network.REDIRECT_GUARD_HOOKS, - ) - response.raise_for_status() - - if status_callback: - status_callback("downloading", "") - - total_size = total_size or float(response.headers.get("content-length", 0)) - pbar = tqdm(total=total_size, unit="B", unit_scale=True, desc="Downloading") - - for chunk in response.iter_content(chunk_size=8192): - if chunk: - buffer.write(chunk) - bytes_downloaded += len(chunk) - pbar.update(len(chunk)) - if progress_callback and total_size > 0: - progress_callback(bytes_downloaded * 100.0 / total_size) - if cancel_flag and cancel_flag.is_set(): - pbar.close() - return None - pbar.close() - - # Validate - check we didn't get HTML instead of file - if ( - total_size > 0 - and bytes_downloaded < total_size * 0.9 - and response.headers.get("content-type", "").startswith("text/html") - ): - logger.warning("Received HTML instead of file") - return None - - logger.debug("Download completed: %s bytes", bytes_downloaded) - - except requests.exceptions.RequestException as e: - status = _get_status_code(e) - retryable = _is_retryable_error(e) - - # Z-Library 403 - try refreshing cookies via bypasser once before giving up - if ( - status == _HTTP_STATUS_FORBIDDEN - and _is_cf_bypass_enabled() - and not zlib_cookie_refresh_attempted - ): - parsed = urlparse(current_url) - if _is_configured_zlib_host(parsed.hostname) and referer: - zlib_cookie_refresh_attempted = True - logger.info("Z-Library 403 - refreshing cookies via referer") - try: - get_bypassed_page(referer, selector, cancel_flag) - time.sleep(0.5) - # Retry with fresh cookies (don't increment attempt) - continue - except _BYPASSER_ERRORS as cookie_err: - logger.warning( - "Z-Library cookie refresh failed: %s", - type(cookie_err).__name__, - ) - - # Non-retryable errors - if status in _HTTP_STATUS_NON_RETRYABLE: - logger.warning("Download failed (%s)", status) - return None - - # Rate limited - skip to next source immediately - # (waiting doesn't help with concurrent downloads hitting the same server) - if status == _HTTP_STATUS_RATE_LIMITED: - network.note_rate_limited(current_url) - logger.info("Rate limited (429) - trying next source") - if status_callback: - status_callback("resolving", "Server busy, trying next") - return None - - # Timeout - don't retry, server likely overloaded - if isinstance(e, requests.exceptions.Timeout): - logger.warning("Timeout - skipping to next source") - if status_callback: - status_callback("resolving", "Server timed out, trying next") - return None - - # Try to resume if we got some data - if bytes_downloaded > 0 and retryable: - resumed = _try_resume( + logger.info("Downloading (attempt %s/%s)", attempt + 1, MAX_DOWNLOAD_RETRIES) + # Try with CF cookies/UA if available + cookies = _apply_cf_bypass(current_url, headers) + response = requests.get( current_url, - buffer, - bytes_downloaded, - total_size, - progress_callback, - cancel_flag, - headers, + stream=True, + proxies=get_proxies(current_url), + timeout=REQUEST_TIMEOUT, + cookies=cookies, + headers=headers, + verify=get_ssl_verify(current_url), + hooks=network.REDIRECT_GUARD_HOOKS, ) - if resumed: - return resumed + response.raise_for_status() - # Try mirror/DNS rotation if nothing downloaded yet - if bytes_downloaded == 0 and retryable: - new_url = _try_rotation(link, current_url, selector) - if new_url: - current_url = new_url - attempt += 1 - continue + if status_callback: + status_callback("downloading", "") - logger.warning("Download error: %s", type(e).__name__) - if attempt < MAX_DOWNLOAD_RETRIES - 1: - time.sleep(_backoff_delay(attempt + 1)) - attempt += 1 - else: - return buffer + total_size = total_size or float(response.headers.get("content-length", 0)) + pbar = tqdm(total=total_size, unit="B", unit_scale=True, desc="Downloading") + + for chunk in response.iter_content(chunk_size=8192): + if chunk: + buffer.write(chunk) + bytes_downloaded += len(chunk) + pbar.update(len(chunk)) + if progress_callback and total_size > 0: + progress_callback(bytes_downloaded * 100.0 / total_size) + if cancel_flag and cancel_flag.is_set(): + pbar.close() + return None + pbar.close() + + # Validate - check we didn't get HTML instead of file + if ( + total_size > 0 + and bytes_downloaded < total_size * 0.9 + and response.headers.get("content-type", "").startswith("text/html") + ): + logger.warning("Received HTML instead of file") + return None + + logger.debug("Download completed: %s bytes", bytes_downloaded) + + except requests.exceptions.RequestException as e: + status = _get_status_code(e) + retryable = _is_retryable_error(e) + + # Z-Library 403 - try refreshing cookies via bypasser once before giving up + if ( + status == _HTTP_STATUS_FORBIDDEN + and _is_cf_bypass_enabled() + and not zlib_cookie_refresh_attempted + ): + parsed = urlparse(current_url) + if _is_configured_zlib_host(parsed.hostname) and referer: + zlib_cookie_refresh_attempted = True + logger.info("Z-Library 403 - refreshing cookies via referer") + try: + get_bypassed_page(referer, selector, cancel_flag) + time.sleep(0.5) + # Retry with fresh cookies (don't increment attempt) + continue + except _BYPASSER_ERRORS as cookie_err: + logger.warning( + "Z-Library cookie refresh failed: %s", + type(cookie_err).__name__, + ) + + # Non-retryable errors + if status in _HTTP_STATUS_NON_RETRYABLE: + logger.warning("Download failed (%s)", status) + return None + + # Rate limited - skip to next source immediately + # (waiting doesn't help with concurrent downloads hitting the same server) + if status == _HTTP_STATUS_RATE_LIMITED: + network.note_rate_limited(current_url) + logger.info("Rate limited (429) - trying next source") + if status_callback: + status_callback("resolving", "Server busy, trying next") + return None + + # Timeout - don't retry, server likely overloaded + if isinstance(e, requests.exceptions.Timeout): + logger.warning("Timeout - skipping to next source") + if status_callback: + status_callback("resolving", "Server timed out, trying next") + return None + + # Try to resume if we got some data + if bytes_downloaded > 0 and retryable: + resumed = _try_resume( + current_url, + buffer, + bytes_downloaded, + total_size, + progress_callback, + cancel_flag, + headers, + ) + if resumed: + returned = True + return resumed + + # Try mirror/DNS rotation if nothing downloaded yet + if bytes_downloaded == 0 and retryable: + new_url = _try_rotation(link, current_url, selector) + if new_url: + current_url = new_url + attempt += 1 + continue + + logger.warning("Download error: %s", type(e).__name__) + if attempt < MAX_DOWNLOAD_RETRIES - 1: + time.sleep(_backoff_delay(attempt + 1)) + attempt += 1 + else: + returned = True + return buffer + + finally: + if not returned: + buffer.close() logger.error("Download failed after %s attempts", MAX_DOWNLOAD_RETRIES) return None @@ -1057,13 +1071,13 @@ def _is_configured_zlib_host(hostname: str | None) -> bool: def _try_resume( url: str, - buffer: BytesIO, + buffer: BinaryIO, start_byte: int, total_size: float, progress_callback: Callable[[float], None] | None, cancel_flag: Event | None, base_headers: dict | None = None, -) -> BytesIO | None: +) -> BinaryIO | None: """Try to resume an interrupted download.""" for attempt in range(MAX_RESUME_ATTEMPTS): logger.info( diff --git a/tests/download/test_http_download_url.py b/tests/download/test_http_download_url.py index 2de96306..a44c63a5 100644 --- a/tests/download/test_http_download_url.py +++ b/tests/download/test_http_download_url.py @@ -336,3 +336,69 @@ def test_default_download_does_not_log_the_rotated_url(monkeypatch): assert http.download_url(source_url) is None assert rotated_url not in str(logger.mock_calls) + + +def test_spooled_download_rolls_to_disk_and_retains_contents(monkeypatch): + http = _prepare_download_test(monkeypatch) + payload = b"a" * (8 * 1024 * 1024 + 1) + monkeypatch.setattr( + http.requests, + "get", + lambda url, **kwargs: _FakeResponse( + 200, headers={"content-length": str(len(payload))}, chunks=[payload], url=url + ), + ) + result = http.download_url("https://example.com/file", spool_to_disk=True) + assert result is not None + with result: + assert result._rolled + result.seek(0) + assert result.read() == payload + + +def test_spooled_download_closes_temporary_file_on_error(monkeypatch): + http = _prepare_download_test(monkeypatch) + created = [] + factory = http.SpooledTemporaryFile + + def track(*args, **kwargs): + stream = factory(*args, **kwargs) + created.append(stream) + return stream + + monkeypatch.setattr(http, "SpooledTemporaryFile", track) + monkeypatch.setattr( + http.requests, + "get", + lambda *args, **kwargs: (_ for _ in ()).throw(requests.exceptions.Timeout()), + ) + assert http.download_url("https://example.com/file", spool_to_disk=True) is None + assert len(created) == 1 and created[0].closed + + +def test_spooled_download_resumes_and_returns_complete_content(monkeypatch): + http = _prepare_download_test(monkeypatch) + responses = iter( + [ + _FakeResponse( + 200, + headers={"content-length": "8"}, + chunks=[b"abcd"], + iter_error=requests.exceptions.ConnectionError("interrupted"), + ), + _FakeResponse(206, chunks=[b"efgh"]), + ] + ) + calls = [] + + def fake_get(url, **kwargs): + calls.append(kwargs) + return next(responses) + + monkeypatch.setattr(http.requests, "get", fake_get) + result = http.download_url("https://example.com/file", spool_to_disk=True) + assert result is not None + with result: + result.seek(0) + assert result.read() == b"abcdefgh" + assert calls[1]["headers"]["Range"] == "bytes=4-" diff --git a/tests/prowlarr/test_debridlink_client.py b/tests/prowlarr/test_debridlink_client.py index b8f9d004..f0eca5bc 100644 --- a/tests/prowlarr/test_debridlink_client.py +++ b/tests/prowlarr/test_debridlink_client.py @@ -617,7 +617,7 @@ class TestDownloadPipeline: state = _state(tmp_path, phase="downloading_http") fetched: list[str] = [] - def fake_download(url, referer=None, cancel_flag=None): + def fake_download(url, referer=None, cancel_flag=None, spool_to_disk=False): fetched.append(url) return _buffer(payload) @@ -678,7 +678,7 @@ class TestDownloadPipeline: state = _state(target, phase="downloading_http") monkeypatch.setattr( "shelfmark.download.clients.debridlink.download_url", - lambda url, referer=None, cancel_flag=None: _buffer(), + lambda url, referer=None, cancel_flag=None, spool_to_disk=False: _buffer(), ) client._process_and_download(state, [{"name": name, "downloadUrl": "https://dl.example/x"}]) @@ -693,7 +693,9 @@ class TestDownloadPipeline: state = _state(tmp_path, phase="downloading_http") monkeypatch.setattr( "shelfmark.download.clients.debridlink.download_url", - lambda url, referer=None, cancel_flag=None: _buffer(next(payloads)), + lambda url, referer=None, cancel_flag=None, spool_to_disk=False: _buffer( + next(payloads) + ), ) client._process_and_download( @@ -712,7 +714,7 @@ class TestDownloadPipeline: state = _state(tmp_path, phase="downloading_http") monkeypatch.setattr( "shelfmark.download.clients.debridlink.download_url", - lambda url, referer=None, cancel_flag=None: None, + lambda url, referer=None, cancel_flag=None, spool_to_disk=False: None, ) client._process_and_download( @@ -730,7 +732,7 @@ class TestDownloadPipeline: monkeypatch.setattr(client, "_fetch_file_list", listing) monkeypatch.setattr( "shelfmark.download.clients.debridlink.download_url", - lambda url, referer=None, cancel_flag=None: _buffer(), + lambda url, referer=None, cancel_flag=None, spool_to_disk=False: _buffer(), ) client._process_and_download(state) @@ -743,7 +745,7 @@ class TestDownloadPipeline: state = _state(tmp_path, phase="downloading_http") fetched: list[str] = [] - def fake_download(url, referer=None, cancel_flag=None): + def fake_download(url, referer=None, cancel_flag=None, spool_to_disk=False): fetched.append(url) state.cancel_event.set() return _buffer()