Handle trailing slashes when parsing a repo url in init

dagshub/common/init.py

@@ -20,6 +20,60 @@
 git = lazy_load("git")
 
 
+def _sanitize_repo_url(parsed: urllib.parse.ParseResult, path: str) -> str:
+    """Rebuild a url from its parsed pieces, keeping only scheme, host and path.
+
+    Any userinfo is dropped: `url` is written into MLflow's tracking URI and
+    into .dvc/config, and .dvc/config is a committed file, so a token pasted
+    into the url as "https://user:token@dagshub.com/owner/repo" would end up in
+    the repository. Credentials for both are set separately from the token.
+    """
+
+    host = parsed.hostname or ""
+    if parsed.port:
+        host = f"{host}:{parsed.port}"
+    if not parsed.scheme and not parsed.netloc:
+        # A bare "owner/repo" - nothing to rebuild around.
+        return path
+    return urllib.parse.urlunparse((parsed.scheme, host, path, "", "", ""))
+
+
+def _parse_repo_url(url: str) -> tuple:
+    """Extract (url, owner, name) from a repository url.
+
+    The owner and name are taken from the url *path*, not from the last two
+    slash-separated pieces of the whole string. Splitting the whole string lets
+    the hostname stand in for the owner: "https://dagshub.com/my-repo" has only
+    one path segment, but its last two pieces are "dagshub.com" and "my-repo",
+    which looks valid and is not.
+
+    The url is returned alongside them, rebuilt from the same parsed segments,
+    so that the value handed to MLflow and DVC agrees with the owner and name
+    that were derived from it. Parsing alone is not enough: empty segments are
+    skipped when reading owner and name, so without this
+    "https://dagshub.com/my-org//my-repo" would yield the right pair while
+    still sending a doubled slash downstream.
+
+    A bare "owner/repo" with no scheme is still accepted, since that is a
+    reasonable thing to pass.
+    """
+
+    parsed = urllib.parse.urlparse(url)
+    # With no scheme, urlparse puts everything in `path`, which is what we want
+    # for a bare "owner/repo"; with one, `netloc` holds the host and is skipped.
+    segments = [segment for segment in parsed.path.split("/") if segment]
+
+    if len(segments) < 2:
+        # The url is reported back redacted - a raw one can carry a token in
+        # its userinfo or query, and this message is likely to be logged.
+        raise ValueError(
+            f"Could not determine the repo owner and name from the url "
+            f"{_sanitize_repo_url(parsed, parsed.path)!r}. "
+            f"Expected a url of the form https://dagshub.com/<owner>/<repo>."
+        )
+    return _sanitize_repo_url(parsed, "/".join(segments)), segments[-2], segments[-1]
+
+
 def init(
     repo_name: Optional[str] = None,
     repo_owner: Optional[str] = None,
@@ -76,11 +130,12 @@ def init(
             repo, branch = determine_repo(root)
             url = repo.repo_url
 
+        # Strip a trailing slash first, so that a URL copied from the browser
+        # address bar doesn't shift every segment by one.
+        url = url.rstrip("/")
         if url.endswith(".git"):
             url = url[:-4]
-        # Extract the owner and name from the repo_url
-        parts = url.split("/")
-        repo_owner, repo_name = parts[-2], parts[-1]
+        url, repo_owner, repo_name = _parse_repo_url(url)
 
     # Create the repo if it wasn't created
     repo_api = RepoAPI(f"{repo_owner}/{repo_name}", host=host)
@@ -93,9 +148,7 @@ def init(
             should_create_under_org = repo_owner != current_user.username
 
         if should_create_under_org:
-            log_message(
-                f'Repository {repo_name} doesn\'t exist, creating it under organization "{repo_owner}".'
-            )
+            log_message(f'Repository {repo_name} doesn\'t exist, creating it under organization "{repo_owner}".')
             create_repo(repo_name, org_name=repo_owner, host=host)
         else:
             log_message(f"Repository {repo_name} doesn't exist, creating it under current user.")

tests/common/test_init.py

@@ -1,3 +1,4 @@
+import os
 from unittest.mock import MagicMock
 
 import pytest
@@ -84,3 +85,99 @@ def test_init_creates_repo_under_org_from_url(
     mock_log_message.assert_any_call(
         'Repository my-repo doesn\'t exist, creating it under organization "my-org".'
     )
+
+
+@pytest.mark.parametrize(
+    "url",
+    [
+        "https://dagshub.com/my-org/my-repo/",
+        "https://dagshub.com/my-org/my-repo.git/",
+        "https://dagshub.com/my-org/my-repo///",
+    ],
+)
+def test_init_from_url_tolerates_trailing_slash(
+    url, mock_repo_api, mock_user_api, mock_create_repo, mock_get_token, mock_log_message
+):
+    """
+    A url copied from the browser address bar carries a trailing slash, which
+    used to shift every segment by one and yield repo_owner="my-repo" with an
+    empty repo_name.
+    """
+    dagshub.init(url=url, mlflow=False, dvc=False)
+
+    mock_repo_api.assert_called_once_with("my-org/my-repo", host="https://dagshub.com")
+    mock_create_repo.assert_called_once_with("my-repo", org_name="my-org", host="https://dagshub.com")
+
+
+@pytest.mark.parametrize(
+    "url",
+    [
+        "https://dagshub.com/",
+        "https://dagshub.com",
+        "my-repo",
+        # One path segment only: the last two slash-separated pieces are
+        # "dagshub.com" and "my-repo", which looks like owner/name and is not.
+        "https://dagshub.com/my-repo",
+        "https://dagshub.com/my-repo/",
+    ],
+)
+def test_init_from_url_without_owner_and_name_raises(url, mock_get_token):
+    """
+    A url that carries no owner/name pair should fail loudly instead of
+    calling the API with empty segments.
+    """
+    with pytest.raises(ValueError, match="Could not determine the repo owner and name"):
+        dagshub.init(url=url, mlflow=False, dvc=False)
+
+
+@pytest.fixture
+def clean_environ(mocker):
+    """Let the test read what init() wrote to os.environ, then put it back."""
+    return mocker.patch.dict(os.environ, {}, clear=False)
+
+
+@pytest.mark.parametrize(
+    "url",
+    [
+        "https://dagshub.com/my-org//my-repo",
+        "https://dagshub.com//my-org/my-repo/",
+        "https://dagshub.com/my-org/my-repo.git",
+    ],
+)
+def test_init_hands_a_canonical_url_to_mlflow(
+    url, clean_environ, mock_repo_api, mock_user_api, mock_create_repo, mock_get_token, mock_log_message
+):
+    """
+    Empty path segments are skipped when reading the owner and name, so they
+    must be skipped in the url too - otherwise the API is called with
+    "my-org/my-repo" while MLflow is pointed at ".../my-org//my-repo.mlflow".
+    """
+    dagshub.init(url=url, mlflow=True, dvc=False)
+
+    mock_repo_api.assert_called_once_with("my-org/my-repo", host="https://dagshub.com")
+    assert os.environ["MLFLOW_TRACKING_URI"] == "https://dagshub.com/my-org/my-repo.mlflow"
+
+
+def test_init_does_not_leak_url_credentials_to_mlflow(
+    clean_environ, mock_repo_api, mock_user_api, mock_create_repo, mock_get_token, mock_log_message
+):
+    """
+    A token pasted into the url must not travel on into the tracking URI - the
+    same url is also written to .dvc/config, which is committed.
+    """
+    dagshub.init(url="https://user:s3cret-token@dagshub.com/my-org/my-repo", mlflow=True, dvc=False)
+
+    assert os.environ["MLFLOW_TRACKING_URI"] == "https://dagshub.com/my-org/my-repo.mlflow"
+    assert "s3cret-token" not in os.environ["MLFLOW_TRACKING_URI"]
+
+
+def test_init_does_not_leak_url_credentials_in_the_error_message(mock_get_token):
+    """
+    The url is quoted back when it cannot be parsed, and that message is likely
+    to be logged, so it must not carry the userinfo it was given.
+    """
+    with pytest.raises(ValueError) as excinfo:
+        dagshub.init(url="https://user:s3cret-token@dagshub.com/my-repo", mlflow=False, dvc=False)
+
+    assert "s3cret-token" not in str(excinfo.value)
+    assert "user:" not in str(excinfo.value)