wehub-resource-sync
393 lines
14 KiB
Python
393 lines
14 KiB
Python
import filecmp
|
|
import hashlib
|
|
import io
|
|
import os
|
|
import shutil
|
|
import stat
|
|
import tarfile
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from pyspark.sql import SparkSession
|
|
|
|
import mlflow
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.pyfunc.dbconnect_artifact_cache import extract_archive_to_dir
|
|
from mlflow.utils import file_utils
|
|
from mlflow.utils.file_utils import (
|
|
TempDir,
|
|
_copy_file_or_tree,
|
|
_handle_readonly_on_windows,
|
|
check_tarfile_security,
|
|
get_parent_dir,
|
|
get_total_file_size,
|
|
local_file_uri_to_path,
|
|
)
|
|
from mlflow.utils.os import is_windows
|
|
|
|
from tests.helper_functions import random_int
|
|
from tests.projects.utils import TEST_PROJECT_DIR
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def spark_session():
|
|
with SparkSession.builder.master("local[*]").getOrCreate() as session:
|
|
yield session
|
|
|
|
|
|
def test_mkdir(tmp_path):
|
|
temp_dir = str(tmp_path)
|
|
new_dir_name = f"mkdir_test_{random_int()}"
|
|
file_utils.mkdir(temp_dir, new_dir_name)
|
|
assert os.listdir(temp_dir) == [new_dir_name]
|
|
|
|
with pytest.raises(OSError, match="bad directory"):
|
|
file_utils.mkdir("/ bad directory @ name ", "ouch")
|
|
|
|
# does not raise if directory exists already
|
|
file_utils.mkdir(temp_dir, new_dir_name)
|
|
|
|
# raises if it exists already but is a file
|
|
dummy_file_path = str(tmp_path.joinpath("dummy_file"))
|
|
with open(dummy_file_path, "a"):
|
|
pass
|
|
|
|
with pytest.raises(OSError, match="exists"):
|
|
file_utils.mkdir(dummy_file_path)
|
|
|
|
|
|
def test_make_tarfile(tmp_path):
|
|
# Tar a local project
|
|
tarfile0 = str(tmp_path.joinpath("first-tarfile"))
|
|
file_utils.make_tarfile(
|
|
output_filename=tarfile0, source_dir=TEST_PROJECT_DIR, archive_name="some-archive"
|
|
)
|
|
# Copy local project into a temp dir
|
|
dst_dir = str(tmp_path.joinpath("project-directory"))
|
|
shutil.copytree(TEST_PROJECT_DIR, dst_dir)
|
|
# Tar the copied project
|
|
tarfile1 = str(tmp_path.joinpath("second-tarfile"))
|
|
file_utils.make_tarfile(
|
|
output_filename=tarfile1, source_dir=dst_dir, archive_name="some-archive"
|
|
)
|
|
# Compare the archives & explicitly verify their SHA256 hashes match (i.e. that
|
|
# changes in file modification timestamps don't affect the archive contents)
|
|
assert filecmp.cmp(tarfile0, tarfile1, shallow=False)
|
|
with open(tarfile0, "rb") as first_tar, open(tarfile1, "rb") as second_tar:
|
|
assert (
|
|
hashlib.sha256(first_tar.read()).hexdigest()
|
|
== hashlib.sha256(second_tar.read()).hexdigest()
|
|
)
|
|
# Extract the TAR and check that its contents match the original directory
|
|
extract_dir = str(tmp_path.joinpath("extracted-tar"))
|
|
os.makedirs(extract_dir)
|
|
with tarfile.open(tarfile0, "r:gz") as handle:
|
|
handle.extractall(path=extract_dir)
|
|
dir_comparison = filecmp.dircmp(os.path.join(extract_dir, "some-archive"), TEST_PROJECT_DIR)
|
|
assert len(dir_comparison.left_only) == 0
|
|
assert len(dir_comparison.right_only) == 0
|
|
assert len(dir_comparison.diff_files) == 0
|
|
assert len(dir_comparison.funny_files) == 0
|
|
|
|
|
|
def test_get_parent_dir(tmp_path):
|
|
child_dir = tmp_path.joinpath("dir")
|
|
child_dir.mkdir()
|
|
assert str(tmp_path) == get_parent_dir(str(child_dir))
|
|
|
|
|
|
def test_file_copy():
|
|
with TempDir() as tmp:
|
|
file_path = tmp.path("test_file.txt")
|
|
copy_path = tmp.path("test_dir1/")
|
|
os.mkdir(copy_path)
|
|
with open(file_path, "a") as f:
|
|
f.write("testing")
|
|
_copy_file_or_tree(file_path, copy_path, "")
|
|
assert filecmp.cmp(file_path, os.path.join(copy_path, "test_file.txt"))
|
|
|
|
|
|
def test_dir_create():
|
|
with TempDir() as tmp:
|
|
file_path = tmp.path("test_file.txt")
|
|
create_dir = tmp.path("test_dir2/")
|
|
with open(file_path, "a") as f:
|
|
f.write("testing")
|
|
name = _copy_file_or_tree(file_path, file_path, create_dir)
|
|
assert filecmp.cmp(file_path, name)
|
|
|
|
|
|
def test_dir_copy():
|
|
with TempDir() as tmp:
|
|
dir_path = tmp.path("test_dir1/")
|
|
copy_path = tmp.path("test_dir2")
|
|
os.mkdir(dir_path)
|
|
with open(os.path.join(dir_path, "test_file.txt"), "a") as f:
|
|
f.write("testing")
|
|
_copy_file_or_tree(dir_path, copy_path, "")
|
|
assert filecmp.dircmp(dir_path, copy_path)
|
|
|
|
|
|
@pytest.mark.skipif(not is_windows(), reason="requires Windows")
|
|
def test_handle_readonly_on_windows(tmp_path):
|
|
tmp_path = tmp_path.joinpath("file")
|
|
with open(tmp_path, "w"):
|
|
pass
|
|
|
|
# Make the file read-only
|
|
os.chmod(tmp_path, stat.S_IREAD | stat.S_IRGRP | stat.S_IROTH)
|
|
# Ensure the file can't be removed
|
|
with pytest.raises(PermissionError, match="Access is denied") as exc:
|
|
os.unlink(tmp_path)
|
|
|
|
_handle_readonly_on_windows(
|
|
os.unlink,
|
|
tmp_path,
|
|
(exc.type, exc.value, exc.traceback),
|
|
)
|
|
assert not os.path.exists(tmp_path)
|
|
|
|
|
|
@pytest.mark.skipif(not is_windows(), reason="This test only passes on Windows")
|
|
@pytest.mark.parametrize(
|
|
("input_uri", "expected_path"),
|
|
[
|
|
(r"\\my_server\my_path\my_sub_path", r"\\my_server\my_path\my_sub_path"),
|
|
],
|
|
)
|
|
def test_local_file_uri_to_path_on_windows(input_uri, expected_path):
|
|
assert local_file_uri_to_path(input_uri) == expected_path
|
|
|
|
|
|
def test_shutil_copytree_without_file_permissions(tmp_path):
|
|
src_dir = tmp_path.joinpath("src-dir")
|
|
src_dir.mkdir()
|
|
dst_dir = tmp_path.joinpath("dst-dir")
|
|
dst_dir.mkdir()
|
|
# Test copying empty directory
|
|
mlflow.utils.file_utils.shutil_copytree_without_file_permissions(src_dir, dst_dir)
|
|
assert len(os.listdir(dst_dir)) == 0
|
|
# Test copying directory with contents
|
|
src_dir.joinpath("subdir").mkdir()
|
|
src_dir.joinpath("subdir/subdir-file.txt").write_text("testing 123")
|
|
src_dir.joinpath("top-level-file.txt").write_text("hi")
|
|
mlflow.utils.file_utils.shutil_copytree_without_file_permissions(src_dir, dst_dir)
|
|
assert set(os.listdir(dst_dir)) == {"top-level-file.txt", "subdir"}
|
|
assert set(os.listdir(dst_dir.joinpath("subdir"))) == {"subdir-file.txt"}
|
|
assert dst_dir.joinpath("subdir/subdir-file.txt").read_text() == "testing 123"
|
|
assert dst_dir.joinpath("top-level-file.txt").read_text() == "hi"
|
|
|
|
|
|
def test_get_total_size_basic(tmp_path):
|
|
subdir = tmp_path.joinpath("subdir")
|
|
subdir.mkdir()
|
|
|
|
def generate_file(path, size_in_bytes):
|
|
with path.open("wb") as fp:
|
|
fp.write(b"\0" * size_in_bytes)
|
|
|
|
file_size_map = {"file1.txt": 11, "file2.txt": 23}
|
|
for name, size in file_size_map.items():
|
|
generate_file(tmp_path.joinpath(name), size)
|
|
generate_file(subdir.joinpath("file3.txt"), 22)
|
|
assert get_total_file_size(tmp_path) == 56
|
|
assert get_total_file_size(subdir) == 22
|
|
|
|
path_not_exists = tmp_path.joinpath("does_not_exist")
|
|
assert get_total_file_size(path_not_exists) is None
|
|
|
|
path_file = tmp_path.joinpath("file1.txt")
|
|
assert get_total_file_size(path_file) is None
|
|
|
|
|
|
def test_check_tarfile_security(tmp_path):
|
|
def create_tar_with_escaped_path(tar_path: str, escaped_path: str, content: bytes) -> None:
|
|
"""Create tar with path traversal entry."""
|
|
with tarfile.open(tar_path, "w:gz") as tar:
|
|
# Add traversal file
|
|
data = io.BytesIO(content)
|
|
info = tarfile.TarInfo(name=escaped_path)
|
|
info.size = len(content)
|
|
tar.addfile(info, data)
|
|
|
|
tar1_path = str(tmp_path.joinpath("file1.tar"))
|
|
create_tar_with_escaped_path(tar1_path, "../pwned2.txt", b"ABX")
|
|
with pytest.raises(
|
|
MlflowException, match="Escaped path destination in the archive file is not allowed"
|
|
):
|
|
check_tarfile_security(tar1_path)
|
|
|
|
def create_tar_with_symlink(
|
|
tar_path: str, link_name: str, link_target: str, file_via_link: str, content: bytes
|
|
) -> None:
|
|
"""Create tar with symlink that points outside, then file through symlink."""
|
|
with tarfile.open(tar_path, "w:gz") as tar:
|
|
# First: create a symlink pointing to parent directory
|
|
link_info = tarfile.TarInfo(name=link_name)
|
|
link_info.type = tarfile.SYMTYPE
|
|
link_info.linkname = link_target
|
|
tar.addfile(link_info)
|
|
# Second: create a file that goes through the symlink
|
|
data = io.BytesIO(content)
|
|
file_info = tarfile.TarInfo(name=file_via_link)
|
|
file_info.size = len(content)
|
|
tar.addfile(file_info, data)
|
|
|
|
tar2_path = str(tmp_path.joinpath("file2.tar"))
|
|
create_tar_with_symlink(
|
|
tar2_path,
|
|
link_name="escape",
|
|
link_target="..",
|
|
file_via_link="escape/pwned.txt",
|
|
content=b"XYZ",
|
|
)
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match="Destination path in the archive file can not go through a symlink",
|
|
):
|
|
check_tarfile_security(tar2_path)
|
|
|
|
def create_tar_with_abs_path(tar_path: str, abs_path: str, content: bytes) -> None:
|
|
"""Create tar with path traversal entry."""
|
|
with tarfile.open(tar_path, "w:gz") as tar:
|
|
# Add traversal file
|
|
data = io.BytesIO(content)
|
|
info = tarfile.TarInfo(name=abs_path)
|
|
info.size = len(content)
|
|
tar.addfile(info, data)
|
|
|
|
tar3_path = str(tmp_path.joinpath("file3.tar"))
|
|
create_tar_with_abs_path(tar3_path, "/tmp/pwned2.txt", b"ABX")
|
|
with pytest.raises(
|
|
MlflowException, match="Absolute path destination in the archive file is not allowed"
|
|
):
|
|
check_tarfile_security(tar3_path)
|
|
|
|
# Symlink with safe target but file going through it
|
|
tar2b_path = str(tmp_path.joinpath("file2b.tar"))
|
|
create_tar_with_symlink(
|
|
tar2b_path,
|
|
link_name="link_dir",
|
|
link_target="subdir",
|
|
file_via_link="link_dir/pwned.txt",
|
|
content=b"XYZ",
|
|
)
|
|
with pytest.raises(
|
|
MlflowException, match="Destination path in the archive file can not go through a symlink"
|
|
):
|
|
check_tarfile_security(tar2b_path)
|
|
|
|
# Backslash-based path traversal in tar (Windows tar slip / path traversal)
|
|
tar4_path = str(tmp_path.joinpath("file4.tar"))
|
|
create_tar_with_escaped_path(tar4_path, "..\\..\\pwned.txt", b"ABX")
|
|
with pytest.raises(
|
|
MlflowException, match="Escaped path destination in the archive file is not allowed"
|
|
):
|
|
check_tarfile_security(tar4_path)
|
|
|
|
def create_tar_with_symlink_only(tar_path: Path, link_name: str, link_target: str) -> None:
|
|
with tarfile.open(tar_path, "w:gz") as tar:
|
|
link_info = tarfile.TarInfo(name=link_name)
|
|
link_info.type = tarfile.SYMTYPE
|
|
link_info.linkname = link_target
|
|
tar.addfile(link_info)
|
|
|
|
# Symlinks with absolute/escaping targets are allowed (virtualenvs use them).
|
|
# Security is enforced by _safe_extractall at extraction time.
|
|
tar5_path = tmp_path / "file5.tar"
|
|
create_tar_with_symlink_only(tar5_path, "python", "/usr/bin/python3")
|
|
check_tarfile_security(tar5_path) # should not raise
|
|
|
|
tar6_path = tmp_path / "file6.tar"
|
|
create_tar_with_symlink_only(tar6_path, "lib", "../../shared/lib")
|
|
check_tarfile_security(tar6_path) # should not raise
|
|
|
|
# Symlink with absolute path as its own name
|
|
tar7_path = tmp_path / "file7.tar"
|
|
create_tar_with_symlink_only(tar7_path, "/tmp/escape", "foo")
|
|
with pytest.raises(
|
|
MlflowException, match="Absolute path destination in the archive file is not allowed"
|
|
):
|
|
check_tarfile_security(tar7_path)
|
|
|
|
# Symlink whose name escapes with ..
|
|
tar8_path = tmp_path / "file8.tar"
|
|
create_tar_with_symlink_only(tar8_path, "../escape", "foo")
|
|
with pytest.raises(
|
|
MlflowException, match="Escaped path destination in the archive file is not allowed"
|
|
):
|
|
check_tarfile_security(tar8_path)
|
|
|
|
# Hard link with escaping target
|
|
def create_tar_with_hardlink(tar_path: Path, name: str, linkname: str) -> None:
|
|
with tarfile.open(tar_path, "w:gz") as tar:
|
|
info = tarfile.TarInfo(name=name)
|
|
info.type = tarfile.LNKTYPE
|
|
info.linkname = linkname
|
|
tar.addfile(info)
|
|
|
|
tar9_path = tmp_path / "file9.tar"
|
|
create_tar_with_hardlink(tar9_path, "legit.txt", "../../etc/passwd")
|
|
with pytest.raises(
|
|
MlflowException, match="Escaped path destination in the archive file is not allowed"
|
|
):
|
|
check_tarfile_security(tar9_path)
|
|
|
|
# Hard link with absolute target
|
|
tar10_path = tmp_path / "file10.tar"
|
|
create_tar_with_hardlink(tar10_path, "legit.txt", "/etc/passwd")
|
|
with pytest.raises(
|
|
MlflowException, match="Absolute path destination in the archive file is not allowed"
|
|
):
|
|
check_tarfile_security(tar10_path)
|
|
|
|
# Windows drive-letter absolute path
|
|
tar11_path = tmp_path / "file11.tar"
|
|
create_tar_with_escaped_path(tar11_path, "C:/Windows/System32/evil.dll", b"ABX")
|
|
with pytest.raises(
|
|
MlflowException, match="Absolute path destination in the archive file is not allowed"
|
|
):
|
|
check_tarfile_security(tar11_path)
|
|
|
|
|
|
def test_extract_archive_to_dir_blocks_traversal(tmp_path):
|
|
# Test that check_tarfile_security blocks path traversal
|
|
mal_tar = tmp_path / "malicious.tar.gz"
|
|
with tarfile.open(mal_tar, "w:gz") as tar:
|
|
info = tarfile.TarInfo("../../escape.txt")
|
|
data = b"owned via tar traversal"
|
|
info.size = len(data)
|
|
tar.addfile(info, fileobj=io.BytesIO(data))
|
|
|
|
dest = tmp_path / "extracted"
|
|
escape_target = tmp_path.parent.parent / "escape.txt"
|
|
with pytest.raises(MlflowException, match="Escaped path destination in the archive file"):
|
|
extract_archive_to_dir(mal_tar, dest)
|
|
assert not escape_target.exists()
|
|
|
|
|
|
def test_safe_extractall_blocks_symlink_escape(tmp_path):
|
|
"""Test that _safe_extractall blocks extraction when a filesystem symlink
|
|
inside dest_dir would cause a member to resolve outside dest_dir.
|
|
"""
|
|
from mlflow.pyfunc.dbconnect_artifact_cache import _safe_extractall
|
|
|
|
dest = tmp_path / "extracted"
|
|
dest.mkdir()
|
|
# Create a symlink inside dest_dir pointing outside
|
|
escape_link = dest / "escape_link"
|
|
escape_link.symlink_to(tmp_path.parent)
|
|
|
|
# Create a tar with a file that goes through the symlink
|
|
mal_tar = tmp_path / "symlink_escape.tar.gz"
|
|
with tarfile.open(mal_tar, "w:gz") as tar:
|
|
info = tarfile.TarInfo("escape_link/pwned.txt")
|
|
data = b"escaped via filesystem symlink"
|
|
info.size = len(data)
|
|
tar.addfile(info, fileobj=io.BytesIO(data))
|
|
|
|
with tarfile.open(mal_tar, "r") as tar:
|
|
with pytest.raises(MlflowException, match="would be extracted outside"):
|
|
_safe_extractall(tar, dest)
|
|
assert not (tmp_path.parent / "pwned.txt").exists()
|