import os import pathlib import posixpath import pytest from mlflow.exceptions import MlflowException from mlflow.store.db.db_types import DATABASE_ENGINES from mlflow.utils.os import is_windows from mlflow.utils.uri import ( add_databricks_profile_info_to_artifact_uri, append_to_uri_path, append_to_uri_query_params, dbfs_hdfs_uri_to_fuse_path, extract_and_normalize_path, extract_db_type_from_uri, get_databricks_profile_uri_from_artifact_uri, get_db_info_from_uri, get_uri_scheme, is_databricks_acled_artifacts_uri, is_databricks_uri, is_fuse_or_uc_volumes_uri, is_http_uri, is_local_uri, is_valid_dbfs_uri, remove_databricks_profile_info_from_artifact_uri, resolve_uri_if_local, strip_scheme, validate_path_is_safe, validate_path_within_directory, ) def test_extract_db_type_from_uri(): uri = "{}://username:password@host:port/database" for legit_db in DATABASE_ENGINES: assert legit_db == extract_db_type_from_uri(uri.format(legit_db)) assert legit_db == get_uri_scheme(uri.format(legit_db)) with_driver = legit_db + "+driver-string" assert legit_db == extract_db_type_from_uri(uri.format(with_driver)) assert legit_db == get_uri_scheme(uri.format(with_driver)) for unsupported_db in ["a", "aa", "sql"]: with pytest.raises(MlflowException, match="Invalid database engine"): extract_db_type_from_uri(unsupported_db) @pytest.mark.parametrize( ("server_uri", "result"), [ ("databricks://aAbB", ("aAbB", None)), ("databricks://aAbB/", ("aAbB", None)), ("databricks://aAbB/path", ("aAbB", None)), ("databricks://profile:prefix", ("profile", "prefix")), ("databricks://profile:prefix/extra", ("profile", "prefix")), ("nondatabricks://profile:prefix", (None, None)), ("databricks://profile", ("profile", None)), ("databricks://profile/", ("profile", None)), ("databricks-uc://profile:prefix", ("profile", "prefix")), ("databricks-uc://profile:prefix/extra", ("profile", "prefix")), ("databricks-uc://profile", ("profile", None)), ("databricks-uc://profile/", ("profile", None)), ], ) def test_get_db_info_from_uri(server_uri, result): assert get_db_info_from_uri(server_uri) == result @pytest.mark.parametrize( "server_uri", ["databricks:/profile:prefix", "databricks:/", "databricks://"], ) def test_get_db_info_from_uri_errors_no_netloc(server_uri): with pytest.raises(MlflowException, match="URI is formatted incorrectly"): get_db_info_from_uri(server_uri) @pytest.mark.parametrize( "server_uri", [ "databricks://profile:prefix:extra", "databricks://profile:prefix:extra ", "databricks://profile:prefix extra", "databricks://profile:prefix ", "databricks://profile ", "databricks://profile:", "databricks://profile: ", ], ) def test_get_db_info_from_uri_errors_invalid_profile(server_uri): with pytest.raises(MlflowException, match="Unsupported Databricks profile"): get_db_info_from_uri(server_uri) def test_is_local_uri(): assert is_local_uri("mlruns") assert is_local_uri("./mlruns") assert is_local_uri("file:///foo/mlruns") assert is_local_uri("file:foo/mlruns") assert is_local_uri("file://./mlruns") assert is_local_uri("file://localhost/mlruns") assert is_local_uri("file://localhost:5000/mlruns") assert is_local_uri("file://127.0.0.1/mlruns") assert is_local_uri("file://127.0.0.1:5000/mlruns") assert is_local_uri("//proc/self/root") assert is_local_uri("/proc/self/root") assert not is_local_uri("https://whatever") assert not is_local_uri("http://whatever") assert not is_local_uri("databricks") assert not is_local_uri("databricks:whatever") assert not is_local_uri("databricks://whatever") with pytest.raises(MlflowException, match="is not a valid remote uri."): is_local_uri("file://myhostname/path/to/file") @pytest.mark.skipif(not is_windows(), reason="Windows-only test") def test_is_local_uri_windows(): assert is_local_uri("C:\\foo\\mlruns") assert is_local_uri("C:/foo/mlruns") assert is_local_uri("file:///C:\\foo\\mlruns") assert not is_local_uri("\\\\server\\aa\\bb") def test_is_databricks_uri(): assert is_databricks_uri("databricks") assert is_databricks_uri("databricks:whatever") assert is_databricks_uri("databricks://whatever") assert not is_databricks_uri("mlruns") assert not is_databricks_uri("http://whatever") def test_is_http_uri(): assert is_http_uri("http://whatever") assert is_http_uri("https://whatever") assert not is_http_uri("file://whatever") assert not is_http_uri("databricks://whatever") assert not is_http_uri("mlruns") def validate_append_to_uri_path_test_cases(cases): for input_uri, input_path, expected_output_uri in cases: assert append_to_uri_path(input_uri, input_path) == expected_output_uri assert append_to_uri_path(input_uri, *posixpath.split(input_path)) == expected_output_uri def test_append_to_uri_path_joins_uri_paths_and_posixpaths_correctly(): validate_append_to_uri_path_test_cases([ ("", "path", "path"), ("", "/path", "/path"), ("path", "", "path/"), ("path", "subpath", "path/subpath"), ("path/", "subpath", "path/subpath"), ("path/", "/subpath", "path/subpath"), ("path", "/subpath", "path/subpath"), ("/path", "/subpath", "/path/subpath"), ("//path", "/subpath", "//path/subpath"), ("///path", "/subpath", "///path/subpath"), ("/path", "/subpath/subdir", "/path/subpath/subdir"), ("file:path", "", "file:path/"), ("file:path/", "", "file:path/"), ("file:path", "subpath", "file:path/subpath"), ("file:path", "/subpath", "file:path/subpath"), ("file:/", "", "file:///"), ("file:/path", "/subpath", "file:///path/subpath"), ("file:///", "", "file:///"), ("file:///", "subpath", "file:///subpath"), ("file:///path", "/subpath", "file:///path/subpath"), ("file:///path/", "subpath", "file:///path/subpath"), ("file:///path", "subpath", "file:///path/subpath"), ("s3://", "", "s3:"), ("s3://", "subpath", "s3:subpath"), ("s3://", "/subpath", "s3:/subpath"), ("s3://host", "subpath", "s3://host/subpath"), ("s3://host", "/subpath", "s3://host/subpath"), ("s3://host/", "subpath", "s3://host/subpath"), ("s3://host/", "/subpath", "s3://host/subpath"), ("s3://host", "subpath/subdir", "s3://host/subpath/subdir"), ]) def test_append_to_uri_path_handles_special_uri_characters_in_posixpaths(): """ Certain characters are treated specially when parsing and interpreting URIs. However, in the case where a URI input for `append_to_uri_path` is simply a POSIX path, these characters should not receive special treatment. This test case verifies that `append_to_uri_path` properly joins POSIX paths containing these characters. """ def create_char_case(special_char): def char_case(*case_args): return tuple(item.format(c=special_char) for item in case_args) return char_case for special_char in [ ".", "-", "+", ":", "?", "@", "&", "$", "%", "/", "[", "]", "(", ")", "*", "'", ",", ]: char_case = create_char_case(special_char) validate_append_to_uri_path_test_cases([ char_case("", "{c}subpath", "{c}subpath"), char_case("", "/{c}subpath", "/{c}subpath"), char_case("dirwith{c}{c}chars", "", "dirwith{c}{c}chars/"), char_case("dirwith{c}{c}chars", "subpath", "dirwith{c}{c}chars/subpath"), char_case("{c}{c}charsdir", "", "{c}{c}charsdir/"), char_case("/{c}{c}charsdir", "", "/{c}{c}charsdir/"), char_case("/{c}{c}charsdir", "subpath", "/{c}{c}charsdir/subpath"), char_case("/{c}{c}charsdir", "subpath", "/{c}{c}charsdir/subpath"), ]) validate_append_to_uri_path_test_cases([ ("#?charsdir:", ":?subpath#", "#?charsdir:/:?subpath#"), ("/#--+charsdir.//:", "/../:?subpath#", "/#--+charsdir.//:/../:?subpath#"), ("$@''(,", ")]*%", "$@''(,/)]*%"), ]) @pytest.mark.parametrize( "uri", [ # query string contains '..' (and its encoded form) are considered invalid "https://example.com?..", "https://example.com?/path/../path/../path", "https://example.com?key=value&../../path", "https://example.com?key=value&%2E%2E%2Fpath", "https://example.com?key=value&%252E%252E%252Fpath", ], ) def test_append_to_uri_throws_for_malicious_query_string_in_uri(uri): with pytest.raises(MlflowException, match=r"Invalid query string"): append_to_uri_path(uri) @pytest.mark.parametrize( ("uri", "existing_query_params", "query_params", "expected"), [ ("https://example.com", "", [("key", "value")], "https://example.com?key=value"), ( "https://example.com", "existing_key=existing_value", [("new_key", "new_value")], "https://example.com?existing_key=existing_value&new_key=new_value", ), ( "https://example.com", "", [("key1", "value1"), ("key2", "value2"), ("key3", "value3")], "https://example.com?key1=value1&key2=value2&key3=value3", ), ( "https://example.com", "", [("key", "value with spaces"), ("key2", "special#characters")], "https://example.com?key=value+with+spaces&key2=special%23characters", ), ("", "", [("key", "value")], "?key=value"), ("https://example.com", "", [], "https://example.com"), ( "https://example.com", "", [("key1", 123), ("key2", 456)], "https://example.com?key1=123&key2=456", ), ( "https://example.com?existing_key=existing_value", "", [("existing_key", "new_value"), ("existing_key", "new_value_2")], "https://example.com?existing_key=existing_value&existing_key=new_value&existing_key=new_value_2", ), ( "s3://bucket/key", "prev1=foo&prev2=bar", [("param1", "value1"), ("param2", "value2")], "s3://bucket/key?prev1=foo&prev2=bar¶m1=value1¶m2=value2", ), ( "s3://bucket/key?existing_param=existing_value", "", [("new_param", "new_value")], "s3://bucket/key?existing_param=existing_value&new_param=new_value", ), ], ) def test_append_to_uri_query_params_appends_as_expected( uri, existing_query_params, query_params, expected ): if existing_query_params: uri += f"?{existing_query_params}" result = append_to_uri_query_params(uri, *query_params) assert result == expected def test_append_to_uri_path_preserves_uri_schemes_hosts_queries_and_fragments(): validate_append_to_uri_path_test_cases([ ("dbscheme+dbdriver:", "", "dbscheme+dbdriver:"), ("dbscheme+dbdriver:", "subpath", "dbscheme+dbdriver:subpath"), ("dbscheme+dbdriver:path", "subpath", "dbscheme+dbdriver:path/subpath"), ("dbscheme+dbdriver://host/path", "/subpath", "dbscheme+dbdriver://host/path/subpath"), ("dbscheme+dbdriver:///path", "subpath", "dbscheme+dbdriver:/path/subpath"), ("dbscheme+dbdriver:?somequery", "subpath", "dbscheme+dbdriver:subpath?somequery"), ("dbscheme+dbdriver:?somequery", "/subpath", "dbscheme+dbdriver:/subpath?somequery"), ("dbscheme+dbdriver:/?somequery", "subpath", "dbscheme+dbdriver:/subpath?somequery"), ("dbscheme+dbdriver://?somequery", "subpath", "dbscheme+dbdriver:subpath?somequery"), ("dbscheme+dbdriver:///?somequery", "/subpath", "dbscheme+dbdriver:/subpath?somequery"), ("dbscheme+dbdriver:#somefrag", "subpath", "dbscheme+dbdriver:subpath#somefrag"), ("dbscheme+dbdriver:#somefrag", "/subpath", "dbscheme+dbdriver:/subpath#somefrag"), ("dbscheme+dbdriver:/#somefrag", "subpath", "dbscheme+dbdriver:/subpath#somefrag"), ("dbscheme+dbdriver://#somefrag", "subpath", "dbscheme+dbdriver:subpath#somefrag"), ("dbscheme+dbdriver:///#somefrag", "/subpath", "dbscheme+dbdriver:/subpath#somefrag"), ( "dbscheme+dbdriver://root:password?creds=creds", "subpath", "dbscheme+dbdriver://root:password/subpath?creds=creds", ), ( "dbscheme+dbdriver://root:password/path/?creds=creds", "/subpath/anotherpath", "dbscheme+dbdriver://root:password/path/subpath/anotherpath?creds=creds", ), ( "dbscheme+dbdriver://root:password///path/?creds=creds", "subpath/anotherpath", "dbscheme+dbdriver://root:password///path/subpath/anotherpath?creds=creds", ), ( "dbscheme+dbdriver://root:password///path/?creds=creds", "/subpath", "dbscheme+dbdriver://root:password///path/subpath?creds=creds", ), ( "dbscheme+dbdriver://root:password#myfragment", "/subpath", "dbscheme+dbdriver://root:password/subpath#myfragment", ), ( "dbscheme+dbdriver://root:password//path/#fragmentwith$pecial@", "subpath/anotherpath", "dbscheme+dbdriver://root:password//path/subpath/anotherpath#fragmentwith$pecial@", ), ( "dbscheme+dbdriver://root:password@host?creds=creds#fragmentwith$pecial@", "subpath", "dbscheme+dbdriver://root:password@host/subpath?creds=creds#fragmentwith$pecial@", ), ( "dbscheme+dbdriver://root:password@host.com/path?creds=creds#*frag@*", "subpath/dir", "dbscheme+dbdriver://root:password@host.com/path/subpath/dir?creds=creds#*frag@*", ), ( "dbscheme-dbdriver://root:password@host.com/path?creds=creds#*frag@*", "subpath/dir", "dbscheme-dbdriver://root:password@host.com/path/subpath/dir?creds=creds#*frag@*", ), ( "dbscheme+dbdriver://root:password@host.com/path?creds=creds,param=value#*frag@*", "subpath/dir", "dbscheme+dbdriver://root:password@host.com/path/subpath/dir?" "creds=creds,param=value#*frag@*", ), ]) def test_extract_and_normalize_path(): base_uri = "databricks/mlflow-tracking/EXP_ID/RUN_ID/artifacts" assert ( extract_and_normalize_path("dbfs:databricks/mlflow-tracking/EXP_ID/RUN_ID/artifacts") == base_uri ) assert ( extract_and_normalize_path("dbfs:/databricks/mlflow-tracking/EXP_ID/RUN_ID/artifacts") == base_uri ) assert ( extract_and_normalize_path("dbfs:///databricks/mlflow-tracking/EXP_ID/RUN_ID/artifacts") == base_uri ) assert ( extract_and_normalize_path( "dbfs:/databricks///mlflow-tracking///EXP_ID///RUN_ID///artifacts/" ) == base_uri ) assert ( extract_and_normalize_path( "dbfs:///databricks///mlflow-tracking//EXP_ID//RUN_ID///artifacts//" ) == base_uri ) assert ( extract_and_normalize_path( "dbfs:databricks///mlflow-tracking//EXP_ID//RUN_ID///artifacts//" ) == base_uri ) def test_is_databricks_acled_artifacts_uri(): assert is_databricks_acled_artifacts_uri( "dbfs:databricks/mlflow-tracking/EXP_ID/RUN_ID/artifacts" ) assert is_databricks_acled_artifacts_uri( "dbfs:/databricks/mlflow-tracking/EXP_ID/RUN_ID/artifacts" ) assert is_databricks_acled_artifacts_uri( "dbfs:///databricks/mlflow-tracking/EXP_ID/RUN_ID/artifacts" ) assert is_databricks_acled_artifacts_uri( "dbfs:/databricks///mlflow-tracking///EXP_ID///RUN_ID///artifacts/" ) assert is_databricks_acled_artifacts_uri( "dbfs:///databricks///mlflow-tracking//EXP_ID//RUN_ID///artifacts//" ) assert is_databricks_acled_artifacts_uri( "dbfs:databricks///mlflow-tracking//EXP_ID//RUN_ID///artifacts//" ) assert not is_databricks_acled_artifacts_uri( "dbfs:/databricks/mlflow//EXP_ID//RUN_ID///artifacts//" ) def _get_databricks_profile_uri_test_cases(): # Each test case is (uri, result, result_scheme) test_case_groups = [ [ # URIs with no databricks profile info -> return None ("ftp://user:pass@realhost:port/path/to/nowhere", None, result_scheme), ("dbfs:/path/to/nowhere", None, result_scheme), ("dbfs://nondatabricks/path/to/nowhere", None, result_scheme), ("dbfs://incorrect:netloc:format/path/to/nowhere", None, result_scheme), # URIs with legit databricks profile info (f"dbfs://{result_scheme}", result_scheme, result_scheme), (f"dbfs://{result_scheme}/", result_scheme, result_scheme), (f"dbfs://{result_scheme}/path/to/nowhere", result_scheme, result_scheme), (f"dbfs://{result_scheme}:port/path/to/nowhere", result_scheme, result_scheme), (f"dbfs://@{result_scheme}/path/to/nowhere", result_scheme, result_scheme), (f"dbfs://@{result_scheme}:port/path/to/nowhere", result_scheme, result_scheme), ( f"dbfs://profile@{result_scheme}/path/to/nowhere", f"{result_scheme}://profile", result_scheme, ), ( f"dbfs://profile@{result_scheme}:port/path/to/nowhere", f"{result_scheme}://profile", result_scheme, ), ( f"dbfs://scope:key_prefix@{result_scheme}/path/abc", f"{result_scheme}://scope:key_prefix", result_scheme, ), ( f"dbfs://scope:key_prefix@{result_scheme}:port/path/abc", f"{result_scheme}://scope:key_prefix", result_scheme, ), # Doesn't care about the scheme of the artifact URI ( f"runs://scope:key_prefix@{result_scheme}/path/abc", f"{result_scheme}://scope:key_prefix", result_scheme, ), ( f"models://scope:key_prefix@{result_scheme}/path/abc", f"{result_scheme}://scope:key_prefix", result_scheme, ), ( f"s3://scope:key_prefix@{result_scheme}/path/abc", f"{result_scheme}://scope:key_prefix", result_scheme, ), ] for result_scheme in ["databricks", "databricks-uc"] ] return [test_case for test_case_group in test_case_groups for test_case in test_case_group] @pytest.mark.parametrize( ("uri", "result", "result_scheme"), _get_databricks_profile_uri_test_cases() ) def test_get_databricks_profile_uri_from_artifact_uri(uri, result, result_scheme): assert get_databricks_profile_uri_from_artifact_uri(uri, result_scheme=result_scheme) == result @pytest.mark.parametrize( "uri", [ # Treats secret key prefixes with ":" to be invalid "dbfs://incorrect:netloc:format@databricks/path/a", "dbfs://scope::key_prefix@databricks/path/abc", "dbfs://scope:key_prefix:@databricks/path/abc", ], ) def test_get_databricks_profile_uri_from_artifact_uri_error_cases(uri): with pytest.raises(MlflowException, match="Unsupported Databricks profile"): get_databricks_profile_uri_from_artifact_uri(uri) @pytest.mark.parametrize( ("uri", "result"), [ # URIs with no databricks profile info should stay the same ( "ftp://user:pass@realhost:port/path/nowhere", "ftp://user:pass@realhost:port/path/nowhere", ), ("dbfs:/path/to/nowhere", "dbfs:/path/to/nowhere"), ("dbfs://nondatabricks/path/to/nowhere", "dbfs://nondatabricks/path/to/nowhere"), ("dbfs://incorrect:netloc:format/path/", "dbfs://incorrect:netloc:format/path/"), # URIs with legit databricks profile info ("dbfs://databricks", "dbfs:"), ("dbfs://databricks/", "dbfs:/"), ("dbfs://databricks/path/to/nowhere", "dbfs:/path/to/nowhere"), ("dbfs://databricks:port/path/to/nowhere", "dbfs:/path/to/nowhere"), ("dbfs://@databricks/path/to/nowhere", "dbfs:/path/to/nowhere"), ("dbfs://@databricks:port/path/to/nowhere", "dbfs:/path/to/nowhere"), ("dbfs://profile@databricks/path/to/nowhere", "dbfs:/path/to/nowhere"), ("dbfs://profile@databricks:port/path/to/nowhere", "dbfs:/path/to/nowhere"), ("dbfs://scope:key_prefix@databricks/path/abc", "dbfs:/path/abc"), ("dbfs://scope:key_prefix@databricks:port/path/abc", "dbfs:/path/abc"), # Treats secret key prefixes with ":" to be valid ("dbfs://incorrect:netloc:format@databricks/path/to/nowhere", "dbfs:/path/to/nowhere"), # Doesn't care about the scheme of the artifact URI ("runs://scope:key_prefix@databricks/path/abc", "runs:/path/abc"), ("models://scope:key_prefix@databricks/path/abc", "models:/path/abc"), ("s3://scope:key_prefix@databricks/path/abc", "s3:/path/abc"), ], ) def test_remove_databricks_profile_info_from_artifact_uri(uri, result): assert remove_databricks_profile_info_from_artifact_uri(uri) == result @pytest.mark.parametrize( ("artifact_uri", "profile_uri", "result"), [ # test various profile URIs ("dbfs:/path/a/b", "databricks", "dbfs://databricks/path/a/b"), ("dbfs:/path/a/b/", "databricks", "dbfs://databricks/path/a/b/"), ("dbfs:/path/a/b/", "databricks://Profile", "dbfs://Profile@databricks/path/a/b/"), ("dbfs:/path/a/b/", "databricks://profile/", "dbfs://profile@databricks/path/a/b/"), ("dbfs:/path/a/b/", "databricks://scope:key", "dbfs://scope:key@databricks/path/a/b/"), ( "dbfs:/path/a/b/", "databricks://scope:key/random_stuff", "dbfs://scope:key@databricks/path/a/b/", ), ("dbfs:/path/a/b/", "nondatabricks://profile", "dbfs:/path/a/b/"), # test various artifact schemes ("runs:/path/a/b/", "databricks://Profile", "runs://Profile@databricks/path/a/b/"), ("runs:/path/a/b/", "nondatabricks://profile", "runs:/path/a/b/"), ("models:/path/a/b/", "databricks://profile", "models://profile@databricks/path/a/b/"), ("models:/path/a/b/", "nondatabricks://Profile", "models:/path/a/b/"), ("s3:/path/a/b/", "databricks://Profile", "s3:/path/a/b/"), ("s3:/path/a/b/", "nondatabricks://profile", "s3:/path/a/b/"), ("ftp:/path/a/b/", "databricks://profile", "ftp:/path/a/b/"), ("ftp:/path/a/b/", "nondatabricks://Profile", "ftp:/path/a/b/"), # test artifact URIs already with authority ("ftp://user:pass@host:port/a/b", "databricks://Profile", "ftp://user:pass@host:port/a/b"), ("ftp://user:pass@host:port/a/b", "nothing://Profile", "ftp://user:pass@host:port/a/b"), ("dbfs://databricks", "databricks://OtherProfile", "dbfs://databricks"), ("dbfs://databricks", "nondatabricks://Profile", "dbfs://databricks"), ("dbfs://databricks/path/a/b", "databricks://OtherProfile", "dbfs://databricks/path/a/b"), ("dbfs://databricks/path/a/b", "nondatabricks://Profile", "dbfs://databricks/path/a/b"), ("dbfs://@databricks/path/a/b", "databricks://OtherProfile", "dbfs://@databricks/path/a/b"), ("dbfs://@databricks/path/a/b", "nondatabricks://Profile", "dbfs://@databricks/path/a/b"), ( "dbfs://profile@databricks/pp", "databricks://OtherProfile", "dbfs://profile@databricks/pp", ), ( "dbfs://profile@databricks/path", "databricks://profile", "dbfs://profile@databricks/path", ), ( "dbfs://profile@databricks/path", "nondatabricks://Profile", "dbfs://profile@databricks/path", ), ], ) def test_add_databricks_profile_info_to_artifact_uri(artifact_uri, profile_uri, result): assert add_databricks_profile_info_to_artifact_uri(artifact_uri, profile_uri) == result @pytest.mark.parametrize( ("artifact_uri", "profile_uri"), [ ("dbfs:/path/a/b", "databricks://not:legit:auth"), ("dbfs:/path/a/b/", "databricks://scope::key"), ("dbfs:/path/a/b/", "databricks://scope:key:/"), ("dbfs:/path/a/b/", "databricks://scope:key "), ], ) def test_add_databricks_profile_info_to_artifact_uri_errors(artifact_uri, profile_uri): with pytest.raises(MlflowException, match="Unsupported Databricks profile"): add_databricks_profile_info_to_artifact_uri(artifact_uri, profile_uri) @pytest.mark.parametrize( ("uri", "result"), [ ("dbfs:/path/a/b", True), ("dbfs://databricks/a/b", True), ("dbfs://@databricks/a/b", True), ("dbfs://profile@databricks/a/b", True), ("dbfs://scope:key@databricks/a/b", True), ("dbfs://scope:key:@databricks/a/b", False), ("dbfs://scope::key@databricks/a/b", False), ("dbfs://profile@notdatabricks/a/b", False), ("dbfs://scope:key@notdatabricks/a/b", False), ("dbfs://scope:key/a/b", False), ("dbfs://notdatabricks/a/b", False), ("s3:/path/a/b", False), ("ftp://user:pass@host:port/path/a/b", False), ("ftp://user:pass@databricks/path/a/b", False), ], ) def test_is_valid_dbfs_uri(uri, result): assert is_valid_dbfs_uri(uri) == result @pytest.mark.parametrize( ("uri", "result"), [ ("/tmp/path", "/dbfs/tmp/path"), ("dbfs:/path", "/dbfs/path"), ("dbfs:/path/a/b", "/dbfs/path/a/b"), ("dbfs:/dbfs/123/abc", "/dbfs/dbfs/123/abc"), ], ) def test_dbfs_hdfs_uri_to_fuse_path(uri, result): assert dbfs_hdfs_uri_to_fuse_path(uri) == result @pytest.mark.parametrize( "path", ["some/relative/local/path", "s3:/some/s3/path", "C:/cool/windows/path"], ) def test_dbfs_hdfs_uri_to_fuse_path_raises(path): with pytest.raises(MlflowException, match="did not start with expected DBFS URI prefix"): dbfs_hdfs_uri_to_fuse_path(path) def _assert_resolve_uri_if_local(input_uri, expected_uri): cwd = pathlib.Path.cwd().as_posix() drive = pathlib.Path.cwd().drive if is_windows(): cwd = f"/{cwd}" drive = f"{drive}/" assert resolve_uri_if_local(input_uri) == expected_uri.format(cwd=cwd, drive=drive) @pytest.mark.skipif(is_windows(), reason="This test fails on Windows") @pytest.mark.parametrize( ("input_uri", "expected_uri"), [ ("my/path", "{cwd}/my/path"), ("#my/path?a=b", "{cwd}/#my/path?a=b"), ("file://localhost/my/path", "file://localhost/my/path"), ("file:///my/path", "file:///{drive}my/path"), ("file:my/path", "file://{cwd}/my/path"), ("/home/my/path", "/home/my/path"), ("dbfs://databricks/a/b", "dbfs://databricks/a/b"), ("s3://host/my/path", "s3://host/my/path"), ], ) def test_resolve_uri_if_local(input_uri, expected_uri): _assert_resolve_uri_if_local(input_uri, expected_uri) @pytest.mark.skipif(not is_windows(), reason="This test only passes on Windows") @pytest.mark.parametrize( ("input_uri", "expected_uri"), [ ("my/path", "file://{cwd}/my/path"), ("#my/path?a=b", "file://{cwd}/#my/path?a=b"), ("\\myhostname/my/path", "file:///{drive}myhostname/my/path"), ("file:///my/path", "file:///{drive}my/path"), ("file:my/path", "file://{cwd}/my/path"), ("/home/my/path", "file:///{drive}home/my/path"), ("dbfs://databricks/a/b", "dbfs://databricks/a/b"), ("s3://host/my/path", "s3://host/my/path"), ], ) def test_resolve_uri_if_local_on_windows(input_uri, expected_uri): _assert_resolve_uri_if_local(input_uri, expected_uri) @pytest.mark.parametrize( "uri", [ "/dbfs/my_path", "dbfs:/my_path", "/Volumes/my_path", "/.fuse-mounts/my_path", "//dbfs////my_path", "///Volumes/", "dbfs://my///path", "/volumes/path/to/file", "/volumes/", "DBFS:/my/path", ], ) def test_correctly_detect_fuse_and_uc_uris(uri): assert is_fuse_or_uc_volumes_uri(uri) @pytest.mark.parametrize( "uri", [ "/My_Volumes/my_path", "s3a:/my_path", "Volumes/my_path", "Volume:/my_path", "dbfs/my_path", "/fuse-mounts/my_path", ], ) def test_negative_detection(uri): assert not is_fuse_or_uc_volumes_uri(uri) @pytest.mark.parametrize( "path", [ "path", "path/", "path/to/file", "dog%step%100%timestamp%100", "dog+step+100+timestamp+100", ], ) def test_validate_path_is_safe_good(path): validate_path_is_safe(path) @pytest.mark.skipif(not is_windows(), reason="This test only passes on Windows") @pytest.mark.parametrize( "path", [ # relative path from current directory of C: drive ".../...//", ], ) def test_validate_path_is_safe_windows_good(path): validate_path_is_safe(path) @pytest.mark.skipif(is_windows(), reason="This test does not pass on Windows") @pytest.mark.parametrize( "path", [ "/path", "../path", "../../path", "./../path", "path/../to/file", "path/../../to/file", "file://a#/..//tmp", "file://a%23/..//tmp/", "/etc/passwd", "/etc/passwd%00.jpg", "/etc/passwd%00.html", "/etc/passwd%00.txt", "/etc/passwd%00.php", "/etc/passwd%00.asp", "/file://etc/passwd", # Encoded paths with '..' "%2E%2E%2Fpath", "%2E%2E%2F%2E%2E%2Fpath", # Some URIs are passed to urllib.parse.urlparse after validation, # which strips out some whitespace characters. If they are further # decoded, this could result in a path that is not safe. # In this example, %2%0952e -> %2\t52e -> %252e -> %2e -> . "%2%0952e%2%0952e/%2%0A52e%2%0A52e/path", ], ) def test_validate_path_is_safe_bad(path): with pytest.raises(MlflowException, match="Invalid path"): validate_path_is_safe(path) @pytest.mark.skipif(not is_windows(), reason="This test only passes on Windows") @pytest.mark.parametrize( "path", [ r"../path", r"../../path", r"./../path", r"path/../to/file", r"path/../../to/file", r"..\path", r"..\..\path", r".\..\path", r"path\..\to\file", r"path\..\..\to\file", # Drive-relative paths r"C:path", r"C:path/", r"C:path/to/file", r"C:../path/to/file", r"C:\path", r"C:/path", r"C:\path\to\file", r"C:\path/to/file", r"C:\path\..\to\file", r"C:/path/../to/file", # UNC(Universal Naming Convention) paths r"\\path\to\file", r"\\path/to/file", r"\\.\\C:\path\to\file", r"\\?\C:\path\to\file", r"\\?\UNC/path/to/file", # Other potential attackable paths r"/etc/password", r"/path", r"/etc/passwd%00.jpg", r"/etc/passwd%00.html", r"/etc/passwd%00.txt", r"/etc/passwd%00.php", r"/etc/passwd%00.asp", r"/Windows/no/such/path", r"/file://etc/passwd", r"/file:c:/passwd", r"/file://d:/windows/win.ini", r"/file://./windows/win.ini", r"file://c:/boot.ini", r"file://C:path", r"file://C:path/", r"file://C:path/to/file", r"file:///C:/Windows/System32/", r"file:///etc/passwd", r"file:///d:/windows/repair/sam", r"file:///proc/version", r"file:///inetpub/wwwroot/global.asa", r"/file://../windows/win.ini", r"../etc/passwd", r"..\Windows\System32\\", r"C:\Windows\System32\\", r"/etc/passwd", r"::Windows\System32", r"..\..\..\..\Windows\System32\\", r"../Windows/System32", r"....\\", r"\\?\C:\Windows\System32\\", r"\\.\C:\Windows\System32\\", r"\\UNC\Server\Share\\", r"\\Server\Share\folder\\", r"\\127.0.0.1\c$\Windows\\", r"\\localhost\c$\Windows\\", r"\\smbserver\share\path\\", r"..\\?\C:\Windows\System32\\", r"C:/Windows/../Windows/System32/", r"C:\Windows\..\Windows\System32\\", r"../../../../../../../../../../../../Windows/System32", r"../../../../../../../../../../../../etc/passwd", r"../../../../../../../../../../../../var/www/html/index.html", r"../../../../../../../../../../../../usr/local/etc/openvpn/server.conf", r"../../../../../../../../../../../../Program Files (x86)", r"/../../../../../../../../../../../../Windows/System32", r"/Windows\../etc/passwd", r"/Windows\..\Windows\System32\\", r"/Windows\..\Windows\System32\cmd.exe", r"/Windows\..\Windows\System32\msconfig.exe", r"/Windows\..\Windows\System32\regedit.exe", r"/Windows\..\Windows\System32\taskmgr.exe", r"/Windows\..\Windows\System32\control.exe", r"/Windows\..\Windows\System32\services.msc", r"/Windows\..\Windows\System32\diskmgmt.msc", r"/Windows\..\Windows\System32\eventvwr.msc", r"/Windows/System32/drivers/etc/hosts", ], ) def test_validate_path_is_safe_windows_bad(path): with pytest.raises(MlflowException, match="Invalid path"): validate_path_is_safe(path) @pytest.mark.parametrize( ("uri", "expected"), [ ("file:///path", "/path"), ("file://host/path", "//host/path"), ("file://host", "//host"), ], ) def test_strip_scheme(uri: str, expected: str): assert strip_scheme(uri) == expected def test_validate_path_within_directory_allows_valid_paths(tmp_path): base_dir = tmp_path / "artifacts" base_dir.mkdir() constructed_path = base_dir / "subdir" / "file.txt" result = validate_path_within_directory(str(base_dir), str(constructed_path)) assert result == str(constructed_path) def test_validate_path_within_directory_blocks_symlink_escape(tmp_path): base_dir = tmp_path / "artifacts" base_dir.mkdir() external_dir = tmp_path / "external" external_dir.mkdir() external_file = external_dir / "secret.txt" external_file.write_text("SECRET") symlink_path = base_dir / "leak" os.symlink(str(external_dir), str(symlink_path)) constructed_path = symlink_path / "secret.txt" with pytest.raises(MlflowException, match="resolved path is outside the artifact directory"): validate_path_within_directory(str(base_dir), str(constructed_path)) def test_validate_path_within_directory_blocks_parent_symlink(tmp_path): base_dir = tmp_path / "artifacts" base_dir.mkdir() symlink_path = base_dir / "parent" os.symlink(str(tmp_path), str(symlink_path)) constructed_path = symlink_path / "artifacts" / ".." / "external" with pytest.raises(MlflowException, match="resolved path is outside the artifact directory"): validate_path_within_directory(str(base_dir), str(constructed_path)) def test_validate_path_within_directory_allows_internal_symlink(tmp_path): base_dir = tmp_path / "artifacts" base_dir.mkdir() real_file = base_dir / "real_file.txt" real_file.write_text("CONTENT") symlink_path = base_dir / "link" os.symlink(str(real_file), str(symlink_path)) result = validate_path_within_directory(str(base_dir), str(symlink_path)) assert result == str(symlink_path) def test_validate_path_within_directory_allows_base_dir_itself(tmp_path): base_dir = tmp_path / "artifacts" base_dir.mkdir() result = validate_path_within_directory(str(base_dir), str(base_dir)) assert result == str(base_dir) def test_validate_path_within_directory_allows_subdirectory_symlink(tmp_path): base_dir = tmp_path / "artifacts" base_dir.mkdir() subdir = base_dir / "subdir" subdir.mkdir() file_in_subdir = subdir / "file.txt" file_in_subdir.write_text("CONTENT") symlink_to_subdir = base_dir / "link_to_subdir" os.symlink(str(subdir), str(symlink_to_subdir)) constructed_path = symlink_to_subdir / "file.txt" result = validate_path_within_directory(str(base_dir), str(constructed_path)) assert result == str(constructed_path)