Add support for Gitaly feature flags

gitaly#1846

Gitaly does not store or remember feature flags. We must pass them in as
metadata on each request. This MR adds plumbing to pass Gitaly feature flags
supplied by gitlab-rails to the Gitaly server as gRPC request metadata.

Signed-off-by: Mateusz Nowotyński <maxmati4@gmail.com>

11 files changed, 101 insertions, 13 deletions

diff --git a/go/internal/command/receivepack/gitalycall.go b/go/internal/command/receivepack/gitalycall.go
index d735f17..60954d9 100644
--- a/go/internal/command/receivepack/gitalycall.go
+++ b/go/internal/command/receivepack/gitalycall.go
@@ -18,6 +18,7 @@ func (c *Command) performGitalyCall(response *accessverifier.Response) error {
 		ServiceName: string(commandargs.ReceivePack),
 		Address:     response.Gitaly.Address,
 		Token:       response.Gitaly.Token,
+		Features:    response.Gitaly.Features,
 	}
 
 	request := &pb.SSHReceivePackRequest{
diff --git a/go/internal/command/receivepack/gitalycall_test.go b/go/internal/command/receivepack/gitalycall_test.go
index eac9218..a27b8a2 100644
--- a/go/internal/command/receivepack/gitalycall_test.go
+++ b/go/internal/command/receivepack/gitalycall_test.go
@@ -14,7 +14,7 @@ import (
 )
 
 func TestReceivePack(t *testing.T) {
-	gitalyAddress, cleanup := testserver.StartGitalyServer(t)
+	gitalyAddress, _, cleanup := testserver.StartGitalyServer(t)
 	defer cleanup()
 
 	requests := requesthandlers.BuildAllowedWithGitalyHandlers(t, gitalyAddress)
diff --git a/go/internal/command/uploadarchive/gitalycall.go b/go/internal/command/uploadarchive/gitalycall.go
index e810ba3..a2bdd84 100644
--- a/go/internal/command/uploadarchive/gitalycall.go
+++ b/go/internal/command/uploadarchive/gitalycall.go
@@ -18,6 +18,7 @@ func (c *Command) performGitalyCall(response *accessverifier.Response) error {
 		ServiceName: string(commandargs.UploadArchive),
 		Address:     response.Gitaly.Address,
 		Token:       response.Gitaly.Token,
+		Features:    response.Gitaly.Features,
 	}
 
 	request := &pb.SSHUploadArchiveRequest{Repository: &response.Gitaly.Repo}
diff --git a/go/internal/command/uploadarchive/gitalycall_test.go b/go/internal/command/uploadarchive/gitalycall_test.go
index 5eb2eae..a95985a 100644
--- a/go/internal/command/uploadarchive/gitalycall_test.go
+++ b/go/internal/command/uploadarchive/gitalycall_test.go
@@ -14,7 +14,7 @@ import (
 )
 
 func TestUploadPack(t *testing.T) {
-	gitalyAddress, cleanup := testserver.StartGitalyServer(t)
+	gitalyAddress, _, cleanup := testserver.StartGitalyServer(t)
 	defer cleanup()
 
 	requests := requesthandlers.BuildAllowedWithGitalyHandlers(t, gitalyAddress)
diff --git a/go/internal/command/uploadpack/gitalycall.go b/go/internal/command/uploadpack/gitalycall.go
index 5dff24a..84f8f32 100644
--- a/go/internal/command/uploadpack/gitalycall.go
+++ b/go/internal/command/uploadpack/gitalycall.go
@@ -18,6 +18,7 @@ func (c *Command) performGitalyCall(response *accessverifier.Response) error {
 		ServiceName: string(commandargs.UploadPack),
 		Address:     response.Gitaly.Address,
 		Token:       response.Gitaly.Token,
+		Features:    response.Gitaly.Features,
 	}
 
 	request := &pb.SSHUploadPackRequest{
diff --git a/go/internal/command/uploadpack/gitalycall_test.go b/go/internal/command/uploadpack/gitalycall_test.go
index eb18aa8..091aca9 100644
--- a/go/internal/command/uploadpack/gitalycall_test.go
+++ b/go/internal/command/uploadpack/gitalycall_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/commandargs"
@@ -14,7 +15,7 @@ import (
 )
 
 func TestUploadPack(t *testing.T) {
-	gitalyAddress, cleanup := testserver.StartGitalyServer(t)
+	gitalyAddress, testServer, cleanup := testserver.StartGitalyServer(t)
 	defer cleanup()
 
 	requests := requesthandlers.BuildAllowedWithGitalyHandlers(t, gitalyAddress)
@@ -37,4 +38,14 @@ func TestUploadPack(t *testing.T) {
 	require.NoError(t, err)
 
 	require.Equal(t, "UploadPack: "+repo, output.String())
+
+	for k, v := range map[string]string{
+		"gitaly-feature-cache_invalidator":        "true",
+		"gitaly-feature-inforef_uploadpack_cache": "false",
+	} {
+		actual := testServer.ReceivedMD[k]
+		assert.Len(t, actual, 1)
+		assert.Equal(t, v, actual[0])
+	}
+	assert.Empty(t, testServer.ReceivedMD["some-other-ff"])
 }
diff --git a/go/internal/gitlabnet/accessverifier/client.go b/go/internal/gitlabnet/accessverifier/client.go
index eb67703..2f19500 100644
--- a/go/internal/gitlabnet/accessverifier/client.go
+++ b/go/internal/gitlabnet/accessverifier/client.go
@@ -31,9 +31,10 @@ type Request struct {
 }
 
 type Gitaly struct {
-	Repo    pb.Repository `json:"repository"`
-	Address string        `json:"address"`
-	Token   string        `json:"token"`
+	Repo     pb.Repository     `json:"repository"`
+	Address  string            `json:"address"`
+	Token    string            `json:"token"`
+	Features map[string]string `json:"features"`
 }
 
 type CustomPayloadData struct {
diff --git a/go/internal/gitlabnet/testserver/gitalyserver.go b/go/internal/gitlabnet/testserver/gitalyserver.go
index 694fd41..6d0c130 100644
--- a/go/internal/gitlabnet/testserver/gitalyserver.go
+++ b/go/internal/gitlabnet/testserver/gitalyserver.go
@@ -10,49 +10,56 @@ import (
 
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
 
 	pb "gitlab.com/gitlab-org/gitaly/proto/go/gitalypb"
 )
 
-type testGitalyServer struct{}
+type TestGitalyServer struct{ ReceivedMD metadata.MD }
 
-func (s *testGitalyServer) SSHReceivePack(stream pb.SSHService_SSHReceivePackServer) error {
+func (s *TestGitalyServer) SSHReceivePack(stream pb.SSHService_SSHReceivePackServer) error {
 	req, err := stream.Recv()
 	if err != nil {
 		return err
 	}
 
+	s.ReceivedMD, _ = metadata.FromIncomingContext(stream.Context())
+
 	response := []byte("ReceivePack: " + req.GlId + " " + req.Repository.GlRepository)
 	stream.Send(&pb.SSHReceivePackResponse{Stdout: response})
 
 	return nil
 }
 
-func (s *testGitalyServer) SSHUploadPack(stream pb.SSHService_SSHUploadPackServer) error {
+func (s *TestGitalyServer) SSHUploadPack(stream pb.SSHService_SSHUploadPackServer) error {
 	req, err := stream.Recv()
 	if err != nil {
 		return err
 	}
 
+	s.ReceivedMD, _ = metadata.FromIncomingContext(stream.Context())
+
 	response := []byte("UploadPack: " + req.Repository.GlRepository)
 	stream.Send(&pb.SSHUploadPackResponse{Stdout: response})
 
 	return nil
 }
 
-func (s *testGitalyServer) SSHUploadArchive(stream pb.SSHService_SSHUploadArchiveServer) error {
+func (s *TestGitalyServer) SSHUploadArchive(stream pb.SSHService_SSHUploadArchiveServer) error {
 	req, err := stream.Recv()
 	if err != nil {
 		return err
 	}
 
+	s.ReceivedMD, _ = metadata.FromIncomingContext(stream.Context())
+
 	response := []byte("UploadArchive: " + req.Repository.GlRepository)
 	stream.Send(&pb.SSHUploadArchiveResponse{Stdout: response})
 
 	return nil
 }
 
-func StartGitalyServer(t *testing.T) (string, func()) {
+func StartGitalyServer(t *testing.T) (string, *TestGitalyServer, func()) {
 	tempDir, _ := ioutil.TempDir("", "gitlab-shell-test-api")
 	gitalySocketPath := path.Join(tempDir, "gitaly.sock")
 
@@ -64,7 +71,8 @@ func StartGitalyServer(t *testing.T) (string, func()) {
 	listener, err := net.Listen("unix", gitalySocketPath)
 	require.NoError(t, err)
 
-	pb.RegisterSSHServiceServer(server, &testGitalyServer{})
+	testServer := TestGitalyServer{}
+	pb.RegisterSSHServiceServer(server, &testServer)
 
 	go server.Serve(listener)
 
@@ -74,5 +82,5 @@ func StartGitalyServer(t *testing.T) (string, func()) {
 		os.RemoveAll(tempDir)
 	}
 
-	return gitalySocketUrl, cleanup
+	return gitalySocketUrl, &testServer, cleanup
 }
diff --git a/go/internal/handler/exec.go b/go/internal/handler/exec.go
index 1f3177d..277d08c 100644
--- a/go/internal/handler/exec.go
+++ b/go/internal/handler/exec.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"strings"
 
 	"gitlab.com/gitlab-org/gitaly/auth"
 	"gitlab.com/gitlab-org/gitaly/client"
@@ -11,6 +12,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
 	"gitlab.com/gitlab-org/labkit/tracing"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
 )
 
 // GitalyHandlerFunc implementations are responsible for making
@@ -29,6 +31,7 @@ type GitalyCommand struct {
 	ServiceName string
 	Address     string
 	Token       string
+	Features    map[string]string
 }
 
 // RunGitalyCommand provides a bootstrap for Gitaly commands executed
@@ -48,6 +51,18 @@ func (gc *GitalyCommand) RunGitalyCommand(handler GitalyHandlerFunc) error {
 	return err
 }
 
+func withOutgoingMetadata(ctx context.Context, features map[string]string) context.Context {
+	md := metadata.New(nil)
+	for k, v := range features {
+		if !strings.HasPrefix(k, "gitaly-feature-") {
+			continue
+		}
+		md.Append(k, v)
+	}
+
+	return metadata.NewOutgoingContext(ctx, md)
+}
+
 func getConn(gc *GitalyCommand) (*GitalyConn, error) {
 	if gc.Address == "" {
 		return nil, fmt.Errorf("no gitaly_address given")
@@ -80,6 +95,7 @@ func getConn(gc *GitalyCommand) (*GitalyConn, error) {
 	)
 
 	ctx, finished := tracing.ExtractFromEnv(context.Background())
+	ctx = withOutgoingMetadata(ctx, gc.Features)
 
 	conn, err := client.Dial(gc.Address, connOpts)
 	if err != nil {
diff --git a/go/internal/handler/exec_test.go b/go/internal/handler/exec_test.go
index e6608ac..97c9f05 100644
--- a/go/internal/handler/exec_test.go
+++ b/go/internal/handler/exec_test.go
@@ -5,8 +5,10 @@ import (
 	"errors"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
 
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
 )
@@ -40,3 +42,45 @@ func TestMissingGitalyAddress(t *testing.T) {
 	err := cmd.RunGitalyCommand(makeHandler(t, nil))
 	require.EqualError(t, err, "no gitaly_address given")
 }
+
+func TestGetConnMetadata(t *testing.T) {
+	tests := []struct {
+		name string
+		gc   *GitalyCommand
+		want map[string]string
+	}{
+		{
+			name: "gitaly_feature_flags",
+			gc: &GitalyCommand{
+				Config:  &config.Config{},
+				Address: "tcp://localhost:9999",
+				Features: map[string]string{
+					"gitaly-feature-cache_invalidator":        "true",
+					"other-ff":                                "true",
+					"gitaly-feature-inforef_uploadpack_cache": "false",
+				},
+			},
+			want: map[string]string{
+				"gitaly-feature-cache_invalidator":        "true",
+				"gitaly-feature-inforef_uploadpack_cache": "false",
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			conn, err := getConn(tt.gc)
+			require.NoError(t, err)
+
+			md, exists := metadata.FromOutgoingContext(conn.ctx)
+			require.True(t, exists)
+			require.Equal(t, len(tt.want), md.Len())
+
+			for k, v := range tt.want {
+				values := md.Get(k)
+				assert.Equal(t, 1, len(values))
+				assert.Equal(t, v, values[0])
+			}
+
+		})
+	}
+}
diff --git a/go/internal/testhelper/requesthandlers/requesthandlers.go b/go/internal/testhelper/requesthandlers/requesthandlers.go
index a7bc427..2931209 100644
--- a/go/internal/testhelper/requesthandlers/requesthandlers.go
+++ b/go/internal/testhelper/requesthandlers/requesthandlers.go
@@ -47,6 +47,11 @@ func BuildAllowedWithGitalyHandlers(t *testing.T, gitalyAddress string) []testse
 						},
 						"address": gitalyAddress,
 						"token":   "token",
+						"features": map[string]string{
+							"gitaly-feature-cache_invalidator":        "true",
+							"gitaly-feature-inforef_uploadpack_cache": "false",
+							"some-other-ff":                           "true",
+						},
 					},
 				}
 				require.NoError(t, json.NewEncoder(w).Encode(body))