diff --git a/python/lib/sift_client/_internal/util/file.py b/python/lib/sift_client/_internal/util/file.py index ac1b83f57..9ab2b3ffc 100644 --- a/python/lib/sift_client/_internal/util/file.py +++ b/python/lib/sift_client/_internal/util/file.py @@ -94,6 +94,7 @@ def download_file( *, rest_client: RestClient, show_progress: bool = False, + extra_headers: dict[str, str] | None = None, ) -> Path: """Download a file from a URL in streaming 4 MiB chunks. @@ -103,6 +104,13 @@ def download_file( rest_client: The SDK rest client to use for the download. show_progress: If True, display a progress bar during download. Defaults to False. + extra_headers: Optional additional headers to include on the request. + Merged on top of the internal ``Authorization: None`` strip, so a + caller-supplied ``Authorization`` (if any) wins. Useful for + presigned URLs whose signature covers a specific ``Host`` header + that differs from the URL's netloc (e.g. an S3 endpoint accessed + through a container-internal alias while the URL was signed + against the browser-facing loopback authority). Returns: The path to the downloaded file. @@ -111,8 +119,13 @@ def download_file( requests.HTTPError: If the download request fails. """ output_path.parent.mkdir(parents=True, exist_ok=True) - # Strip the session's default Authorization header, presigned URLs carry their own auth - with rest_client.get(signed_url, stream=True, headers={"Authorization": None}) as response: + # Strip the session's default Authorization header (presigned URLs carry + # their own auth). Any ``extra_headers`` merge on top so callers can + # override the Host header, add a range, etc. + headers: dict[str, str | None] = {"Authorization": None} + if extra_headers: + headers.update(extra_headers) + with rest_client.get(signed_url, stream=True, headers=headers) as response: response.raise_for_status() total_bytes = int(response.headers.get("Content-Length", 0)) or None with alive_bar( diff --git a/python/lib/sift_client/_tests/_internal/util/__init__.py b/python/lib/sift_client/_tests/_internal/util/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/lib/sift_client/_tests/_internal/util/test_file.py b/python/lib/sift_client/_tests/_internal/util/test_file.py new file mode 100644 index 000000000..1d14014f0 --- /dev/null +++ b/python/lib/sift_client/_tests/_internal/util/test_file.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +from unittest.mock import MagicMock + +from sift_client._internal.util.file import download_file + + +class _FakeResponse: + """Minimal ``requests.Response`` stand-in for the streamed-download tests.""" + + def __init__(self, chunks: list[bytes]): + self._chunks = chunks + self.headers: dict[str, str] = {} + + def raise_for_status(self) -> None: ... + + def iter_content(self, chunk_size: int): + return iter(self._chunks) + + def __enter__(self): + return self + + def __exit__(self, *_): + return False + + +def _install_rest_client_stub(response: _FakeResponse) -> tuple[MagicMock, dict]: + """Return a rest_client mock whose ``.get(...)`` records call kwargs.""" + captured: dict = {} + rest_client = MagicMock() + + def _fake_get(url, **kwargs): + captured["url"] = url + captured.update(kwargs) + return response + + rest_client.get.side_effect = _fake_get + return rest_client, captured + + +class TestDownloadFileExtraHeaders: + """``extra_headers`` merges on top of the internal ``Authorization: None`` + strip so callers can override the Host header (or add a range, etc.) for + presigned URLs whose signature covers a specific ``Host`` distinct from the + URL's authority. Without this the caller has to bypass ``download_file`` + entirely for signed-URL edge cases (e.g. local minio via a container- + internal alias) and reimplement streaming + auth stripping from scratch. + """ + + def test_no_extra_headers_strips_only_authorization(self, tmp_path): + rest_client, captured = _install_rest_client_stub(_FakeResponse([b"hello ", b"world"])) + target = tmp_path / "out.bin" + + download_file( + "https://example.com/file?sig=1", + target, + rest_client=rest_client, + ) + + assert captured["headers"] == {"Authorization": None} + assert target.read_bytes() == b"hello world" + + def test_extra_headers_are_forwarded(self, tmp_path): + rest_client, captured = _install_rest_client_stub(_FakeResponse([b"payload"])) + target = tmp_path / "out.bin" + + download_file( + "https://example.com/file?sig=1", + target, + rest_client=rest_client, + extra_headers={"Host": "signed.example.com", "X-Trace": "abc"}, + ) + + assert captured["headers"] == { + "Authorization": None, + "Host": "signed.example.com", + "X-Trace": "abc", + } + assert target.read_bytes() == b"payload" + + def test_extra_headers_can_override_authorization_strip(self, tmp_path): + # The strip is a default, not a hard invariant -- if a caller has a + # concrete reason to pass an Authorization on the download request + # (e.g. a proxy-scoped bearer), let their value win. + rest_client, captured = _install_rest_client_stub(_FakeResponse([b"payload"])) + target = tmp_path / "out.bin" + + download_file( + "https://example.com/file?sig=1", + target, + rest_client=rest_client, + extra_headers={"Authorization": "Bearer proxy-token"}, + ) + + assert captured["headers"]["Authorization"] == "Bearer proxy-token"