Merge branch 'po-refactor-ff-helper' into 'master'

Refactor context feature flag helpers

See merge request gitlab-org/gitaly!1802

12 files changed, 76 insertions, 40 deletions

diff --git a/cmd/gitaly-ssh/receive_pack.go b/cmd/gitaly-ssh/receive_pack.go
index a91bc4c91..b0f44326b 100644
--- a/cmd/gitaly-ssh/receive_pack.go
+++ b/cmd/gitaly-ssh/receive_pack.go
@@ -23,7 +23,7 @@ func receivePack(ctx context.Context, conn *grpc.ClientConn, req string) (int32,
 	defer cancel()
 
 	if request.GetGitProtocol() == git.ProtocolV2 {
-		ctx = featureflag.ContextWithFeatureFlag(ctx, featureflag.UseGitProtocolV2)
+		ctx = featureflag.OutgoingCtxWithFeatureFlag(ctx, featureflag.UseGitProtocolV2)
 	}
 
 	return client.ReceivePack(ctx, conn, os.Stdin, os.Stdout, os.Stderr, &request)
diff --git a/cmd/gitaly-ssh/upload_pack.go b/cmd/gitaly-ssh/upload_pack.go
index d6ef54687..37a62f64b 100644
--- a/cmd/gitaly-ssh/upload_pack.go
+++ b/cmd/gitaly-ssh/upload_pack.go
@@ -32,7 +32,7 @@ func uploadPack(ctx context.Context, conn *grpc.ClientConn, req string) (int32,
 	defer cancel()
 
 	if request.GetGitProtocol() == git.ProtocolV2 {
-		ctx = featureflag.ContextWithFeatureFlag(ctx, featureflag.UseGitProtocolV2)
+		ctx = featureflag.OutgoingCtxWithFeatureFlag(ctx, featureflag.UseGitProtocolV2)
 	}
 
 	return client.UploadPack(ctx, conn, os.Stdin, os.Stdout, os.Stderr, &request)
diff --git a/internal/git/log/commit_test.go b/internal/git/log/commit_test.go
index 8ca3c1ec4..5eff5400d 100644
--- a/internal/git/log/commit_test.go
+++ b/internal/git/log/commit_test.go
@@ -165,7 +165,7 @@ func TestGetCommitCatfile(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
 			if tc.featureOn {
-				ctx = featureflag.ContextWithFeatureFlag(ctx, featureflag.CommitWithoutBatchCheck)
+				ctx = featureflag.OutgoingCtxWithFeatureFlag(ctx, featureflag.CommitWithoutBatchCheck)
 			}
 
 			c, err := GetCommitCatfile(ctx, c, tc.revision)
diff --git a/internal/git/protocol_test.go b/internal/git/protocol_test.go
index 99993a135..18583b381 100644
--- a/internal/git/protocol_test.go
+++ b/internal/git/protocol_test.go
@@ -6,7 +6,6 @@ import (
 
 	"github.com/stretchr/testify/require"
 	"gitlab.com/gitlab-org/gitaly/internal/metadata/featureflag"
-	"google.golang.org/grpc/metadata"
 )
 
 type fakeProtocolMessage struct {
@@ -18,17 +17,7 @@ func (f fakeProtocolMessage) GetGitProtocol() string {
 }
 
 func enableProtocolV2(ctx context.Context) context.Context {
-	// TODO: replace this implementation with a helper function that deals
-	// with the incoming context
-	flag := featureflag.UseGitProtocolV2
-
-	md, ok := metadata.FromIncomingContext(ctx)
-	if !ok {
-		md = metadata.New(map[string]string{featureflag.HeaderKey(flag): "true"})
-	}
-	md.Set(featureflag.HeaderKey(flag), "true")
-
-	return metadata.NewIncomingContext(ctx, md)
+	return featureflag.IncomingCtxWithFeatureFlag(ctx, featureflag.UseGitProtocolV2)
 }
 
 func TestAddGitProtocolEnv(t *testing.T) {
diff --git a/internal/metadata/featureflag/context.go b/internal/metadata/featureflag/context.go
new file mode 100644
index 000000000..2068683cc
--- /dev/null
+++ b/internal/metadata/featureflag/context.go
@@ -0,0 +1,31 @@
+package featureflag
+
+import (
+	"context"
+
+	"google.golang.org/grpc/metadata"
+)
+
+// OutgoingCtxWithFeatureFlag is used to enable a feature flag in the outgoing
+// context metadata. The returned context is meant to be used in a client where
+// the outcoming context is transferred to an incoming context.
+func OutgoingCtxWithFeatureFlag(ctx context.Context, flag string) context.Context {
+	md, ok := metadata.FromOutgoingContext(ctx)
+	if !ok {
+		md = metadata.New(map[string]string{})
+	}
+	md.Set(HeaderKey(flag), "true")
+	return metadata.NewOutgoingContext(ctx, md)
+}
+
+// IncomingCtxWithFeatureFlag is used to enable a feature flag in the incoming
+// context. This is NOT meant for use in clients that transfer the context
+// across process boundaries.
+func IncomingCtxWithFeatureFlag(ctx context.Context, flag string) context.Context {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		md = metadata.New(map[string]string{})
+	}
+	md.Set(HeaderKey(flag), "true")
+	return metadata.NewIncomingContext(ctx, md)
+}
diff --git a/internal/metadata/featureflag/context_test.go b/internal/metadata/featureflag/context_test.go
new file mode 100644
index 000000000..3422aae59
--- /dev/null
+++ b/internal/metadata/featureflag/context_test.go
@@ -0,0 +1,33 @@
+package featureflag
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/metadata"
+)
+
+func TestIncomingCtxWithFeatureFlag(t *testing.T) {
+	ctx := context.Background()
+	require.False(t, IsEnabled(ctx, UseGitProtocolV2))
+
+	ctx = IncomingCtxWithFeatureFlag(ctx, UseGitProtocolV2)
+	require.True(t, IsEnabled(ctx, UseGitProtocolV2))
+}
+
+func TestOutgoingCtxWithFeatureFlag(t *testing.T) {
+	ctx := context.Background()
+	require.False(t, IsEnabled(ctx, UseGitProtocolV2))
+
+	ctx = OutgoingCtxWithFeatureFlag(ctx, UseGitProtocolV2)
+	require.False(t, IsEnabled(ctx, UseGitProtocolV2))
+
+	// simulate an outgoing context leaving the process boundary and then
+	// becoming an incoming context in a new process boundary
+	md, ok := metadata.FromOutgoingContext(ctx)
+	require.True(t, ok)
+
+	ctx = metadata.NewIncomingContext(context.Background(), md)
+	require.True(t, IsEnabled(ctx, UseGitProtocolV2))
+}
diff --git a/internal/metadata/featureflag/test_utils.go b/internal/metadata/featureflag/test_utils.go
deleted file mode 100644
index 8131b82a4..000000000
--- a/internal/metadata/featureflag/test_utils.go
+++ /dev/null
@@ -1,19 +0,0 @@
-package featureflag
-
-import (
-	"context"
-
-	"google.golang.org/grpc/metadata"
-)
-
-// ContextWithFeatureFlag is used in tests to enable a feature flag in the context metadata
-func ContextWithFeatureFlag(ctx context.Context, flag string) context.Context {
-	md, ok := metadata.FromIncomingContext(ctx)
-	if !ok {
-		md = metadata.New(map[string]string{HeaderKey(flag): "true"})
-	} else {
-		md.Set(HeaderKey(flag), "true")
-	}
-
-	return metadata.NewOutgoingContext(ctx, md)
-}
diff --git a/internal/middleware/cache/cache_test.go b/internal/middleware/cache/cache_test.go
index b4aa56307..0d9d126de 100644
--- a/internal/middleware/cache/cache_test.go
+++ b/internal/middleware/cache/cache_test.go
@@ -46,7 +46,7 @@ func TestInvalidators(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
 	defer cancel()
 
-	ctx = featureflag.ContextWithFeatureFlag(ctx, featureflag.CacheInvalidator)
+	ctx = featureflag.OutgoingCtxWithFeatureFlag(ctx, featureflag.CacheInvalidator)
 
 	svc := &testSvc{}
 
diff --git a/internal/service/repository/archive_test.go b/internal/service/repository/archive_test.go
index a5afbc1f1..0743c484d 100644
--- a/internal/service/repository/archive_test.go
+++ b/internal/service/repository/archive_test.go
@@ -214,7 +214,9 @@ func TestGetArchiveFailure(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
 			ctx, cancel := testhelper.Context()
-			ctx = featureflag.ContextWithFeatureFlag(ctx, featureflag.CacheInvalidator)
+			// enable cache invalidator to test against bad inputs
+			// See https://gitlab.com/gitlab-org/gitaly/issues/2262
+			ctx = featureflag.OutgoingCtxWithFeatureFlag(ctx, featureflag.CacheInvalidator)
 			defer cancel()
 
 			req := &gitalypb.GetArchiveRequest{
diff --git a/internal/service/smarthttp/inforefs_test.go b/internal/service/smarthttp/inforefs_test.go
index d633da2cf..64ebbe6a9 100644
--- a/internal/service/smarthttp/inforefs_test.go
+++ b/internal/service/smarthttp/inforefs_test.go
@@ -80,7 +80,7 @@ func TestSuccessfulInfoRefsUploadPackWithGitProtocol(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	ctx = featureflag.ContextWithFeatureFlag(ctx, featureflag.UseGitProtocolV2)
+	ctx = featureflag.OutgoingCtxWithFeatureFlag(ctx, featureflag.UseGitProtocolV2)
 
 	c, err := client.InfoRefsUploadPack(ctx, rpcRequest)
 
diff --git a/internal/service/smarthttp/receive_pack_test.go b/internal/service/smarthttp/receive_pack_test.go
index 7f0eb50ed..46b874034 100644
--- a/internal/service/smarthttp/receive_pack_test.go
+++ b/internal/service/smarthttp/receive_pack_test.go
@@ -85,7 +85,7 @@ func TestSuccessfulReceivePackRequestWithGitProtocol(t *testing.T) {
 	ctx, cancel := testhelper.Context()
 	defer cancel()
 
-	ctx = featureflag.ContextWithFeatureFlag(ctx, featureflag.UseGitProtocolV2)
+	ctx = featureflag.OutgoingCtxWithFeatureFlag(ctx, featureflag.UseGitProtocolV2)
 
 	stream, err := client.PostReceivePack(ctx)
 	require.NoError(t, err)
diff --git a/internal/service/smarthttp/upload_pack_test.go b/internal/service/smarthttp/upload_pack_test.go
index 05ffbd7a4..a7777e595 100644
--- a/internal/service/smarthttp/upload_pack_test.go
+++ b/internal/service/smarthttp/upload_pack_test.go
@@ -173,7 +173,7 @@ func TestUploadPackRequestWithGitProtocol(t *testing.T) {
 	ctx, cancel := testhelper.Context()
 	defer cancel()
 
-	ctx = featureflag.ContextWithFeatureFlag(ctx, featureflag.UseGitProtocolV2)
+	ctx = featureflag.OutgoingCtxWithFeatureFlag(ctx, featureflag.UseGitProtocolV2)
 
 	testRepo := testhelper.TestRepository()
 	testRepoPath, err := helper.GetRepoPath(testRepo)
@@ -384,7 +384,7 @@ func TestUploadPackRequestForPartialCloneSuccess(t *testing.T) {
 	_, err = makePostUploadPackRequest(ctx, t, serverSocketPath, req, &requestBufferForFailed)
 	require.Error(t, err, "trying to use filters without the feature flag should result in an error")
 
-	ctx = featureflag.ContextWithFeatureFlag(ctx, featureflag.UploadPackFilter)
+	ctx = featureflag.OutgoingCtxWithFeatureFlag(ctx, featureflag.UploadPackFilter)
 
 	responseBuffer, err := makePostUploadPackRequest(ctx, t, serverSocketPath, req, &requestBuffer)
 	require.NoError(t, err)