Merge branch 'backup_gitaly_servers_env' into 'master'

gitaly-backup: Allow using env GITALY_SERVERS to setup gitaly connection info

See merge request https://gitlab.com/gitlab-org/gitaly/-/merge_requests/5384

Merged-by: Quang-Minh Nguyen <qmnguyen@gitlab.com>
Approved-by: Justin Tobler <jtobler@gitlab.com>
Approved-by: Quang-Minh Nguyen <qmnguyen@gitlab.com>
Reviewed-by: James Fargher <proglottis@gmail.com>
Reviewed-by: Quang-Minh Nguyen <qmnguyen@gitlab.com>
Reviewed-by: Justin Tobler <jtobler@gitlab.com>
Co-authored-by: James Fargher <jfargher@gitlab.com>

5 files changed, 146 insertions, 9 deletions

diff --git a/cmd/gitaly-backup/main.go b/cmd/gitaly-backup/main.go
index 808c67a3a..eed6a6d63 100644
--- a/cmd/gitaly-backup/main.go
+++ b/cmd/gitaly-backup/main.go
@@ -6,6 +6,7 @@ import (
 	"io"
 	"os"
 
+	"gitlab.com/gitlab-org/gitaly/v15/internal/gitaly/storage"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/log"
 )
 
@@ -41,7 +42,12 @@ func main() {
 	subcmd.Flags(subcmdFlags)
 	_ = subcmdFlags.Parse(flags.Args()[2:])
 
-	if err := subcmd.Run(context.Background(), os.Stdin, os.Stdout); err != nil {
+	ctx, err := storage.InjectGitalyServersEnv(context.Background())
+	if err != nil {
+		logger.Fatalf("%s", err)
+	}
+
+	if err := subcmd.Run(ctx, os.Stdin, os.Stdout); err != nil {
 		logger.Fatalf("%s", err)
 	}
 }
diff --git a/internal/backup/backup.go b/internal/backup/backup.go
index fc3ca0ac8..11bb812fb 100644
--- a/internal/backup/backup.go
+++ b/internal/backup/backup.go
@@ -149,6 +149,10 @@ type CreateRequest struct {
 
 // Create creates a repository backup.
 func (mgr *Manager) Create(ctx context.Context, req *CreateRequest) error {
+	if err := setContextServerInfo(ctx, &req.Server, req.Repository.GetStorageName()); err != nil {
+		return fmt.Errorf("manager: %w", err)
+	}
+
 	if isEmpty, err := mgr.isEmpty(ctx, req.Server, req.Repository); err != nil {
 		return fmt.Errorf("manager: %w", err)
 	} else if isEmpty {
@@ -196,6 +200,10 @@ type RestoreRequest struct {
 
 // Restore restores a repository from a backup.
 func (mgr *Manager) Restore(ctx context.Context, req *RestoreRequest) error {
+	if err := setContextServerInfo(ctx, &req.Server, req.Repository.GetStorageName()); err != nil {
+		return fmt.Errorf("manager: %w", err)
+	}
+
 	if err := mgr.removeRepository(ctx, req.Server, req.Repository); err != nil {
 		return fmt.Errorf("manager: %w", err)
 	}
@@ -236,6 +244,17 @@ func (mgr *Manager) Restore(ctx context.Context, req *RestoreRequest) error {
 	return nil
 }
 
+// setContextServerInfo overwrites server with gitaly connection info from ctx metadata when server is zero.
+func setContextServerInfo(ctx context.Context, server *storage.ServerInfo, storageName string) error {
+	if !server.Zero() {
+		return nil
+	}
+
+	var err error
+	*server, err = storage.ExtractGitalyServer(ctx, storageName)
+	return err
+}
+
 func (mgr *Manager) isEmpty(ctx context.Context, server storage.ServerInfo, repo *gitalypb.Repository) (bool, error) {
 	repoClient, err := mgr.newRepoClient(ctx, server)
 	if err != nil {
diff --git a/internal/backup/backup_test.go b/internal/backup/backup_test.go
index 753af831e..50f25d93e 100644
--- a/internal/backup/backup_test.go
+++ b/internal/backup/backup_test.go
@@ -503,6 +503,40 @@ func testManagerRestore(t *testing.T, ctx context.Context) {
 	}
 }
 
+func TestManager_CreateRestore_contextServerInfo(t *testing.T) {
+	t.Parallel()
+
+	cfg := testcfg.Build(t)
+	testcfg.BuildGitalyHooks(t, cfg)
+	cfg.SocketPath = testserver.RunGitalyServer(t, cfg, nil, setup.RegisterAll)
+
+	ctx := testhelper.Context(t)
+
+	repo, repoPath := gittest.CreateRepository(t, ctx, cfg)
+	commitID := gittest.WriteCommit(t, cfg, repoPath, gittest.WithBranch("main"))
+	gittest.WriteTag(t, cfg, repoPath, "v1.0.0", commitID.Revision())
+
+	backupRoot := testhelper.TempDir(t)
+
+	pool := client.NewPool()
+	defer testhelper.MustClose(t, pool)
+
+	sink := NewFilesystemSink(backupRoot)
+	locator, err := ResolveLocator("pointer", sink)
+	require.NoError(t, err)
+
+	fsBackup := NewManager(sink, locator, pool, "unused-backup-id")
+
+	ctx = testhelper.MergeIncomingMetadata(ctx, testcfg.GitalyServersMetadataFromCfg(t, cfg))
+
+	require.NoError(t, fsBackup.Create(ctx, &CreateRequest{
+		Repository: repo,
+	}))
+	require.NoError(t, fsBackup.Restore(ctx, &RestoreRequest{
+		Repository: repo,
+	}))
+}
+
 func TestResolveSink(t *testing.T) {
 	ctx := testhelper.Context(t)
 
diff --git a/internal/gitaly/storage/servers.go b/internal/gitaly/storage/servers.go
index c9351fc87..68f710286 100644
--- a/internal/gitaly/storage/servers.go
+++ b/internal/gitaly/storage/servers.go
@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 
+	"gitlab.com/gitlab-org/gitaly/v15/internal/helper/env"
 	"google.golang.org/grpc/metadata"
 )
 
@@ -19,6 +20,11 @@ type ServerInfo struct {
 	Token   string `json:"token"`
 }
 
+// Zero returns true when no attributes have been set.
+func (si ServerInfo) Zero() bool {
+	return si == ServerInfo{}
+}
+
 // GitalyServers hold Gitaly servers info like {"default":{"token":"x","address":"y"}},
 // to be passed in `gitaly-servers` metadata.
 type GitalyServers map[string]ServerInfo
@@ -39,15 +45,9 @@ func ExtractGitalyServers(ctx context.Context) (gitalyServersInfo GitalyServers,
 		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)
+	if err := unmarshalGitalyServers(gitalyServersJSONEncoded[0], &gitalyServersInfo); err != nil {
+		return nil, err
 	}
-
 	return
 }
 
@@ -82,3 +82,33 @@ func InjectGitalyServers(ctx context.Context, name, address, token string) (cont
 
 	return metadata.AppendToOutgoingContext(ctx, "gitaly-servers", base64.StdEncoding.EncodeToString(gitalyServersJSON)), nil
 }
+
+// InjectGitalyServersEnv injects the GITALY_SERVERS env var into an incoming
+// context.
+func InjectGitalyServersEnv(ctx context.Context) (context.Context, error) {
+	rawServers := env.GetString("GITALY_SERVERS", "")
+	if rawServers == "" {
+		return ctx, nil
+	}
+
+	// Make sure we fail early if the value in the env var cannot be interpreted.
+	if err := unmarshalGitalyServers(rawServers, &GitalyServers{}); err != nil {
+		return nil, fmt.Errorf("injecting GITALY_SERVERS: %w", err)
+	}
+
+	md := metadata.Pairs("gitaly-servers", rawServers)
+	return metadata.NewIncomingContext(ctx, md), nil
+}
+
+func unmarshalGitalyServers(encoded string, servers *GitalyServers) error {
+	gitalyServersJSON, err := base64.StdEncoding.DecodeString(encoded)
+	if err != nil {
+		return fmt.Errorf("failed decoding base64: %v", err)
+	}
+
+	if err := json.Unmarshal(gitalyServersJSON, servers); err != nil {
+		return fmt.Errorf("failed unmarshalling json: %v", err)
+	}
+
+	return nil
+}
diff --git a/internal/gitaly/storage/servers_test.go b/internal/gitaly/storage/servers_test.go
index 6aa003447..be5e2d2e0 100644
--- a/internal/gitaly/storage/servers_test.go
+++ b/internal/gitaly/storage/servers_test.go
@@ -102,3 +102,51 @@ func TestInjectGitalyServers(t *testing.T) {
 		require.Equal(t, []string{"bar"}, md["foo"])
 	})
 }
+
+func TestInjectGitalyServersEnv(t *testing.T) {
+	t.Run("GITALY_SERVERS not set", func(t *testing.T) {
+		testhelper.Unsetenv(t, "GITALY_SERVERS")
+
+		ctx := testhelper.Context(t)
+
+		newCtx, err := storage.InjectGitalyServersEnv(ctx)
+
+		require.NoError(t, err)
+		require.Equal(t, ctx, newCtx)
+	})
+
+	for _, tc := range []struct {
+		desc             string
+		gitalyServersEnv string
+		expectedErr      string
+		expectedInfo     storage.GitalyServers
+	}{
+		{
+			desc:             "GITALY_SERVERS invalid",
+			gitalyServersEnv: base64.StdEncoding.EncodeToString([]byte("definitely not JSON")),
+			expectedErr:      "injecting GITALY_SERVERS: failed unmarshalling json: invalid character 'd' looking for beginning of value",
+		},
+		{
+			desc:             "GITALY_SERVERS set",
+			gitalyServersEnv: base64.StdEncoding.EncodeToString([]byte(`{"default":{"address":"unix:///tmp/sock","token":"hunter1"}}`)),
+			expectedInfo:     storage.GitalyServers{"default": storage.ServerInfo{Address: "unix:///tmp/sock", Token: "hunter1"}},
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			t.Setenv("GITALY_SERVERS", tc.gitalyServersEnv)
+
+			ctx, err := storage.InjectGitalyServersEnv(testhelper.Context(t))
+			if tc.expectedErr == "" {
+				require.NoError(t, err)
+			} else {
+				require.EqualError(t, err, tc.expectedErr)
+				return
+			}
+
+			info, err := storage.ExtractGitalyServers(ctx)
+			require.NoError(t, err)
+
+			require.Equal(t, tc.expectedInfo, info)
+		})
+	}
+}