git: Drop `GetBranch()` function in favor of Reference helpers

The `GetBranch()` function is a small wrapper around `GetReference()`
which simply resolves the given string to a fully qualified reference
first. This commit implements an alternative set of helper functions to
create a Reference from a potentially unqualified branch name and get a
Reference's branch name and converts the codebase to use them instead of
`GetBranch()`.

6 files changed, 111 insertions, 75 deletions

diff --git a/internal/git/reference.go b/internal/git/reference.go
index 23b7b1a67..fee318bf8 100644
--- a/internal/git/reference.go
+++ b/internal/git/reference.go
@@ -3,6 +3,7 @@ package git
 import (
 	"bytes"
 	"context"
+	"strings"
 )
 
 // Revision represents anything that resolves to either a commit, multiple
@@ -21,6 +22,21 @@ func (r Revision) String() string {
 // Revision does and must always contain a fully qualified reference.
 type ReferenceName string
 
+// NewBranchReferenceName returns a new ReferenceName for the given branch. The
+// branch may either be a fully qualified branch name ("refs/heads/master") or
+// an unqualified name ("master"). In the latter case, the returned
+// ReferenceName will still be a fully qualified reference by prepending
+// "refs/heads".
+func NewBranchReferenceName(branch string) ReferenceName {
+	if strings.HasPrefix(branch, "refs/heads/") {
+		return ReferenceName(branch)
+	}
+	if strings.HasPrefix(branch, "heads/") {
+		return ReferenceName("refs/" + branch)
+	}
+	return ReferenceName("refs/heads/" + branch)
+}
+
 // String returns the string representation of the ReferenceName.
 func (r ReferenceName) String() string {
 	return string(r)
@@ -32,6 +48,16 @@ func (r ReferenceName) Revision() Revision {
 	return Revision(r)
 }
 
+// Branch returns `true` and the branch name if the reference is a branch. E.g.
+// if ReferenceName is "refs/heads/master", it will return "master". If it is
+// not a branch, `false` is returned.
+func (r ReferenceName) Branch() (string, bool) {
+	if strings.HasPrefix(r.String(), "refs/heads/") {
+		return r.String()[len("refs/heads/"):], true
+	}
+	return "", false
+}
+
 // Reference represents a Git reference.
 type Reference struct {
 	// Name is the name of the reference
diff --git a/internal/git/reference_test.go b/internal/git/reference_test.go
index cd209fc10..7b3b58738 100644
--- a/internal/git/reference_test.go
+++ b/internal/git/reference_test.go
@@ -51,3 +51,77 @@ func TestCheckRefFormat(t *testing.T) {
 		})
 	}
 }
+
+func TestReferenceName_NewBranchReferenceName(t *testing.T) {
+	for _, tc := range []struct {
+		desc      string
+		reference string
+		expected  string
+	}{
+		{
+			desc:      "fully qualified reference",
+			reference: "refs/heads/master",
+			expected:  "refs/heads/master",
+		},
+		{
+			desc:      "partly qualified reference",
+			reference: "heads/master",
+			expected:  "refs/heads/master",
+		},
+		{
+			desc:      "unqualified reference",
+			reference: "master",
+			expected:  "refs/heads/master",
+		},
+		{
+			desc:      "weird branch name",
+			reference: "refs/master",
+			expected:  "refs/heads/refs/master",
+		},
+		{
+			desc:      "tag is treated as a branch",
+			reference: "refs/tags/master",
+			expected:  "refs/heads/refs/tags/master",
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			ref := NewBranchReferenceName(tc.reference)
+			require.Equal(t, ref.String(), tc.expected)
+		})
+	}
+}
+
+func TestReferenceName_Branch(t *testing.T) {
+	for _, tc := range []struct {
+		desc      string
+		reference string
+		expected  string
+	}{
+		{
+			desc:      "fully qualified reference",
+			reference: "refs/heads/master",
+			expected:  "master",
+		},
+		{
+			desc:      "nested branch",
+			reference: "refs/heads/foo/master",
+			expected:  "foo/master",
+		},
+		{
+			desc:      "unqualified branch is not a branch",
+			reference: "master",
+			expected:  "",
+		},
+		{
+			desc:      "tag is not a branch",
+			reference: "refs/tags/master",
+			expected:  "",
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			branch, ok := ReferenceName(tc.reference).Branch()
+			require.Equal(t, tc.expected, branch)
+			require.Equal(t, tc.expected != "", ok)
+		})
+	}
+}
diff --git a/internal/git/repository.go b/internal/git/repository.go
index 6541f0ba0..b80820856 100644
--- a/internal/git/repository.go
+++ b/internal/git/repository.go
@@ -405,19 +405,6 @@ func (repo *LocalRepository) getReferences(ctx context.Context, pattern string,
 	return refs, nil
 }
 
-// GetBranch looks up and returns the given branch. Returns a
-// ErrReferenceNotFound if it wasn't found.
-func (repo *LocalRepository) GetBranch(ctx context.Context, branch string) (Reference, error) {
-	if strings.HasPrefix(branch, "refs/heads/") {
-		return repo.GetReference(ctx, ReferenceName(branch))
-	}
-
-	if strings.HasPrefix(branch, "heads/") {
-		branch = strings.TrimPrefix(branch, "heads/")
-	}
-	return repo.GetReference(ctx, ReferenceName("refs/heads/"+branch))
-}
-
 // GetBranches returns all branches.
 func (repo *LocalRepository) GetBranches(ctx context.Context) ([]Reference, error) {
 	return repo.GetReferences(ctx, "refs/heads/")
diff --git a/internal/git/repository_test.go b/internal/git/repository_test.go
index 663acdeff..1ff223ce5 100644
--- a/internal/git/repository_test.go
+++ b/internal/git/repository_test.go
@@ -121,55 +121,6 @@ func TestLocalRepository_GetReference(t *testing.T) {
 	}
 }
 
-func TestLocalRepository_GetBranch(t *testing.T) {
-	ctx, cancel := testhelper.Context()
-	defer cancel()
-
-	testRepo, _, cleanup := testhelper.NewTestRepo(t)
-	defer cleanup()
-
-	repo := NewRepository(testRepo)
-
-	testcases := []struct {
-		desc     string
-		ref      string
-		expected Reference
-	}{
-		{
-			desc:     "fully qualified master branch",
-			ref:      "refs/heads/master",
-			expected: NewReference("refs/heads/master", MasterID),
-		},
-		{
-			desc:     "half-qualified master branch",
-			ref:      "heads/master",
-			expected: NewReference("refs/heads/master", MasterID),
-		},
-		{
-			desc:     "fully qualified master branch",
-			ref:      "master",
-			expected: NewReference("refs/heads/master", MasterID),
-		},
-		{
-			desc:     "nonexistent branch",
-			ref:      "nonexistent",
-			expected: Reference{},
-		},
-	}
-
-	for _, tc := range testcases {
-		t.Run(tc.desc, func(t *testing.T) {
-			ref, err := repo.GetBranch(ctx, tc.ref)
-			if tc.expected.Name == "" {
-				require.True(t, errors.Is(err, ErrReferenceNotFound))
-			} else {
-				require.NoError(t, err)
-				require.Equal(t, tc.expected, ref)
-			}
-		})
-	}
-}
-
 func TestLocalRepository_GetReferences(t *testing.T) {
 	ctx, cancel := testhelper.Context()
 	defer cancel()
diff --git a/internal/gitaly/service/operations/submodules.go b/internal/gitaly/service/operations/submodules.go
index 8900ea95f..32eaf6a99 100644
--- a/internal/gitaly/service/operations/submodules.go
+++ b/internal/gitaly/service/operations/submodules.go
@@ -87,7 +87,7 @@ func (s *Server) userUpdateSubmodule(ctx context.Context, req *gitalypb.UserUpda
 		}, nil
 	}
 
-	branchRef, err := repo.GetBranch(ctx, string(req.GetBranch()))
+	branchRef, err := repo.GetReference(ctx, git.NewBranchReferenceName(string(req.GetBranch())))
 	if err != nil {
 		if errors.Is(err, git.ErrReferenceNotFound) {
 			return nil, helper.ErrInvalidArgumentf("Cannot find branch")
diff --git a/internal/gitaly/service/ref/branches.go b/internal/gitaly/service/ref/branches.go
index 6da18fb49..066e00174 100644
--- a/internal/gitaly/service/ref/branches.go
+++ b/internal/gitaly/service/ref/branches.go
@@ -3,7 +3,6 @@ package ref
 import (
 	"context"
 	"errors"
-	"strings"
 
 	"gitlab.com/gitlab-org/gitaly/internal/git"
 	"gitlab.com/gitlab-org/gitaly/internal/git/log"
@@ -13,20 +12,14 @@ import (
 )
 
 func (s *server) FindBranch(ctx context.Context, req *gitalypb.FindBranchRequest) (*gitalypb.FindBranchResponse, error) {
-	refName := string(req.GetName())
-	if len(refName) == 0 {
+	if len(req.GetName()) == 0 {
 		return nil, status.Errorf(codes.InvalidArgument, "Branch name cannot be empty")
 	}
 
-	if strings.HasPrefix(refName, "refs/heads/") {
-		refName = strings.TrimPrefix(refName, "refs/heads/")
-	} else if strings.HasPrefix(refName, "heads/") {
-		refName = strings.TrimPrefix(refName, "heads/")
-	}
-
 	repo := req.GetRepository()
 
-	branch, err := git.NewRepository(repo).GetBranch(ctx, refName)
+	branchName := git.NewBranchReferenceName(string(req.GetName()))
+	branchRef, err := git.NewRepository(repo).GetReference(ctx, branchName)
 	if err != nil {
 		if errors.Is(err, git.ErrReferenceNotFound) {
 			return &gitalypb.FindBranchResponse{}, nil
@@ -34,14 +27,19 @@ func (s *server) FindBranch(ctx context.Context, req *gitalypb.FindBranchRequest
 		return nil, err
 	}
 
-	commit, err := log.GetCommit(ctx, s.locator, repo, git.Revision(branch.Target))
+	commit, err := log.GetCommit(ctx, s.locator, repo, git.Revision(branchRef.Target))
 	if err != nil {
 		return nil, err
 	}
 
+	branch, ok := branchName.Branch()
+	if !ok {
+		return nil, status.Errorf(codes.InvalidArgument, "reference is not a branch")
+	}
+
 	return &gitalypb.FindBranchResponse{
 		Branch: &gitalypb.Branch{
-			Name:         []byte(refName),
+			Name:         []byte(branch),
 			TargetCommit: commit,
 		},
 	}, nil