xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-11-12 22:57:03 +08:00
Code Issues Packages Projects Releases Wiki Activity

Restrict downloading to safetensor files only.

Browse Source
This commit is contained in:
Robin Huang 2024-08-13 12:29:35 -07:00
parent bc85c0dc38
commit 5cf97f2a58
2 changed files with 30 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
model_filemanager/download_models.py
View File

@ -204,7 +204,7 @@ def validate_model_subdirectory(model_subdirectory: str) -> bool:
return True return True
def validate_filename(filename): def validate_filename(filename: str)-> bool:
""" """
Validate a filename to ensure it's safe and doesn't contain any path traversal attempts. Validate a filename to ensure it's safe and doesn't contain any path traversal attempts.
@ -214,6 +214,9 @@ def validate_filename(filename):
Returns: Returns:
bool: True if the filename is valid, False otherwise bool: True if the filename is valid, False otherwise
""" """
if not filename.lower().endswith(('.sft', '.safetensors')):
return False
# Check if the filename is empty, None, or just whitespace # Check if the filename is empty, None, or just whitespace
if not filename or not filename.strip(): if not filename or not filename.strip():
return False return False

54
tests-unit/prompt_server_test/download_models_test.py
View File

@ -59,22 +59,22 @@ async def test_download_model_success():
mock_open.return_value.__enter__.return_value = mock_file mock_open.return_value.__enter__.return_value = mock_file
time_values = itertools.count(0, 0.1) time_values = itertools.count(0, 0.1)
with patch('model_filemanager.create_model_path', return_value=('models/checkpoints/model.bin', 'checkpoints/model.bin')), \ with patch('model_filemanager.create_model_path', return_value=('models/checkpoints/model.sft', 'checkpoints/model.sft')), \
patch('model_filemanager.check_file_exists', return_value=None), \ patch('model_filemanager.check_file_exists', return_value=None), \
patch('builtins.open', mock_open), \ patch('builtins.open', mock_open), \
patch('time.time', side_effect=time_values): # Simulate time passing patch('time.time', side_effect=time_values): # Simulate time passing
result = await download_model( result = await download_model(
mock_make_request, mock_make_request,
'model.bin', 'model.sft',
'http://example.com/model.bin', 'http://example.com/model.sft',
'checkpoints', 'checkpoints',
mock_progress_callback mock_progress_callback
) )
# Assert the result # Assert the result
assert isinstance(result, DownloadModelStatus) assert isinstance(result, DownloadModelStatus)
assert result.message == 'Successfully downloaded model.bin' assert result.message == 'Successfully downloaded model.sft'
assert result.status == 'completed' assert result.status == 'completed'
assert result.already_existed is False assert result.already_existed is False
@ -83,14 +83,14 @@ async def test_download_model_success():
# Check initial call # Check initial call
mock_progress_callback.assert_any_call( mock_progress_callback.assert_any_call(
'checkpoints/model.bin', 'checkpoints/model.sft',
DownloadModelStatus(DownloadStatusType.PENDING, 0, "Starting download of model.bin", False) DownloadModelStatus(DownloadStatusType.PENDING, 0, "Starting download of model.sft", False)
) )
# Check final call # Check final call
mock_progress_callback.assert_any_call( mock_progress_callback.assert_any_call(
'checkpoints/model.bin', 'checkpoints/model.sft',
DownloadModelStatus(DownloadStatusType.COMPLETED, 100, "Successfully downloaded model.bin", False) DownloadModelStatus(DownloadStatusType.COMPLETED, 100, "Successfully downloaded model.sft", False)
) )
# Verify file writing # Verify file writing
@ -99,7 +99,7 @@ async def test_download_model_success():
mock_file.write.assert_any_call(b'c' * 200) mock_file.write.assert_any_call(b'c' * 200)
# Verify request was made # Verify request was made
mock_make_request.assert_called_once_with('http://example.com/model.bin') mock_make_request.assert_called_once_with('http://example.com/model.sft')
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_download_model_url_request_failure(): async def test_download_model_url_request_failure():
@ -160,8 +160,8 @@ async def test_download_model_invalid_model_subdirectory():
result = await download_model( result = await download_model(
mock_make_request, mock_make_request,
'model.bin', 'model.sft',
'http://example.com/model.bin', 'http://example.com/model.sft',
'../bad_path', '../bad_path',
mock_progress_callback mock_progress_callback
) )
@ -178,7 +178,7 @@ def test_create_model_path(tmp_path, monkeypatch):
mock_models_dir = tmp_path / "models" mock_models_dir = tmp_path / "models"
monkeypatch.setattr('folder_paths.models_dir', str(mock_models_dir)) monkeypatch.setattr('folder_paths.models_dir', str(mock_models_dir))
model_name = "test_model.bin" model_name = "test_model.sft"
model_directory = "test_dir" model_directory = "test_dir"
file_path, relative_path = create_model_path(model_name, model_directory, mock_models_dir) file_path, relative_path = create_model_path(model_name, model_directory, mock_models_dir)
@ -190,30 +190,30 @@ def test_create_model_path(tmp_path, monkeypatch):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_check_file_exists_when_file_exists(tmp_path): async def test_check_file_exists_when_file_exists(tmp_path):
file_path = tmp_path / "existing_model.bin" file_path = tmp_path / "existing_model.sft"
file_path.touch() # Create an empty file file_path.touch() # Create an empty file
mock_callback = AsyncMock() mock_callback = AsyncMock()
result = await check_file_exists(str(file_path), "existing_model.bin", mock_callback, "test/existing_model.bin") result = await check_file_exists(str(file_path), "existing_model.sft", mock_callback, "test/existing_model.sft")
assert result is not None assert result is not None
assert result.status == "completed" assert result.status == "completed"
assert result.message == "existing_model.bin already exists" assert result.message == "existing_model.sft already exists"
assert result.already_existed is True assert result.already_existed is True
mock_callback.assert_called_once_with( mock_callback.assert_called_once_with(
"test/existing_model.bin", "test/existing_model.sft",
DownloadModelStatus(DownloadStatusType.COMPLETED, 100, "existing_model.bin already exists", already_existed=True) DownloadModelStatus(DownloadStatusType.COMPLETED, 100, "existing_model.sft already exists", already_existed=True)
) )
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_check_file_exists_when_file_does_not_exist(tmp_path): async def test_check_file_exists_when_file_does_not_exist(tmp_path):
file_path = tmp_path / "non_existing_model.bin" file_path = tmp_path / "non_existing_model.sft"
mock_callback = AsyncMock() mock_callback = AsyncMock()
result = await check_file_exists(str(file_path), "non_existing_model.bin", mock_callback, "test/non_existing_model.bin") result = await check_file_exists(str(file_path), "non_existing_model.sft", mock_callback, "test/non_existing_model.sft")
assert result is None assert result is None
mock_callback.assert_not_called() mock_callback.assert_not_called()
@ -229,15 +229,15 @@ async def test_track_download_progress_no_content_length():
with patch('builtins.open', mock_open): with patch('builtins.open', mock_open):
result = await track_download_progress( result = await track_download_progress(
mock_response, '/mock/path/model.bin', 'model.bin', mock_response, '/mock/path/model.sft', 'model.sft',
mock_callback, 'models/model.bin', interval=0.1 mock_callback, 'models/model.sft', interval=0.1
) )
assert result.status == "completed" assert result.status == "completed"
# Check that progress was reported even without knowing the total size # Check that progress was reported even without knowing the total size
mock_callback.assert_any_call( mock_callback.assert_any_call(
'models/model.bin', 'models/model.sft',
DownloadModelStatus(DownloadStatusType.IN_PROGRESS, 0, "Downloading model.bin", already_existed=False) DownloadModelStatus(DownloadStatusType.IN_PROGRESS, 0, "Downloading model.sft", already_existed=False)
) )
@pytest.mark.asyncio @pytest.mark.asyncio
@ -256,8 +256,8 @@ async def test_track_download_progress_interval():
with patch('builtins.open', mock_open), \ with patch('builtins.open', mock_open), \
patch('time.time', mock_time): patch('time.time', mock_time):
await track_download_progress( await track_download_progress(
mock_response, '/mock/path/model.bin', 'model.bin', mock_response, '/mock/path/model.sft', 'model.sft',
mock_callback, 'models/model.bin', interval=1.0 mock_callback, 'models/model.sft', interval=1.0
) )
# Print out the actual call count and the arguments of each call for debugging # Print out the actual call count and the arguments of each call for debugging
@ -303,12 +303,10 @@ def test_empty_subdirectory():
@pytest.mark.parametrize("filename, expected", [ @pytest.mark.parametrize("filename, expected", [
("valid_model.safetensors", True), ("valid_model.safetensors", True),
("valid_model.sft", True), ("valid_model.sft", True),
("another-valid_model.ckpt", True),
("valid model.safetensors", True), # Test with space ("valid model.safetensors", True), # Test with space
("UPPERCASE_MODEL.SAFETENSORS", True), ("UPPERCASE_MODEL.SAFETENSORS", True),
("model_with.multiple.dots.pt", True), ("model_with.multiple.dots.pt", False),
("", False), # Empty string ("", False), # Empty string
(None, False), # None value
("../../../etc/passwd", False), # Path traversal attempt ("../../../etc/passwd", False), # Path traversal attempt
("/etc/passwd", False), # Absolute path ("/etc/passwd", False), # Absolute path
("\\windows\\system32\\config\\sam", False), # Windows path ("\\windows\\system32\\config\\sam", False), # Windows path