Skip to content

Commit

Permalink
Fix URI local path traversal (#11473)
Browse files Browse the repository at this point in the history
  • Loading branch information
daniellok-db committed Mar 20, 2024
1 parent f588a9b commit 695d287
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 7 deletions.
8 changes: 4 additions & 4 deletions mlflow/server/handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -599,13 +599,13 @@ def _create_experiment():
tags = [ExperimentTag(tag.key, tag.value) for tag in request_message.tags]

# Validate query string in artifact location to prevent attacks
parsed_artifact_locaion = urllib.parse.urlparse(request_message.artifact_location)
if parsed_artifact_locaion.fragment:
parsed_artifact_location = urllib.parse.urlparse(request_message.artifact_location)
if parsed_artifact_location.fragment or parsed_artifact_location.params:
raise MlflowException(
"'artifact_location' URL can't include fragment part.",
"'artifact_location' URL can't include fragments or params.",
error_code=INVALID_PARAMETER_VALUE,
)
validate_query_string(parsed_artifact_locaion.query)
validate_query_string(parsed_artifact_location.query)
experiment_id = _get_tracking_store().create_experiment(
request_message.name, request_message.artifact_location, tags
)
Expand Down
15 changes: 12 additions & 3 deletions tests/server/test_handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -804,20 +804,29 @@ def test_delete_artifact_mlflow_artifacts_throws_for_malicious_path(enable_serve
assert json_response["message"] == "Invalid path"


def test_local_file_read_write_by_pass_vulnerability():
@pytest.mark.parametrize(
"uri",
[
"http://host#/abc/etc/",
"http://host/;..%2F..%2Fetc",
],
)
def test_local_file_read_write_by_pass_vulnerability(uri):
request = mock.MagicMock()
request.method = "POST"
request.content_type = "application/json; charset=utf-8"
request.get_json = mock.MagicMock()
request.get_json.return_value = {
"name": "hello",
"artifact_location": "http://host#/abc/etc/",
"artifact_location": uri,
}
msg = _get_request_message(CreateExperiment(), flask_request=request)
with mock.patch("mlflow.server.handlers._get_request_message", return_value=msg):
response = _create_experiment()
json_response = json.loads(response.get_data())
assert json_response["message"] == "'artifact_location' URL can't include fragment part."
assert (
json_response["message"] == "'artifact_location' URL can't include fragments or params."
)

# Test if source is a local filesystem path, `_validate_source` validates that the run
# artifact_uri is also a local filesystem path.
Expand Down

0 comments on commit 695d287

Please sign in to comment.