diff --git a/kikotools/tools/model_downloader/base.py b/kikotools/tools/model_downloader/base.py index eb5946b..928840d 100644 --- a/kikotools/tools/model_downloader/base.py +++ b/kikotools/tools/model_downloader/base.py @@ -88,7 +88,7 @@ def safe_server_filename(name: Optional[str], default: str) -> str: class _SafeRedirectHandler(urllib.request.HTTPRedirectHandler): - """Only follow http(s) redirects and drop credentials on cross-host hops.""" + """Only follow http(s) redirects; drop credentials on cross-host or https-to-http hops.""" def redirect_request(self, req, fp, code, msg, headers, newurl): if urlparse(newurl).scheme.lower() not in ALLOWED_SCHEMES: @@ -96,8 +96,9 @@ class _SafeRedirectHandler(urllib.request.HTTPRedirectHandler): newurl, code, "Redirect to disallowed URL scheme", headers, fp ) new_request = super().redirect_request(req, fp, code, msg, headers, newurl) + old, new = urlparse(req.full_url), urlparse(newurl) if new_request is not None and ( - urlparse(newurl).netloc.lower() != urlparse(req.full_url).netloc.lower() + new.netloc.lower() != old.netloc.lower() or new.scheme.lower() != "https" ): new_request.remove_header("Authorization") return new_request diff --git a/kikotools/tools/model_downloader/civitai.py b/kikotools/tools/model_downloader/civitai.py index 66a1f7b..9cfb156 100644 --- a/kikotools/tools/model_downloader/civitai.py +++ b/kikotools/tools/model_downloader/civitai.py @@ -156,7 +156,10 @@ class CivitAIDownloader(BaseDownloader): # Downloads go through civitai.com whichever site the link came from if "/api/download/models/" in url: - url = urllib.parse.urlunparse(parsed_url._replace(netloc="civitai.com")) + # https only: the request carries the token + url = urllib.parse.urlunparse( + parsed_url._replace(scheme="https", netloc="civitai.com") + ) else: ids = self._parse_civitai_url(url) diff --git a/kikotools/tools/model_downloader/huggingface.py b/kikotools/tools/model_downloader/huggingface.py index 31e18c4..0815e90 100644 --- a/kikotools/tools/model_downloader/huggingface.py +++ b/kikotools/tools/model_downloader/huggingface.py @@ -5,7 +5,7 @@ import os import urllib.request import urllib.error from typing import Optional -from urllib.parse import urlparse, quote +from urllib.parse import urlparse, urlunparse, quote from .base import ( USER_AGENT, @@ -131,6 +131,8 @@ class HuggingFaceDownloader(BaseDownloader): logger.info("Converted URL to: %s", download_url) else: download_url = url + # https only: the request may carry the server's token + download_url = urlunparse(urlparse(download_url)._replace(scheme="https")) if not filename: if url_info["filename"]: diff --git a/tests/unit/tools/model_downloader/test_downloads.py b/tests/unit/tools/model_downloader/test_downloads.py index f528f7e..a325a66 100644 --- a/tests/unit/tools/model_downloader/test_downloads.py +++ b/tests/unit/tools/model_downloader/test_downloads.py @@ -282,6 +282,12 @@ class TestRedirectHandler: ) assert new.get_header("Authorization") == "Bearer s" + def test_https_to_http_drops_authorization(self): + new = self.redirect( + "https://a.com/x", "http://a.com/y", {"Authorization": "Bearer s"} + ) + assert not new.has_header("Authorization") + def test_disallowed_scheme_blocked(self): with pytest.raises(urllib.error.HTTPError): self.redirect("https://a.com/x", "file:///etc/passwd", {}) @@ -510,6 +516,14 @@ class TestCustom: class TestHuggingFace: + def test_http_link_is_fetched_over_https(self, root, out, opener): + _, op = opener + op.open.return_value = FakeResponse(b"d", headers={"Content-Length": "1"}) + make(HuggingFaceDownloader, root, token="hf_x").download( + "http://huggingface.co/u/r/resolve/main/m.safetensors", out + ) + assert op.open.call_args.args[0].full_url.startswith("https://") + def test_download_blob_url_converted(self, root, out, opener): _, op = opener op.open.return_value = FakeResponse(b"d", headers={"Content-Length": "1"}) @@ -794,6 +808,16 @@ class TestCivitAI: == "https://civitai.com/api/download/models/77?type=Model" ) + def test_http_api_link_is_fetched_over_https(self, root, out, opener): + _, op = opener + op.open.return_value = FakeResponse(b"d", status=200) + make(CivitAIDownloader, root, token="t").download( + "http://civitai.red/api/download/models/77", out, filename="m.safetensors" + ) + assert op.open.call_args.args[0].full_url == ( + "https://civitai.com/api/download/models/77" + ) + def test_red_model_page_resolves_version(self, root, out, opener): _, op = opener op.open.return_value = FakeResponse(b"d", status=200)