Skip to content

Commit

Permalink
Apply os.path.basename before . or .. check in `FTPArtifactRepo…
Browse files Browse the repository at this point in the history
…sitory` (#10657)

Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>
  • Loading branch information
harupy committed Dec 12, 2023
1 parent 2121149 commit b9ab9ed
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 1 deletion.
2 changes: 1 addition & 1 deletion mlflow/store/artifact/ftp_artifact_repo.py
Original file line number Diff line number Diff line change
Expand Up @@ -106,10 +106,10 @@ def list_artifacts(self, path=None):
if not self._is_dir(ftp, list_dir):
return []
artifact_files = ftp.nlst(list_dir)
artifact_files = list(filter(lambda x: x != "." and x != "..", artifact_files))
# Make sure artifact_files is a list of file names because ftp.nlst
# may return absolute paths.
artifact_files = [os.path.basename(f) for f in artifact_files]
artifact_files = list(filter(lambda x: x != "." and x != "..", artifact_files))
infos = []
for file_name in artifact_files:
file_path = file_name if path is None else posixpath.join(path, file_name)
Expand Down
12 changes: 12 additions & 0 deletions tests/store/artifact/test_ftp_artifact_repo.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,18 @@ def test_list_artifacts(ftp_mock):
assert artifacts[1].file_size is None


def test_list_artifacts_malicious_path(ftp_mock):
artifact_root_path = "/experiment_id/run_id/"
repo = FTPArtifactRepository("ftp://test_ftp" + artifact_root_path)
repo.get_ftp_client = MagicMock()
call_mock = MagicMock(return_value=ftp_mock)
repo.get_ftp_client.return_value = MagicMock(__enter__=call_mock)
ftp_mock.nlst = MagicMock(return_value=[".", "/.", "/..", "//.."])

artifacts = repo.list_artifacts(path=None)
assert artifacts == []


def test_list_artifacts_when_ftp_nlst_returns_absolute_paths(ftp_mock):
artifact_root_path = "/experiment_id/run_id/"
repo = FTPArtifactRepository("ftp://test_ftp" + artifact_root_path)
Expand Down

0 comments on commit b9ab9ed

Please sign in to comment.