1package storage_test
2
3import (
4	"context"
5	"encoding/base64"
6	"encoding/json"
7	"strings"
8	"testing"
9
10	"github.com/stretchr/testify/require"
11	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/storage"
12	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper"
13	"google.golang.org/grpc/metadata"
14)
15
16func TestExtractGitalyServers(t *testing.T) {
17	ctxOuter, cancel := testhelper.Context()
18	defer cancel()
19
20	testCases := []struct {
21		desc     string
22		metadata metadata.MD
23		info     storage.GitalyServers
24	}{
25		{
26			desc:     "no gitaly-servers metadata",
27			metadata: metadata.Pairs("not-gitaly-servers", "definitely not JSON camouflaging in base64"),
28		},
29		{
30			desc:     "metadata not encoded in base64",
31			metadata: metadata.Pairs("gitaly-servers", "definitely not base64"),
32		},
33		{
34			desc:     "encoded metadata is not JSON",
35			metadata: metadata.Pairs("gitaly-servers", base64.StdEncoding.EncodeToString([]byte("definitely not JSON"))),
36		},
37		{
38			desc:     "encoded JSON is not of the expected format",
39			metadata: metadata.Pairs("gitaly-servers", base64.StdEncoding.EncodeToString([]byte(`{"default":"string"}`))),
40		},
41		{
42			desc:     "properly-encoded string",
43			metadata: metadata.Pairs("gitaly-servers", base64.StdEncoding.EncodeToString([]byte(`{"default":{"address":"unix:///tmp/sock","token":"hunter1"}}`))),
44			info:     storage.GitalyServers{"default": storage.ServerInfo{Address: "unix:///tmp/sock", Token: "hunter1"}},
45		},
46	}
47
48	for _, testCase := range testCases {
49		t.Run(testCase.desc, func(t *testing.T) {
50			ctx := metadata.NewIncomingContext(ctxOuter, testCase.metadata)
51
52			info, err := storage.ExtractGitalyServers(ctx)
53			if testCase.info == nil {
54				require.Error(t, err)
55			} else {
56				require.NoError(t, err)
57				require.Equal(t, testCase.info, info)
58			}
59		})
60	}
61}
62
63func TestInjectGitalyServers(t *testing.T) {
64	check := func(t *testing.T, ctx context.Context) {
65		t.Helper()
66
67		newCtx, err := storage.InjectGitalyServers(ctx, "gitaly-1", "1.1.1.1", "secret")
68		require.NoError(t, err)
69
70		md, found := metadata.FromOutgoingContext(newCtx)
71		require.True(t, found)
72
73		gs, found := md["gitaly-servers"]
74		require.True(t, found)
75
76		require.Len(t, gs, 1)
77
78		var servers map[string]interface{}
79		require.NoError(t, json.NewDecoder(base64.NewDecoder(base64.StdEncoding, strings.NewReader(gs[0]))).Decode(&servers), "received %s", gs[0])
80		require.EqualValues(t, map[string]interface{}{"gitaly-1": map[string]interface{}{"address": "1.1.1.1", "token": "secret"}}, servers)
81	}
82
83	t.Run("brand new context", func(t *testing.T) {
84		ctx := context.Background()
85
86		check(t, ctx)
87	})
88
89	t.Run("context with existing outgoing metadata should not be re-written", func(t *testing.T) {
90		existing := metadata.New(map[string]string{"foo": "bar"})
91
92		ctx := metadata.NewOutgoingContext(context.Background(), existing)
93		check(t, ctx)
94
95		md, found := metadata.FromOutgoingContext(ctx)
96		require.True(t, found)
97		require.Equal(t, []string{"bar"}, md["foo"])
98	})
99}
100