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