Skip to content

Commit

Permalink
Use validated path after path validation (#10666)
Browse files Browse the repository at this point in the history
Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>
  • Loading branch information
harupy committed Dec 13, 2023
1 parent a98a341 commit 5044878
Show file tree
Hide file tree
Showing 3 changed files with 15 additions and 17 deletions.
26 changes: 11 additions & 15 deletions mlflow/server/handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -548,7 +548,7 @@ def get_artifact_handler():
request_dict = parser.parse(query_string, normalized=True)
run_id = request_dict.get("run_id") or request_dict.get("run_uuid")
path = request_dict["path"]
validate_path_is_safe(path)
path = validate_path_is_safe(path)
run = _get_tracking_store().get_run(run_id)

if _is_servable_proxied_run_artifact_root(run.info.artifact_uri):
Expand Down Expand Up @@ -941,7 +941,7 @@ def _list_artifacts():
response_message = ListArtifacts.Response()
if request_message.HasField("path"):
path = request_message.path
validate_path_is_safe(path)
path = validate_path_is_safe(path)
else:
path = None
run_id = request_message.run_id or request_message.run_uuid
Expand Down Expand Up @@ -1236,7 +1236,7 @@ def upload_artifact_handler():
message="Request must specify path.",
error_code=INVALID_PARAMETER_VALUE,
)
validate_path_is_safe(path)
path = validate_path_is_safe(path)

if request.content_length and request.content_length > 10 * 1024 * 1024:
raise MlflowException(
Expand Down Expand Up @@ -1660,7 +1660,7 @@ def get_model_version_artifact_handler():
name = request_dict.get("name")
version = request_dict.get("version")
path = request_dict["path"]
validate_path_is_safe(path)
path = validate_path_is_safe(path)
artifact_uri = _get_model_registry_store().get_model_version_download_uri(name, version)
if _is_servable_proxied_run_artifact_root(artifact_uri):
artifact_repo = _get_artifact_repo_mlflow_artifacts()
Expand Down Expand Up @@ -1886,7 +1886,7 @@ def _download_artifact(artifact_path):
A request handler for `GET /mlflow-artifacts/artifacts/<artifact_path>` to download an artifact
from `artifact_path` (a relative path from the root artifact directory).
"""
validate_path_is_safe(artifact_path)
artifact_path = validate_path_is_safe(artifact_path)
tmp_dir = tempfile.TemporaryDirectory()
artifact_repo = _get_artifact_repo_mlflow_artifacts()
dst = artifact_repo.download_artifacts(artifact_path, tmp_dir.name)
Expand All @@ -1911,7 +1911,7 @@ def _upload_artifact(artifact_path):
A request handler for `PUT /mlflow-artifacts/artifacts/<artifact_path>` to upload an artifact
to `artifact_path` (a relative path from the root artifact directory).
"""
validate_path_is_safe(artifact_path)
artifact_path = validate_path_is_safe(artifact_path)
head, tail = posixpath.split(artifact_path)
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = os.path.join(tmp_dir, tail)
Expand All @@ -1937,11 +1937,7 @@ def _list_artifacts_mlflow_artifacts():
(a relative path from the root artifact directory).
"""
request_message = _get_request_message(ListArtifactsMlflowArtifacts())
if request_message.HasField("path"):
validate_path_is_safe(request_message.path)
path = request_message.path
else:
path = None
path = validate_path_is_safe(request_message.path) if request_message.HasField("path") else None
artifact_repo = _get_artifact_repo_mlflow_artifacts()
files = []
for file_info in artifact_repo.list_artifacts(path):
Expand All @@ -1962,7 +1958,7 @@ def _delete_artifact_mlflow_artifacts(artifact_path):
A request handler for `DELETE /mlflow-artifacts/artifacts?path=<value>` to delete artifacts in
`path` (a relative path from the root artifact directory).
"""
validate_path_is_safe(artifact_path)
artifact_path = validate_path_is_safe(artifact_path)
_get_request_message(DeleteArtifact())
artifact_repo = _get_artifact_repo_mlflow_artifacts()
artifact_repo.delete_artifacts(artifact_path)
Expand All @@ -1984,7 +1980,7 @@ def _create_multipart_upload_artifact(artifact_path):
A request handler for `POST /mlflow-artifacts/mpu/create` to create a multipart upload
to `artifact_path` (a relative path from the root artifact directory).
"""
validate_path_is_safe(artifact_path)
artifact_path = validate_path_is_safe(artifact_path)

request_message = _get_request_message(
CreateMultipartUpload(),
Expand Down Expand Up @@ -2017,7 +2013,7 @@ def _complete_multipart_upload_artifact(artifact_path):
A request handler for `POST /mlflow-artifacts/mpu/complete` to complete a multipart upload
to `artifact_path` (a relative path from the root artifact directory).
"""
validate_path_is_safe(artifact_path)
artifact_path = validate_path_is_safe(artifact_path)

request_message = _get_request_message(
CompleteMultipartUpload(),
Expand Down Expand Up @@ -2050,7 +2046,7 @@ def _abort_multipart_upload_artifact(artifact_path):
A request handler for `POST /mlflow-artifacts/mpu/abort` to abort a multipart upload
to `artifact_path` (a relative path from the root artifact directory).
"""
validate_path_is_safe(artifact_path)
artifact_path = validate_path_is_safe(artifact_path)

request_message = _get_request_message(
AbortMultipartUpload(),
Expand Down
4 changes: 2 additions & 2 deletions mlflow/store/artifact/http_artifact_repo.py
Original file line number Diff line number Diff line change
Expand Up @@ -89,9 +89,9 @@ def list_artifacts(self, path=None):
augmented_raise_for_status(resp)
file_infos = []
for f in resp.json().get("files", []):
validate_path_is_safe(f["path"])
validated_path = validate_path_is_safe(f["path"])
file_info = FileInfo(
posixpath.join(path, f["path"]) if path else f["path"],
posixpath.join(path, validated_path) if path else validated_path,
f["is_dir"],
int(f["file_size"]) if ("file_size" in f) else None,
)
Expand Down
2 changes: 2 additions & 0 deletions mlflow/utils/uri.py
Original file line number Diff line number Diff line change
Expand Up @@ -453,6 +453,8 @@ def validate_path_is_safe(path):
):
raise exc

return path


def validate_query_string(query):
query = _decode(query)
Expand Down

0 comments on commit 5044878

Please sign in to comment.