storage: Extract InjectGitalyServersEnv into storage package

Colocates this function with other functions that manipulate this value
on context.

2 files changed, 68 insertions, 0 deletions

diff --git a/internal/gitaly/storage/servers.go b/internal/gitaly/storage/servers.go
index f4f2408c5..f621925ef 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"
 )
 
@@ -87,3 +88,22 @@ 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
+	}
+
+	md := metadata.Pairs("gitaly-servers", rawServers)
+	ctx = metadata.NewIncomingContext(ctx, md)
+
+	// Make sure we fail early if the value in the env var cannot be interpreted.
+	if _, err := ExtractGitalyServers(ctx); err != nil {
+		return nil, fmt.Errorf("injecting GITALY_SERVERS: %w", err)
+	}
+
+	return ctx, 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)
+		})
+	}
+}