858 lines
34 KiB
Python
858 lines
34 KiB
Python
import io
|
|
import os
|
|
import stat
|
|
import tarfile
|
|
import tempfile
|
|
import unittest
|
|
import zipfile
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import requests
|
|
|
|
from core.aria2_installer import (
|
|
ARIA2_INSTALL_ROOT,
|
|
Aria2InstallError,
|
|
_safe_extract_tar,
|
|
_safe_extract_zip,
|
|
)
|
|
from core.download.api import context as downloader_module
|
|
from core.download.api import (
|
|
download_file,
|
|
download_model,
|
|
is_allowed_model_download_filename,
|
|
sanitize_download_filename,
|
|
)
|
|
from core.download.aria2_backend import Aria2Error
|
|
from core.download.validation import _sanitize_download_error
|
|
from core.network_utils import (
|
|
UnsafeUrlError,
|
|
host_matches_domain,
|
|
request_public_url,
|
|
validate_public_http_url,
|
|
)
|
|
from core.path_utils import (
|
|
get_safe_metadata_sidecar_path,
|
|
get_safe_model_resolver_sidecar_path,
|
|
is_path_in_configured_model_roots,
|
|
is_path_within,
|
|
)
|
|
|
|
|
|
class DummyFolderPaths:
|
|
def __init__(self, roots):
|
|
self.folder_names_and_paths = {
|
|
"checkpoints": (roots, {".safetensors"}),
|
|
}
|
|
|
|
def get_folder_paths(self, category):
|
|
value = self.folder_names_and_paths.get(category, ([], set()))
|
|
return list(value[0])
|
|
|
|
|
|
class PathSecurityTests(unittest.TestCase):
|
|
def test_is_path_within_rejects_sibling_directory_with_shared_prefix(self):
|
|
with tempfile.TemporaryDirectory() as temp_root:
|
|
model_root = os.path.join(temp_root, "models")
|
|
sibling_root = os.path.join(temp_root, "models2")
|
|
os.makedirs(model_root)
|
|
os.makedirs(sibling_root)
|
|
|
|
self.assertTrue(is_path_within(model_root, model_root))
|
|
self.assertTrue(is_path_within(os.path.join(model_root, "model.safetensors"), model_root))
|
|
self.assertFalse(is_path_within(os.path.join(sibling_root, "model.safetensors"), model_root))
|
|
|
|
def test_external_model_root_is_allowed(self):
|
|
with tempfile.TemporaryDirectory() as model_root:
|
|
model_path = os.path.join(model_root, "model.safetensors")
|
|
with open(model_path, "wb") as handle:
|
|
handle.write(b"model")
|
|
|
|
self.assertTrue(
|
|
is_path_in_configured_model_roots(
|
|
model_path,
|
|
DummyFolderPaths([model_root]),
|
|
)
|
|
)
|
|
|
|
with tempfile.TemporaryDirectory() as model_root, tempfile.TemporaryDirectory() as other_root:
|
|
outside_path = os.path.join(other_root, "model.safetensors")
|
|
with open(outside_path, "wb") as handle:
|
|
handle.write(b"model")
|
|
|
|
self.assertFalse(
|
|
is_path_in_configured_model_roots(
|
|
outside_path,
|
|
DummyFolderPaths([model_root]),
|
|
)
|
|
)
|
|
|
|
def test_download_filename_is_forced_to_basename(self):
|
|
self.assertEqual(
|
|
"evil.safetensors",
|
|
sanitize_download_filename("../nested/evil.safetensors"),
|
|
)
|
|
self.assertEqual(
|
|
"evil.safetensors",
|
|
sanitize_download_filename(r"C:\temp\evil.safetensors"),
|
|
)
|
|
|
|
def test_download_target_stays_inside_model_directory(self):
|
|
with tempfile.TemporaryDirectory() as model_root:
|
|
with patch("core.download.api.context.get_download_directory", return_value=model_root):
|
|
with patch("core.download.api.context.download_file") as mock_download_file:
|
|
expected_path = os.path.join(model_root, "evil.safetensors")
|
|
mock_download_file.return_value = {
|
|
"success": True,
|
|
"download_id": "pathsecurity",
|
|
"path": expected_path,
|
|
}
|
|
|
|
result = download_model(
|
|
"https://example.com/evil.safetensors",
|
|
"../evil.safetensors",
|
|
"checkpoints",
|
|
download_id="pathsecurity",
|
|
)
|
|
|
|
self.assertTrue(result["success"])
|
|
self.assertEqual(
|
|
os.path.join(model_root, "evil.safetensors"),
|
|
mock_download_file.call_args.args[1],
|
|
)
|
|
|
|
def test_download_subfolder_accepts_both_separator_styles(self):
|
|
with tempfile.TemporaryDirectory() as model_root, patch(
|
|
"core.download.api.context.get_download_directory", return_value=model_root
|
|
), patch("core.download.api.context.download_file") as mock_download_file:
|
|
expected_path = os.path.join(
|
|
model_root,
|
|
"Pony",
|
|
"Styles",
|
|
"model.safetensors",
|
|
)
|
|
mock_download_file.return_value = {
|
|
"success": True,
|
|
"download_id": "portable-subfolder",
|
|
"path": expected_path,
|
|
}
|
|
|
|
result = download_model(
|
|
"https://example.com/model.safetensors",
|
|
"model.safetensors",
|
|
"loras",
|
|
download_id="portable-subfolder",
|
|
subfolder=r"Pony\Styles",
|
|
)
|
|
|
|
self.assertTrue(result["success"])
|
|
self.assertEqual(expected_path, mock_download_file.call_args.args[1])
|
|
|
|
def test_non_model_executable_download_is_rejected(self):
|
|
self.assertFalse(is_allowed_model_download_filename("payload.exe"))
|
|
self.assertTrue(is_allowed_model_download_filename("model.safetensors"))
|
|
|
|
def test_metadata_path_is_always_model_resolver_sidecar(self):
|
|
with tempfile.TemporaryDirectory() as model_root:
|
|
model_path = os.path.join(model_root, "model.safetensors")
|
|
with open(model_path, "wb") as handle:
|
|
handle.write(b"model")
|
|
|
|
metadata_path = get_safe_model_resolver_sidecar_path(model_path)
|
|
self.assertEqual(
|
|
os.path.join(
|
|
model_root,
|
|
"model.safetensors.modelresolver.json",
|
|
),
|
|
metadata_path,
|
|
)
|
|
self.assertNotEqual(model_path, metadata_path)
|
|
|
|
def test_safe_sidecar_paths_reject_empty_input(self):
|
|
for get_safe_sidecar_path in (
|
|
get_safe_metadata_sidecar_path,
|
|
get_safe_model_resolver_sidecar_path,
|
|
):
|
|
with self.subTest(function=get_safe_sidecar_path.__name__), self.assertRaisesRegex(
|
|
ValueError, "A model path is required"
|
|
):
|
|
get_safe_sidecar_path(" ")
|
|
|
|
def test_safe_sidecar_paths_reject_sidecars_outside_model_directory(self):
|
|
with tempfile.TemporaryDirectory() as model_root:
|
|
model_path = os.path.join(model_root, "model.safetensors")
|
|
outside_path = os.path.join(
|
|
os.path.dirname(model_root),
|
|
"outside.modelresolver.json",
|
|
)
|
|
with patch(
|
|
"core.path_utils.get_metadata_sidecar_path",
|
|
return_value=outside_path,
|
|
), patch(
|
|
"core.path_utils.get_model_resolver_sidecar_path",
|
|
return_value=outside_path,
|
|
):
|
|
for get_safe_sidecar_path in (
|
|
get_safe_metadata_sidecar_path,
|
|
get_safe_model_resolver_sidecar_path,
|
|
):
|
|
with self.subTest(function=get_safe_sidecar_path.__name__), self.assertRaisesRegex(
|
|
ValueError,
|
|
"Metadata path is outside the model directory",
|
|
):
|
|
get_safe_sidecar_path(model_path)
|
|
|
|
def test_private_and_local_download_urls_are_rejected(self):
|
|
for url in (
|
|
"http://127.0.0.1:8188/internal",
|
|
"http://10.0.0.1/model.safetensors",
|
|
"http://[::1]/model.safetensors",
|
|
"file:///etc/passwd",
|
|
):
|
|
with self.subTest(url=url):
|
|
with self.assertRaises(UnsafeUrlError):
|
|
validate_public_http_url(url)
|
|
|
|
def test_provider_domain_matching_does_not_accept_lookalikes(self):
|
|
self.assertTrue(host_matches_domain("huggingface.co", "huggingface.co"))
|
|
self.assertTrue(host_matches_domain("cdn.huggingface.co", "huggingface.co"))
|
|
self.assertFalse(
|
|
host_matches_domain("huggingface.co.attacker.example", "huggingface.co")
|
|
)
|
|
self.assertFalse(host_matches_domain("evilcivitai.com", "civitai.com"))
|
|
|
|
def test_redirect_to_loopback_is_blocked_before_second_request(self):
|
|
redirect_response = MagicMock()
|
|
redirect_response.status_code = 302
|
|
redirect_response.headers = {
|
|
"Location": "http://127.0.0.1:8188/private"
|
|
}
|
|
|
|
with patch(
|
|
"core.network_utils.socket.getaddrinfo",
|
|
return_value=[(2, 1, 6, "", ("93.184.216.34", 443))],
|
|
), patch(
|
|
"core.network_utils.requests.get",
|
|
return_value=redirect_response,
|
|
) as mock_get:
|
|
with self.assertRaises(UnsafeUrlError):
|
|
request_public_url(
|
|
"GET",
|
|
"https://example.com/model.safetensors",
|
|
headers={"Authorization": "Bearer secret"},
|
|
)
|
|
|
|
mock_get.assert_called_once()
|
|
redirect_response.close.assert_called_once()
|
|
|
|
def test_authorization_is_preserved_only_for_explicit_xet_bridge(self):
|
|
redirect_response = MagicMock()
|
|
redirect_response.status_code = 302
|
|
redirect_response.headers = {
|
|
"Location": "https://cas-bridge.xethub.hf.co/object?signature=test"
|
|
}
|
|
final_response = MagicMock()
|
|
final_response.status_code = 200
|
|
final_response.headers = {}
|
|
|
|
with patch(
|
|
"core.network_utils.socket.getaddrinfo",
|
|
return_value=[(2, 1, 6, "", ("93.184.216.34", 443))],
|
|
), patch(
|
|
"core.network_utils.requests.get",
|
|
side_effect=[redirect_response, final_response],
|
|
) as mock_get:
|
|
response, final_url, final_headers = request_public_url(
|
|
"GET",
|
|
"https://huggingface.co/example/resolve/main/model.safetensors",
|
|
headers={
|
|
"Authorization": "Bearer secret",
|
|
"Cookie": "session=secret",
|
|
},
|
|
trusted_sensitive_redirect_hosts={"cas-bridge.xethub.hf.co"},
|
|
trusted_sensitive_redirect_headers={"authorization"},
|
|
)
|
|
|
|
self.assertIs(response, final_response)
|
|
self.assertEqual("https://cas-bridge.xethub.hf.co/object?signature=test", final_url)
|
|
self.assertEqual("Bearer secret", final_headers["Authorization"])
|
|
self.assertNotIn("Cookie", final_headers)
|
|
self.assertEqual("Bearer secret", mock_get.call_args_list[1].kwargs["headers"]["Authorization"])
|
|
self.assertNotIn("Cookie", mock_get.call_args_list[1].kwargs["headers"])
|
|
|
|
def test_aria2_backend_does_not_get_replaced_by_hf_xet(self):
|
|
expected = {"success": True, "download_id": "aria-xet"}
|
|
with patch(
|
|
"core.download.api.context._download_backend_from_settings",
|
|
return_value="aria2",
|
|
), patch(
|
|
"core.download.api.context._download_huggingface_xet",
|
|
) as mock_xet, patch(
|
|
"core.download.api.context.download_file_with_aria2",
|
|
return_value=expected,
|
|
) as mock_aria2:
|
|
result = download_file(
|
|
"https://huggingface.co/example/resolve/main/model.safetensors",
|
|
"model.safetensors",
|
|
"aria-xet",
|
|
headers={"Authorization": "Bearer secret"},
|
|
)
|
|
|
|
self.assertEqual(expected, result)
|
|
mock_xet.assert_not_called()
|
|
mock_aria2.assert_called_once()
|
|
|
|
def test_aria2_cannot_follow_an_unvalidated_redirect_after_preflight(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
destination = os.path.join(temp_dir, "model.safetensors")
|
|
Path(destination).write_bytes(b"model")
|
|
with patch(
|
|
"core.download.api.context.load_settings",
|
|
return_value={"download_backend": "aria2"},
|
|
), patch(
|
|
"core.download.api.context._ensure_aria2_daemon",
|
|
), patch(
|
|
"core.download.api.context._resolve_download_url_for_aria2",
|
|
return_value=(
|
|
"https://cas-bridge.xethub.hf.co/object?signature=test",
|
|
{"Authorization": "Bearer secret"},
|
|
),
|
|
), patch(
|
|
"core.download.api.context._aria2_rpc",
|
|
return_value="test-gid",
|
|
) as mock_rpc, patch(
|
|
"core.download.api.context._aria2_tell_status",
|
|
return_value={
|
|
"status": "complete",
|
|
"completedLength": "5",
|
|
"totalLength": "5",
|
|
},
|
|
), patch(
|
|
"core.download.api.context.write_model_resolver_metadata",
|
|
return_value="",
|
|
), patch(
|
|
"core.download.api.context._schedule_aria2_idle_stop",
|
|
):
|
|
result = downloader_module.download_file_with_aria2(
|
|
"https://huggingface.co/example/resolve/main/model.safetensors",
|
|
destination,
|
|
"aria-no-redirect",
|
|
)
|
|
|
|
self.assertTrue(result["success"])
|
|
add_uri_params = mock_rpc.call_args_list[0].args[1]
|
|
self.assertEqual("0", add_uri_params[1]["max-redirect"])
|
|
|
|
def test_aria2_resume_does_not_replace_uri_of_partial_download(self):
|
|
transfer = {
|
|
"gid": "test-gid",
|
|
"path": "model.safetensors",
|
|
}
|
|
rpc_calls = []
|
|
|
|
def fake_rpc(method, params):
|
|
rpc_calls.append((method, params))
|
|
return "OK"
|
|
|
|
with patch.dict(
|
|
downloader_module.aria2_transfers,
|
|
{"resume-xet": transfer},
|
|
clear=True,
|
|
), patch.dict(
|
|
downloader_module.aria2_desired_states,
|
|
{"resume-xet": {"status": "downloading", "seq": 1, "running": True}},
|
|
clear=True,
|
|
), patch.dict(
|
|
downloader_module.download_progress,
|
|
{"resume-xet": {"status": "paused", "speed": 0}},
|
|
clear=True,
|
|
), patch(
|
|
"core.download.api.context._resolve_download_url_for_aria2",
|
|
) as mock_resolve, patch(
|
|
"core.download.api.context._aria2_rpc",
|
|
side_effect=fake_rpc,
|
|
):
|
|
downloader_module._run_aria2_desired_state_worker("resume-xet")
|
|
|
|
self.assertEqual([("aria2.unpause", ["test-gid"])], rpc_calls)
|
|
mock_resolve.assert_not_called()
|
|
|
|
def test_aria2_status_poll_retries_transient_connection_reset(self):
|
|
expected_status = {
|
|
"gid": "test-gid",
|
|
"status": "active",
|
|
"completedLength": "1024",
|
|
}
|
|
transient_error = requests.exceptions.ConnectionError(
|
|
"Connection reset by peer"
|
|
)
|
|
|
|
with patch(
|
|
"core.download.api.context._aria2_rpc",
|
|
side_effect=[transient_error, expected_status],
|
|
) as mock_rpc, patch("core.download.api.context.time.sleep") as mock_sleep:
|
|
status = downloader_module._aria2_tell_status("test-gid")
|
|
|
|
self.assertEqual(expected_status, status)
|
|
self.assertEqual(2, mock_rpc.call_count)
|
|
mock_sleep.assert_called_once_with(
|
|
downloader_module.ARIA2_STATUS_RPC_RETRY_DELAY
|
|
)
|
|
|
|
def test_aria2_can_resume_partial_file_after_rpc_failure(self):
|
|
with tempfile.TemporaryDirectory() as model_root:
|
|
destination = os.path.join(model_root, "model.safetensors")
|
|
Path(destination).write_bytes(b"partial")
|
|
Path(f"{destination}.aria2").write_bytes(b"control")
|
|
expected = {
|
|
"success": True,
|
|
"download_id": "resume-partial",
|
|
"path": destination,
|
|
}
|
|
|
|
with patch(
|
|
"core.download.api.context.get_download_directory",
|
|
return_value=model_root,
|
|
), patch(
|
|
"core.download.api.context._download_backend_from_settings",
|
|
return_value="aria2",
|
|
), patch(
|
|
"core.download.api.context.download_file",
|
|
return_value=expected,
|
|
) as mock_download:
|
|
result = download_model(
|
|
"https://huggingface.co/example/resolve/main/model.safetensors",
|
|
"model.safetensors",
|
|
"checkpoints",
|
|
download_id="resume-partial",
|
|
)
|
|
|
|
self.assertEqual(expected, result)
|
|
mock_download.assert_called_once()
|
|
|
|
def test_arbitrary_aria2_executable_path_is_rejected(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
executable = os.path.join(temp_dir, "aria2c.exe")
|
|
with open(executable, "wb") as handle:
|
|
handle.write(b"not an executable")
|
|
|
|
with patch("core.download.api.context.shutil.which", return_value=None):
|
|
with self.assertRaises(Aria2Error):
|
|
downloader_module._resolve_aria2c_executable(
|
|
{"aria2c_path": executable}
|
|
)
|
|
|
|
def test_managed_aria2_executable_path_is_allowed(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
executable = os.path.join(temp_dir, "aria2c.exe")
|
|
with open(executable, "wb") as handle:
|
|
handle.write(b"managed executable")
|
|
|
|
with patch("core.download.api.context.MANAGED_ARIA2_ROOT", Path(temp_dir)):
|
|
resolved = downloader_module._resolve_aria2c_executable(
|
|
{"aria2c_path": executable}
|
|
)
|
|
self.assertEqual(os.path.realpath(executable), resolved)
|
|
|
|
def test_default_aria2_root_matches_installer_root(self):
|
|
self.assertEqual(
|
|
os.path.realpath(str(downloader_module.MANAGED_ARIA2_ROOT)),
|
|
os.path.realpath(str(ARIA2_INSTALL_ROOT)),
|
|
)
|
|
|
|
def test_symlinked_managed_aria2_executable_path_is_allowed(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
managed_root = Path(temp_dir) / "managed"
|
|
alias_root = Path(temp_dir) / "alias"
|
|
managed_root.mkdir()
|
|
executable = managed_root / "aria2c.exe"
|
|
executable.write_bytes(b"managed executable")
|
|
|
|
try:
|
|
os.symlink(managed_root, alias_root, target_is_directory=True)
|
|
except (OSError, NotImplementedError) as exc:
|
|
self.skipTest(f"directory symlinks are unavailable: {exc}")
|
|
|
|
with patch("core.download.api.context.MANAGED_ARIA2_ROOT", managed_root):
|
|
resolved = downloader_module._resolve_aria2c_executable(
|
|
{"aria2c_path": str(alias_root / executable.name)}
|
|
)
|
|
|
|
self.assertEqual(os.path.realpath(executable), resolved)
|
|
|
|
def test_huggingface_xet_download_bypasses_signed_http_bridge(self):
|
|
file_metadata = MagicMock()
|
|
file_metadata.size = 5
|
|
file_metadata.xet_file_data = SimpleNamespace(
|
|
file_hash="test-hash",
|
|
refresh_route="https://huggingface.co/xet-token",
|
|
)
|
|
session = MagicMock()
|
|
group = MagicMock()
|
|
session.new_file_download_group.return_value = group
|
|
group.__enter__.return_value = group
|
|
|
|
def fake_start_download_file(file_info, destination_path):
|
|
Path(destination_path).write_bytes(b"model")
|
|
progress_callback = session.new_file_download_group.call_args.kwargs[
|
|
"progress_callback"
|
|
]
|
|
progress_callback(
|
|
SimpleNamespace(
|
|
total_bytes_completed=5,
|
|
total_bytes_completion_rate=5000,
|
|
total_transfer_bytes_completed=5,
|
|
total_transfer_bytes=5,
|
|
total_transfer_bytes_completion_rate=2.5,
|
|
),
|
|
{},
|
|
)
|
|
|
|
group.start_download_file.side_effect = fake_start_download_file
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
destination = os.path.join(temp_dir, "model.safetensors")
|
|
with patch(
|
|
"core.download.api.context.validate_public_http_url",
|
|
side_effect=lambda url: url,
|
|
), patch(
|
|
"huggingface_hub.file_download.get_hf_file_metadata",
|
|
return_value=file_metadata,
|
|
), patch(
|
|
"huggingface_hub.utils._xet.get_xet_session",
|
|
return_value=session,
|
|
create=True,
|
|
), patch(
|
|
"huggingface_hub.utils._xet.xet_headers_without_auth",
|
|
side_effect=lambda headers: headers,
|
|
create=True,
|
|
), patch(
|
|
"hf_xet.XetSession",
|
|
create=True,
|
|
), patch(
|
|
"hf_xet.XetFileInfo",
|
|
side_effect=lambda hash, file_size: SimpleNamespace(
|
|
hash=hash,
|
|
file_size=file_size,
|
|
),
|
|
create=True,
|
|
), patch(
|
|
"core.download.api.context.write_model_resolver_metadata",
|
|
return_value=os.path.join(
|
|
temp_dir,
|
|
"model.safetensors.modelresolver.json",
|
|
),
|
|
):
|
|
result = downloader_module._download_huggingface_xet(
|
|
"https://huggingface.co/example/model/resolve/main/model.safetensors",
|
|
destination,
|
|
"xet-test",
|
|
)
|
|
|
|
self.assertTrue(result["success"])
|
|
self.assertEqual(5, result["size"])
|
|
self.assertEqual(b"model", Path(destination).read_bytes())
|
|
group.start_download_file.assert_called_once()
|
|
self.assertEqual(
|
|
200,
|
|
session.new_file_download_group.call_args.kwargs[
|
|
"progress_interval_ms"
|
|
],
|
|
)
|
|
|
|
def test_huggingface_xet_persists_provider_sha256_in_metadata(self):
|
|
expected_sha256 = "a" * 64
|
|
file_metadata = MagicMock()
|
|
file_metadata.size = 5
|
|
file_metadata.xet_file_data = SimpleNamespace(
|
|
file_hash="test-xet-hash",
|
|
refresh_route="https://huggingface.co/xet-token",
|
|
)
|
|
session = MagicMock()
|
|
group = MagicMock()
|
|
session.new_file_download_group.return_value = group
|
|
group.__enter__.return_value = group
|
|
|
|
def fake_start_download_file(file_info, destination_path):
|
|
Path(destination_path).write_bytes(b"model")
|
|
|
|
group.start_download_file.side_effect = fake_start_download_file
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
destination = os.path.join(temp_dir, "model.safetensors")
|
|
metadata_path = os.path.join(
|
|
temp_dir,
|
|
"model.safetensors.modelresolver.json",
|
|
)
|
|
with patch(
|
|
"core.download.api.context.validate_public_http_url",
|
|
side_effect=lambda url: url,
|
|
), patch(
|
|
"huggingface_hub.file_download.get_hf_file_metadata",
|
|
return_value=file_metadata,
|
|
), patch(
|
|
"huggingface_hub.utils._xet.get_xet_session",
|
|
return_value=session,
|
|
create=True,
|
|
), patch(
|
|
"huggingface_hub.utils._xet.xet_headers_without_auth",
|
|
side_effect=lambda headers: headers,
|
|
create=True,
|
|
), patch(
|
|
"hf_xet.XetSession",
|
|
create=True,
|
|
), patch(
|
|
"hf_xet.XetFileInfo",
|
|
side_effect=lambda file_hash, file_size: SimpleNamespace(
|
|
hash=file_hash,
|
|
file_size=file_size,
|
|
),
|
|
create=True,
|
|
), patch(
|
|
"core.sources.huggingface.get_huggingface_file_sha256",
|
|
return_value=expected_sha256,
|
|
), patch(
|
|
"core.download.api.context.write_model_resolver_metadata",
|
|
return_value=metadata_path,
|
|
) as write_metadata:
|
|
result = downloader_module._download_huggingface_xet(
|
|
"https://huggingface.co/example/model/resolve/main/model.safetensors",
|
|
destination,
|
|
"xet-sha256-test",
|
|
)
|
|
|
|
self.assertTrue(result["success"])
|
|
self.assertEqual(expected_sha256, result["sha256"])
|
|
self.assertEqual("huggingface_api", result["sha256_source"])
|
|
metadata_call = write_metadata.call_args.args
|
|
self.assertEqual(expected_sha256, metadata_call[1]["sha256"])
|
|
self.assertEqual(expected_sha256, metadata_call[1]["hashes"]["SHA256"])
|
|
progress = downloader_module.get_progress("xet-sha256-test")
|
|
self.assertEqual(expected_sha256, progress["sha256"])
|
|
self.assertEqual("huggingface_api", progress["sha256_source"])
|
|
|
|
downloader_module.download_progress.pop("xet-sha256-test", None)
|
|
|
|
def test_huggingface_xet_detailed_progress_reports_speed_and_fraction(self):
|
|
download_id = "xet-progress-test"
|
|
total_size = 8 * 1024**3
|
|
downloaded = 64 * 1024**2
|
|
adapter = downloader_module._HuggingFaceXetProgressAdapter(
|
|
download_id,
|
|
total_size,
|
|
0,
|
|
)
|
|
total_update = SimpleNamespace(
|
|
total_bytes_completed=downloaded,
|
|
total_bytes_completion_rate=1024**3,
|
|
total_transfer_bytes_completed=48 * 1024**2,
|
|
total_transfer_bytes=6 * 1024**3,
|
|
total_transfer_bytes_completion_rate=10 * 1024**2,
|
|
)
|
|
|
|
with patch.dict(
|
|
downloader_module.download_progress,
|
|
{download_id: {}},
|
|
clear=True,
|
|
):
|
|
adapter(total_update, [])
|
|
progress = downloader_module.download_progress[download_id]
|
|
|
|
self.assertEqual("downloading", progress["status"])
|
|
self.assertEqual(downloaded, progress["downloaded"])
|
|
self.assertEqual(10 * 1024**2, progress["speed"])
|
|
self.assertEqual(48 * 1024**2, progress["transfer_downloaded"])
|
|
self.assertEqual(6 * 1024**3, progress["transfer_total_size"])
|
|
self.assertEqual(0.8, progress["transfer_progress"])
|
|
self.assertEqual(0.8, progress["progress"])
|
|
|
|
def test_huggingface_xet_without_refresh_helper_uses_native_fallback(self):
|
|
legacy_hf_xet = SimpleNamespace(
|
|
PyItemProgressUpdate=object(),
|
|
PyTotalProgressUpdate=object(),
|
|
PyXetDownloadInfo=object(),
|
|
download_files=MagicMock(),
|
|
)
|
|
real_import = __import__
|
|
|
|
def import_without_refresh_helper(
|
|
name,
|
|
globals=None,
|
|
locals=None,
|
|
fromlist=(),
|
|
level=0,
|
|
):
|
|
if name == "hf_xet":
|
|
return legacy_hf_xet
|
|
if (
|
|
name == "huggingface_hub.utils"
|
|
and "refresh_xet_connection_info" in fromlist
|
|
):
|
|
raise ImportError("refresh_xet_connection_info is unavailable")
|
|
return real_import(name, globals, locals, fromlist, level)
|
|
|
|
incomplete_path = Path("model.safetensors.xet-part")
|
|
xet_file_data = SimpleNamespace(file_hash="test-hash")
|
|
progress_adapter = MagicMock()
|
|
|
|
with patch(
|
|
"builtins.__import__",
|
|
side_effect=import_without_refresh_helper,
|
|
), patch("huggingface_hub.file_download.xet_get") as mock_xet_get:
|
|
downloader_module._run_huggingface_xet_transfer(
|
|
incomplete_path,
|
|
xet_file_data,
|
|
{"Authorization": "Bearer test"},
|
|
1024,
|
|
"model.safetensors",
|
|
progress_adapter,
|
|
)
|
|
|
|
mock_xet_get.assert_called_once_with(
|
|
incomplete_path=incomplete_path,
|
|
xet_file_data=xet_file_data,
|
|
headers={"Authorization": "Bearer test"},
|
|
expected_size=1024,
|
|
displayed_filename="model.safetensors",
|
|
_tqdm_bar=progress_adapter,
|
|
)
|
|
legacy_hf_xet.download_files.assert_not_called()
|
|
|
|
def test_cancel_download_calls_active_huggingface_xet_handle(self):
|
|
download_id = "xet-cancel-handle-test"
|
|
handle = MagicMock()
|
|
|
|
with patch.dict(
|
|
downloader_module.download_progress,
|
|
{download_id: {"status": "downloading"}},
|
|
clear=True,
|
|
), patch.dict(
|
|
downloader_module.xet_transfers,
|
|
{download_id: {"handle": handle, "partial_path": "model.xet-part"}},
|
|
clear=True,
|
|
):
|
|
downloader_module.cancel_download(download_id)
|
|
|
|
handle.cancel.assert_called_once_with()
|
|
downloader_module.cancelled_downloads.discard(download_id)
|
|
|
|
def test_huggingface_xet_partial_cleanup_retries_windows_file_lock(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
partial_path = os.path.join(temp_dir, "model.safetensors.xet-part")
|
|
Path(partial_path).write_bytes(b"partial")
|
|
real_remove = os.remove
|
|
attempts = 0
|
|
|
|
def remove_after_lock_released(path):
|
|
nonlocal attempts
|
|
attempts += 1
|
|
if attempts == 1:
|
|
raise PermissionError("file is still in use")
|
|
real_remove(path)
|
|
|
|
with patch("core.download.api.context.os.remove", side_effect=remove_after_lock_released), patch(
|
|
"core.download.api.context.time.sleep"
|
|
):
|
|
removed = downloader_module._delete_xet_partial_file(partial_path)
|
|
|
|
self.assertTrue(removed)
|
|
self.assertEqual(2, attempts)
|
|
self.assertFalse(os.path.exists(partial_path))
|
|
|
|
def test_huggingface_xet_does_not_infer_network_speed_from_file_progress(self):
|
|
download_id = "xet-no-derived-speed-test"
|
|
adapter = downloader_module._HuggingFaceXetProgressAdapter(
|
|
download_id,
|
|
100 * 1024**2,
|
|
0,
|
|
)
|
|
|
|
with patch.dict(
|
|
downloader_module.download_progress,
|
|
{download_id: {}},
|
|
clear=True,
|
|
):
|
|
adapter(
|
|
SimpleNamespace(
|
|
total_bytes_completed=10 * 1024**2,
|
|
total_bytes_completion_rate=1024**3,
|
|
total_transfer_bytes_completed=0,
|
|
total_transfer_bytes=80 * 1024**2,
|
|
total_transfer_bytes_completion_rate=None,
|
|
),
|
|
{},
|
|
)
|
|
progress = downloader_module.download_progress[download_id]
|
|
|
|
self.assertEqual(10 * 1024**2, progress["downloaded"])
|
|
self.assertEqual(0, progress["speed"])
|
|
|
|
def test_signed_download_query_is_removed_from_error(self):
|
|
error = _sanitize_download_error(
|
|
"403 for url: https://cas-bridge.xethub.hf.co/object?Policy=secret&Signature=secret"
|
|
)
|
|
self.assertEqual(
|
|
"403 for url: https://cas-bridge.xethub.hf.co/object",
|
|
error,
|
|
)
|
|
|
|
def test_tar_links_are_rejected(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
archive_path = os.path.join(temp_dir, "aria2.tar")
|
|
destination = Path(temp_dir) / "extract"
|
|
destination.mkdir()
|
|
with tarfile.open(archive_path, "w") as archive:
|
|
payload = tarfile.TarInfo("safe.txt")
|
|
payload.size = 4
|
|
archive.addfile(payload, io.BytesIO(b"safe"))
|
|
link = tarfile.TarInfo("escape")
|
|
link.type = tarfile.SYMTYPE
|
|
link.linkname = "../../outside"
|
|
archive.addfile(link)
|
|
|
|
with self.assertRaises(Aria2InstallError):
|
|
_safe_extract_tar(Path(archive_path), destination)
|
|
|
|
def test_zip_paths_outside_destination_are_rejected(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
archive_path = Path(temp_dir) / "aria2.zip"
|
|
destination = Path(temp_dir) / "extract"
|
|
destination.mkdir()
|
|
with zipfile.ZipFile(archive_path, "w") as archive:
|
|
archive.writestr("../outside.txt", "unsafe")
|
|
|
|
with self.assertRaises(Aria2InstallError):
|
|
_safe_extract_zip(archive_path, destination)
|
|
|
|
def test_zip_links_are_rejected(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
archive_path = Path(temp_dir) / "aria2.zip"
|
|
destination = Path(temp_dir) / "extract"
|
|
destination.mkdir()
|
|
link = zipfile.ZipInfo("escape")
|
|
link.create_system = 3
|
|
link.external_attr = (stat.S_IFLNK | 0o777) << 16
|
|
with zipfile.ZipFile(archive_path, "w") as archive:
|
|
archive.writestr(link, "../../outside")
|
|
|
|
with self.assertRaises(Aria2InstallError):
|
|
_safe_extract_zip(archive_path, destination)
|
|
|
|
def test_tar_paths_outside_destination_are_rejected(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
archive_path = Path(temp_dir) / "aria2.tar"
|
|
destination = Path(temp_dir) / "extract"
|
|
destination.mkdir()
|
|
with tarfile.open(archive_path, "w") as archive:
|
|
payload = tarfile.TarInfo("../outside.txt")
|
|
payload.size = 6
|
|
archive.addfile(payload, io.BytesIO(b"unsafe"))
|
|
|
|
with self.assertRaises(Aria2InstallError):
|
|
_safe_extract_tar(archive_path, destination)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|