helper: Move server metadata into storage package

The helpers to inject and extract Gitaly server information from the
context are currently located in the generic "helper" package, even
though required definitions are part of the "storage" package already.
Move them over into the storage package to move related code closer to
each other.

12 files changed, 92 insertions, 95 deletions

diff --git a/internal/git/remoterepo/repository.go b/internal/git/remoterepo/repository.go
index db91bf147..69fa197f4 100644
--- a/internal/git/remoterepo/repository.go
+++ b/internal/git/remoterepo/repository.go
@@ -6,7 +6,7 @@ import (
 
 	"gitlab.com/gitlab-org/gitaly/v14/client"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/git"
-	"gitlab.com/gitlab-org/gitaly/v14/internal/helper"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/storage"
 	"gitlab.com/gitlab-org/gitaly/v14/proto/go/gitalypb"
 	"google.golang.org/grpc"
 )
@@ -19,7 +19,7 @@ type Repo struct {
 
 // New creates a new remote Repository from its protobuf representation.
 func New(ctx context.Context, repo *gitalypb.Repository, pool *client.Pool) (*Repo, error) {
-	server, err := helper.ExtractGitalyServer(ctx, repo.GetStorageName())
+	server, err := storage.ExtractGitalyServer(ctx, repo.GetStorageName())
 	if err != nil {
 		return nil, fmt.Errorf("remote repository: %w", err)
 	}
diff --git a/internal/git/remoterepo/repository_test.go b/internal/git/remoterepo/repository_test.go
index c04c9d38d..b209ef79d 100644
--- a/internal/git/remoterepo/repository_test.go
+++ b/internal/git/remoterepo/repository_test.go
@@ -11,7 +11,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/service"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/service/commit"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/service/repository"
-	"gitlab.com/gitlab-org/gitaly/v14/internal/helper"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/storage"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/metadata"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper/testcfg"
@@ -44,7 +44,7 @@ func TestRepository(t *testing.T) {
 	ctx, cancel := testhelper.Context()
 	defer cancel()
 
-	ctx, err := helper.InjectGitalyServers(ctx, "default", serverSocketPath, cfg.Auth.Token)
+	ctx, err := storage.InjectGitalyServers(ctx, "default", serverSocketPath, cfg.Auth.Token)
 	require.NoError(t, err)
 
 	pool := client.NewPool()
diff --git a/internal/gitaly/service/remote/fetch_internal_remote.go b/internal/gitaly/service/remote/fetch_internal_remote.go
index 7eafcf23b..37e5c716c 100644
--- a/internal/gitaly/service/remote/fetch_internal_remote.go
+++ b/internal/gitaly/service/remote/fetch_internal_remote.go
@@ -11,7 +11,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v14/internal/git/localrepo"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/config"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/service/ref"
-	"gitlab.com/gitlab-org/gitaly/v14/internal/helper"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/storage"
 	"gitlab.com/gitlab-org/gitaly/v14/proto/go/gitalypb"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -100,7 +100,7 @@ func (s *server) FetchInternalRemote(ctx context.Context, req *gitalypb.FetchInt
 
 // getRemoteDefaultBranch gets the default branch of a repository hosted on another Gitaly node.
 func getRemoteDefaultBranch(ctx context.Context, repo *gitalypb.Repository, conns *client.Pool) ([]byte, error) {
-	serverInfo, err := helper.ExtractGitalyServer(ctx, repo.StorageName)
+	serverInfo, err := storage.ExtractGitalyServer(ctx, repo.StorageName)
 	if err != nil {
 		return nil, fmt.Errorf("getRemoteDefaultBranch: %w", err)
 	}
diff --git a/internal/gitaly/service/remote/fetch_internal_remote_test.go b/internal/gitaly/service/remote/fetch_internal_remote_test.go
index 191724728..b4c5df9a4 100644
--- a/internal/gitaly/service/remote/fetch_internal_remote_test.go
+++ b/internal/gitaly/service/remote/fetch_internal_remote_test.go
@@ -20,7 +20,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/service/hook"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/service/ref"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/service/ssh"
-	"gitlab.com/gitlab-org/gitaly/v14/internal/helper"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/storage"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper/testcfg"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper/testserver"
@@ -200,7 +200,7 @@ func TestSuccessfulFetchInternalRemote(t *testing.T) {
 	ctx, cancel := testhelper.Context()
 	t.Cleanup(cancel)
 
-	ctx, err := helper.InjectGitalyServers(ctx, remoteRepo.GetStorageName(), remoteAddr, "")
+	ctx, err := storage.InjectGitalyServers(ctx, remoteRepo.GetStorageName(), remoteAddr, "")
 	require.NoError(t, err)
 
 	//nolint:staticcheck
diff --git a/internal/gitaly/service/repository/replicate.go b/internal/gitaly/service/repository/replicate.go
index b55d3fc96..c372d344a 100644
--- a/internal/gitaly/service/repository/replicate.go
+++ b/internal/gitaly/service/repository/replicate.go
@@ -326,7 +326,7 @@ func (s *server) writeFile(ctx context.Context, path string, mode os.FileMode, r
 
 // newRepoClient creates a new RepositoryClient that talks to the gitaly of the source repository
 func (s *server) newRepoClient(ctx context.Context, storageName string) (gitalypb.RepositoryServiceClient, error) {
-	gitalyServerInfo, err := helper.ExtractGitalyServer(ctx, storageName)
+	gitalyServerInfo, err := storage.ExtractGitalyServer(ctx, storageName)
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/gitaly/service/repository/server_test.go b/internal/gitaly/service/repository/server_test.go
index 0c8cc0f22..546eb84dd 100644
--- a/internal/gitaly/service/repository/server_test.go
+++ b/internal/gitaly/service/repository/server_test.go
@@ -5,7 +5,7 @@ import (
 
 	"github.com/stretchr/testify/require"
 	"gitlab.com/gitlab-org/gitaly/v14/client"
-	"gitlab.com/gitlab-org/gitaly/v14/internal/helper"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/storage"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper"
 	"google.golang.org/grpc/metadata"
 )
@@ -21,7 +21,7 @@ func TestGetConnectionByStorage(t *testing.T) {
 	defer cancel()
 
 	storageName, address := "default", "unix:///fake/address/wont/work"
-	injectedCtx, err := helper.InjectGitalyServers(ctx, storageName, address, "token")
+	injectedCtx, err := storage.InjectGitalyServers(ctx, storageName, address, "token")
 	require.NoError(t, err)
 
 	md, ok := metadata.FromOutgoingContext(injectedCtx)
diff --git a/internal/gitaly/storage/servers.go b/internal/gitaly/storage/servers.go
index 87f3aae3c..c9351fc87 100644
--- a/internal/gitaly/storage/servers.go
+++ b/internal/gitaly/storage/servers.go
@@ -1,5 +1,15 @@
 package storage
 
+import (
+	"context"
+	"encoding/base64"
+	"encoding/json"
+	"errors"
+	"fmt"
+
+	"google.golang.org/grpc/metadata"
+)
+
 // ServerInfo contains information about how to reach a Gitaly server or a
 // Praefect virtual storage. This is necessary for Gitaly RPC's involving more
 // than one Gitaly. Without this information, Gitaly would not know how to reach
@@ -12,3 +22,63 @@ type ServerInfo struct {
 // GitalyServers hold Gitaly servers info like {"default":{"token":"x","address":"y"}},
 // to be passed in `gitaly-servers` metadata.
 type GitalyServers map[string]ServerInfo
+
+// ErrEmptyMetadata indicates that the gRPC metadata was not found in the
+// context
+var ErrEmptyMetadata = errors.New("empty metadata")
+
+// ExtractGitalyServers extracts `storage.GitalyServers` from an incoming context.
+func ExtractGitalyServers(ctx context.Context) (gitalyServersInfo GitalyServers, err error) {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		return nil, ErrEmptyMetadata
+	}
+
+	gitalyServersJSONEncoded := md["gitaly-servers"]
+	if len(gitalyServersJSONEncoded) == 0 {
+		return nil, fmt.Errorf("empty gitaly-servers metadata")
+	}
+
+	gitalyServersJSON, err := base64.StdEncoding.DecodeString(gitalyServersJSONEncoded[0])
+	if err != nil {
+		return nil, fmt.Errorf("failed decoding base64: %v", err)
+	}
+
+	if err := json.Unmarshal(gitalyServersJSON, &gitalyServersInfo); err != nil {
+		return nil, fmt.Errorf("failed unmarshalling json: %v", err)
+	}
+
+	return
+}
+
+// ExtractGitalyServer extracts server information for a specific storage
+func ExtractGitalyServer(ctx context.Context, storageName string) (ServerInfo, error) {
+	gitalyServers, err := ExtractGitalyServers(ctx)
+	if err != nil {
+		return ServerInfo{}, err
+	}
+
+	gitalyServer, ok := gitalyServers[storageName]
+	if !ok {
+		return ServerInfo{}, errors.New("storage name not found")
+	}
+
+	return gitalyServer, nil
+}
+
+// InjectGitalyServers injects gitaly-servers metadata into an outgoing context
+func InjectGitalyServers(ctx context.Context, name, address, token string) (context.Context, error) {
+	gitalyServers := GitalyServers{
+		name: {
+			Address: address,
+			Token:   token,
+		},
+	}
+
+	gitalyServersJSON, err := json.Marshal(gitalyServers)
+	if err != nil {
+		return nil, err
+	}
+
+	return metadata.AppendToOutgoingContext(ctx, "gitaly-servers", base64.StdEncoding.EncodeToString(gitalyServersJSON)), nil
+}
diff --git a/internal/helper/storage_test.go b/internal/gitaly/storage/servers_test.go
index 486efc12d..5c6def8da 100644
--- a/internal/helper/storage_test.go
+++ b/internal/gitaly/storage/servers_test.go
@@ -1,4 +1,4 @@
-package helper
+package storage_test
 
 import (
 	"context"
@@ -49,7 +49,7 @@ func TestExtractGitalyServers(t *testing.T) {
 		t.Run(testCase.desc, func(t *testing.T) {
 			ctx := metadata.NewIncomingContext(ctxOuter, testCase.metadata)
 
-			info, err := ExtractGitalyServers(ctx)
+			info, err := storage.ExtractGitalyServers(ctx)
 			if testCase.info == nil {
 				require.Error(t, err)
 			} else {
@@ -64,7 +64,7 @@ func TestInjectGitalyServers(t *testing.T) {
 	check := func(t *testing.T, ctx context.Context) {
 		t.Helper()
 
-		newCtx, err := InjectGitalyServers(ctx, "gitaly-1", "1.1.1.1", "secret")
+		newCtx, err := storage.InjectGitalyServers(ctx, "gitaly-1", "1.1.1.1", "secret")
 		require.NoError(t, err)
 
 		md, found := metadata.FromOutgoingContext(newCtx)
diff --git a/internal/gitalyssh/gitalyssh.go b/internal/gitalyssh/gitalyssh.go
index b882f7c78..4e8a1b410 100644
--- a/internal/gitalyssh/gitalyssh.go
+++ b/internal/gitalyssh/gitalyssh.go
@@ -8,7 +8,7 @@ import (
 	"strings"
 
 	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/config"
-	"gitlab.com/gitlab-org/gitaly/v14/internal/helper"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/storage"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/metadata/featureflag"
 	gitalyx509 "gitlab.com/gitlab-org/gitaly/v14/internal/x509"
 	"gitlab.com/gitlab-org/gitaly/v14/proto/go/gitalypb"
@@ -43,7 +43,7 @@ func commandEnv(ctx context.Context, cfg config.Cfg, storageName, command string
 		return nil, status.Errorf(codes.Internal, "commandEnv: marshalling payload failed: %v", err)
 	}
 
-	serversInfo, err := helper.ExtractGitalyServers(ctx)
+	serversInfo, err := storage.ExtractGitalyServers(ctx)
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "commandEnv: extracting Gitaly servers: %v", err)
 	}
diff --git a/internal/helper/storage.go b/internal/helper/storage.go
deleted file mode 100644
index 3bbf277c2..000000000
--- a/internal/helper/storage.go
+++ /dev/null
@@ -1,72 +0,0 @@
-package helper
-
-import (
-	"context"
-	"encoding/base64"
-	"encoding/json"
-	"errors"
-	"fmt"
-
-	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/storage"
-	"google.golang.org/grpc/metadata"
-)
-
-// ErrEmptyMetadata indicates that the gRPC metadata was not found in the
-// context
-var ErrEmptyMetadata = errors.New("empty metadata")
-
-// ExtractGitalyServers extracts `storage.GitalyServers` from an incoming context.
-func ExtractGitalyServers(ctx context.Context) (gitalyServersInfo storage.GitalyServers, err error) {
-	md, ok := metadata.FromIncomingContext(ctx)
-	if !ok {
-		return nil, ErrEmptyMetadata
-	}
-
-	gitalyServersJSONEncoded := md["gitaly-servers"]
-	if len(gitalyServersJSONEncoded) == 0 {
-		return nil, fmt.Errorf("empty gitaly-servers metadata")
-	}
-
-	gitalyServersJSON, err := base64.StdEncoding.DecodeString(gitalyServersJSONEncoded[0])
-	if err != nil {
-		return nil, fmt.Errorf("failed decoding base64: %v", err)
-	}
-
-	if err := json.Unmarshal(gitalyServersJSON, &gitalyServersInfo); err != nil {
-		return nil, fmt.Errorf("failed unmarshalling json: %v", err)
-	}
-
-	return
-}
-
-// ExtractGitalyServer extracts server information for a specific storage
-func ExtractGitalyServer(ctx context.Context, storageName string) (storage.ServerInfo, error) {
-	gitalyServers, err := ExtractGitalyServers(ctx)
-	if err != nil {
-		return storage.ServerInfo{}, err
-	}
-
-	gitalyServer, ok := gitalyServers[storageName]
-	if !ok {
-		return storage.ServerInfo{}, errors.New("storage name not found")
-	}
-
-	return gitalyServer, nil
-}
-
-// InjectGitalyServers injects gitaly-servers metadata into an outgoing context
-func InjectGitalyServers(ctx context.Context, name, address, token string) (context.Context, error) {
-	gitalyServers := storage.GitalyServers{
-		name: {
-			Address: address,
-			Token:   token,
-		},
-	}
-
-	gitalyServersJSON, err := json.Marshal(gitalyServers)
-	if err != nil {
-		return nil, err
-	}
-
-	return metadata.AppendToOutgoingContext(ctx, "gitaly-servers", base64.StdEncoding.EncodeToString(gitalyServersJSON)), nil
-}
diff --git a/internal/metadata/metadata_test.go b/internal/metadata/metadata_test.go
index 6068c8fe8..71413549e 100644
--- a/internal/metadata/metadata_test.go
+++ b/internal/metadata/metadata_test.go
@@ -6,21 +6,20 @@ import (
 
 	"github.com/stretchr/testify/require"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/storage"
-	"gitlab.com/gitlab-org/gitaly/v14/internal/helper"
 )
 
 func TestOutgoingToIncoming(t *testing.T) {
 	ctx := context.Background()
-	ctx, err := helper.InjectGitalyServers(ctx, "a", "b", "c")
+	ctx, err := storage.InjectGitalyServers(ctx, "a", "b", "c")
 	require.NoError(t, err)
 
-	_, err = helper.ExtractGitalyServer(ctx, "a")
-	require.Equal(t, helper.ErrEmptyMetadata, err,
+	_, err = storage.ExtractGitalyServer(ctx, "a")
+	require.Equal(t, storage.ErrEmptyMetadata, err,
 		"server should not be found in the incoming context")
 
 	ctx = OutgoingToIncoming(ctx)
 
-	info, err := helper.ExtractGitalyServer(ctx, "a")
+	info, err := storage.ExtractGitalyServer(ctx, "a")
 	require.NoError(t, err)
 	require.Equal(t, storage.ServerInfo{Address: "b", Token: "c"}, info)
 }
diff --git a/internal/praefect/replicator.go b/internal/praefect/replicator.go
index 8c251ab11..8d3cbdf4f 100644
--- a/internal/praefect/replicator.go
+++ b/internal/praefect/replicator.go
@@ -10,7 +10,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/sirupsen/logrus"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/service/repository"
-	"gitlab.com/gitlab-org/gitaly/v14/internal/helper"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/storage"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/middleware/metadatahandler"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/config"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/datastore"
@@ -714,7 +714,7 @@ func (r ReplMgr) processReplicationEvent(ctx context.Context, event datastore.Re
 			return fmt.Errorf("no connection to source node %q/%q", event.Job.VirtualStorage, event.Job.SourceNodeStorage)
 		}
 
-		ctx, err = helper.InjectGitalyServers(ctx, event.Job.SourceNodeStorage, source.Address, source.Token)
+		ctx, err = storage.InjectGitalyServers(ctx, event.Job.SourceNodeStorage, source.Address, source.Token)
 		if err != nil {
 			return fmt.Errorf("inject Gitaly servers into context: %w", err)
 		}