From 20e37eaf65b9a7c90ac4f6a6a49e670e88f994b3 Mon Sep 17 00:00:00 2001
From: Jingfang Liu <jingfangliu@google.com>
Date: Thu, 13 Dec 2018 10:35:03 -0800
Subject: [PATCH] improve url parsing function in gitcloner

---
 pkg/loader/gitcloner.go      | 19 +++++++++++++++----
 pkg/loader/gitcloner_test.go | 14 ++++++++++++++
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/pkg/loader/gitcloner.go b/pkg/loader/gitcloner.go
index 506dfe02b..01d21d4cd 100644
--- a/pkg/loader/gitcloner.go
+++ b/pkg/loader/gitcloner.go
@@ -44,7 +44,7 @@ func isRepoUrl(arg string) bool {
 		(strings.HasPrefix(arg, "git::") ||
 			strings.HasPrefix(arg, "gh:") ||
 			strings.HasPrefix(arg, "github.com") ||
-			strings.HasPrefix(arg, "git@github.com:") ||
+			strings.HasPrefix(arg, "git@") ||
 			strings.Index(arg, "github.com/") > -1 ||
 			isAzureHost(arg) || isAWSHost(arg))
 }
@@ -106,7 +106,10 @@ func parseUrl(n string) (
 	return
 }
 
-const refQuery = "?ref="
+const (
+	refQuery  = "?ref="
+	gitSuffix = ".git"
+)
 
 // From strings like git@github.com:someOrg/someRepo.git or
 // https://github.com/someOrg/someRepo?ref=someHash, extract
@@ -116,8 +119,16 @@ func parseGithubUrl(n string) (
 	host, n = parseHostSpec(n)
 	host = normalizeGitHostSpec(host)
 
-	if strings.HasSuffix(n, ".git") {
-		n = n[0 : len(n)-len(".git")]
+	if strings.HasSuffix(n, gitSuffix) {
+		repo = n[0 : len(n)-len(gitSuffix)]
+		return
+	}
+	if strings.Contains(n, gitSuffix) {
+		index := strings.Index(n, gitSuffix)
+		repo = n[0:index]
+		n = n[index+len(gitSuffix):]
+		path, gitRef = peelQuery(n)
+		return
 	}
 	i := strings.Index(n, "/")
 	if i < 1 {
diff --git a/pkg/loader/gitcloner_test.go b/pkg/loader/gitcloner_test.go
index 732d074ee..672cf16a8 100644
--- a/pkg/loader/gitcloner_test.go
+++ b/pkg/loader/gitcloner_test.go
@@ -52,6 +52,14 @@ func TestIsRepoURL(t *testing.T) {
 			input:    "git::https://gitlab.com/org/repo",
 			expected: true,
 		},
+		{
+			input:    "git@gitlab2.sqtools.ru:10022/infra/kubernetes/thanos-base.git?ref=v0.1.0",
+			expected: true,
+		},
+		{
+			input:    "git@bitbucket.org:org/repo.git",
+			expected: true,
+		},
 		{
 			input:    "/github.com/org/repo",
 			expected: false,
@@ -283,6 +291,12 @@ func TestParseUrl(t *testing.T) {
 			path:  "somedir",
 			ref:   "",
 		},
+		{
+			input: "git@gitlab2.sqtools.ru:10022/infra/kubernetes/thanos-base.git?ref=v0.1.0",
+			repo:  "git@gitlab2.sqtools.ru:10022/infra/kubernetes/thanos-base.git",
+			path:  "",
+			ref:   "v0.1.0",
+		},
 	}
 	for _, testcase := range testcases {
 		repo, path, ref, err := parseUrl(testcase.input)