Merge branch 'pass-ff-to-gitaly' into 'master'

Add support for Gitaly feature flags

See merge request gitlab-org/gitlab-shell!351

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{

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)

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}

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)

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{

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/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"])
 }

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 {

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
 }

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/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 {

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/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])
+			}
+
+		})
+	}
+}

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))