Refactor server metadata to be more type safe

Prevents improper access to dynamic map keys
by unmarshalling the schema to stricter struct
with fields.

Part of https://gitlab.com/gitlab-org/gitaly/-/issues/3061

8 files changed, 31 insertions, 22 deletions

diff --git a/changelogs/unreleased/po-refacto-server-metadata.yml b/changelogs/unreleased/po-refacto-server-metadata.yml
new file mode 100644
index 000000000..e6432875b
--- /dev/null
+++ b/changelogs/unreleased/po-refacto-server-metadata.yml
@@ -0,0 +1,5 @@
+---
+title: Refactor server metadata to be more type safe
+merge_request: 2624
+author:
+type: other
diff --git a/internal/gitaly/service/remote/fetch_internal_remote.go b/internal/gitaly/service/remote/fetch_internal_remote.go
index 311fb1da5..1cc272144 100644
--- a/internal/gitaly/service/remote/fetch_internal_remote.go
+++ b/internal/gitaly/service/remote/fetch_internal_remote.go
@@ -75,18 +75,16 @@ func (s *server) FetchInternalRemote(ctx context.Context, req *gitalypb.FetchInt
 
 // getRemoteDefaultBranch gets the default branch of a repository hosted on another Gitaly node
 func (s *server) getRemoteDefaultBranch(ctx context.Context, repo *gitalypb.Repository) ([]byte, error) {
-	server, err := helper.ExtractGitalyServer(ctx, repo.StorageName)
+	serverInfo, err := helper.ExtractGitalyServer(ctx, repo.StorageName)
 	if err != nil {
 		return nil, fmt.Errorf("getRemoteDefaultBranch: %w", err)
 	}
 
-	address := server["address"]
-	if address == "" {
+	if serverInfo.Address == "" {
 		return nil, errors.New("getRemoteDefaultBranch: empty Gitaly address")
 	}
-	token := server["token"]
 
-	conn, err := s.conns.Dial(ctx, address, token)
+	conn, err := s.conns.Dial(ctx, serverInfo.Address, serverInfo.Token)
 	if err != nil {
 		return nil, fmt.Errorf("getRemoteDefaultBranch: %w", err)
 	}
diff --git a/internal/gitaly/service/repository/replicate.go b/internal/gitaly/service/repository/replicate.go
index 75a7508fc..eebeb7d75 100644
--- a/internal/gitaly/service/repository/replicate.go
+++ b/internal/gitaly/service/repository/replicate.go
@@ -237,7 +237,7 @@ func (s *server) newRepoClient(ctx context.Context, storageName string) (gitalyp
 		return nil, err
 	}
 
-	conn, err := s.conns.Dial(ctx, gitalyServerInfo["address"], gitalyServerInfo["token"])
+	conn, err := s.conns.Dial(ctx, gitalyServerInfo.Address, gitalyServerInfo.Token)
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/gitalyssh/gitalyssh.go b/internal/gitalyssh/gitalyssh.go
index ff4abba24..249c8ed9a 100644
--- a/internal/gitalyssh/gitalyssh.go
+++ b/internal/gitalyssh/gitalyssh.go
@@ -52,20 +52,17 @@ func commandEnv(ctx context.Context, storageName, command string, message proto.
 		return nil, status.Errorf(codes.InvalidArgument, "commandEnv: no storage info for %s", storageName)
 	}
 
-	address := storageInfo["address"]
-	if address == "" {
+	if storageInfo.Address == "" {
 		return nil, status.Errorf(codes.InvalidArgument, "commandEnv: empty gitaly address")
 	}
 
-	token := storageInfo["token"]
-
 	featureFlagPairs := featureflag.AllFlags(ctx)
 
 	return []string{
 		fmt.Sprintf("GITALY_PAYLOAD=%s", payload),
 		fmt.Sprintf("GIT_SSH_COMMAND=%s %s", gitalySSHPath(), command),
-		fmt.Sprintf("GITALY_ADDRESS=%s", address),
-		fmt.Sprintf("GITALY_TOKEN=%s", token),
+		fmt.Sprintf("GITALY_ADDRESS=%s", storageInfo.Address),
+		fmt.Sprintf("GITALY_TOKEN=%s", storageInfo.Token),
 		fmt.Sprintf("GITALY_FEATUREFLAGS=%s", strings.Join(featureFlagPairs, ",")),
 		fmt.Sprintf("CORRELATION_ID=%s", getCorrelationID(ctx)),
 		// Pass through the SSL_CERT_* variables that indicate which
diff --git a/internal/helper/storage.go b/internal/helper/storage.go
index 0a465756f..254e51fa6 100644
--- a/internal/helper/storage.go
+++ b/internal/helper/storage.go
@@ -36,15 +36,15 @@ func ExtractGitalyServers(ctx context.Context) (gitalyServersInfo storage.Gitaly
 }
 
 // ExtractGitalyServer extracts server information for a specific storage
-func ExtractGitalyServer(ctx context.Context, storageName string) (map[string]string, error) {
+func ExtractGitalyServer(ctx context.Context, storageName string) (storage.ServerInfo, error) {
 	gitalyServers, err := ExtractGitalyServers(ctx)
 	if err != nil {
-		return nil, err
+		return storage.ServerInfo{}, err
 	}
 
 	gitalyServer, ok := gitalyServers[storageName]
 	if !ok {
-		return nil, errors.New("storage name not found")
+		return storage.ServerInfo{}, errors.New("storage name not found")
 	}
 
 	return gitalyServer, nil
@@ -64,8 +64,8 @@ func IncomingToOutgoing(ctx context.Context) context.Context {
 func InjectGitalyServers(ctx context.Context, name, address, token string) (context.Context, error) {
 	gitalyServers := storage.GitalyServers{
 		name: {
-			"address": address,
-			"token":   token,
+			Address: address,
+			Token:   token,
 		},
 	}
 
diff --git a/internal/helper/storage_test.go b/internal/helper/storage_test.go
index 1e50e04cd..44e49b218 100644
--- a/internal/helper/storage_test.go
+++ b/internal/helper/storage_test.go
@@ -41,7 +41,7 @@ func TestExtractGitalyServers(t *testing.T) {
 		{
 			desc:     "properly-encoded string",
 			metadata: metadata.Pairs("gitaly-servers", base64.StdEncoding.EncodeToString([]byte(`{"default":{"address":"unix:///tmp/sock","token":"hunter1"}}`))),
-			info:     storage.GitalyServers{"default": {"address": "unix:///tmp/sock", "token": "hunter1"}},
+			info:     storage.GitalyServers{"default": storage.ServerInfo{Address: "unix:///tmp/sock", Token: "hunter1"}},
 		},
 	}
 
diff --git a/internal/storage/servers.go b/internal/storage/servers.go
index 8dbdf3675..87f3aae3c 100644
--- a/internal/storage/servers.go
+++ b/internal/storage/servers.go
@@ -1,5 +1,14 @@
 package storage
 
+// 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
+// the remote peer.
+type ServerInfo struct {
+	Address string `json:"address"`
+	Token   string `json:"token"`
+}
+
 // GitalyServers hold Gitaly servers info like {"default":{"token":"x","address":"y"}},
 // to be passed in `gitaly-servers` metadata.
-type GitalyServers map[string]map[string]string
+type GitalyServers map[string]ServerInfo
diff --git a/internal/testhelper/testhelper.go b/internal/testhelper/testhelper.go
index 134367ead..8a11b8c28 100644
--- a/internal/testhelper/testhelper.go
+++ b/internal/testhelper/testhelper.go
@@ -133,9 +133,9 @@ func GitlabTestStoragePath() string {
 // inter-gitaly operations.
 func GitalyServersMetadata(t testing.TB, serverSocketPath string) metadata.MD {
 	gitalyServers := storage.GitalyServers{
-		"default": {
-			"address": serverSocketPath,
-			"token":   RepositoryAuthToken,
+		"default": storage.ServerInfo{
+			Address: serverSocketPath,
+			Token:   RepositoryAuthToken,
 		},
 	}