From 8542d85d27d3c8fb381179b235a04a10ece6a4c2 Mon Sep 17 00:00:00 2001
From: "Jacob Vosmaer (GitLab)" <jacob@gitlab.com>
Date: Mon, 13 Nov 2017 15:42:06 +0000
Subject: Implement RepositoryService::FetchSourceBranch

---
 internal/rubyserver/proxy.go           | 19 ++++++++++++++
 internal/rubyserver/proxy_test.go      | 46 ++++++++++++++++++++++++++++++++++
 internal/rubyserver/rubyserver.go      |  7 +-----
 internal/rubyserver/rubyserver_test.go |  2 --
 internal/rubyserver/testhelper_test.go | 22 ++++++++++++++++
 5 files changed, 88 insertions(+), 8 deletions(-)
 create mode 100644 internal/rubyserver/proxy_test.go
 create mode 100644 internal/rubyserver/testhelper_test.go

(limited to 'internal/rubyserver')

diff --git a/internal/rubyserver/proxy.go b/internal/rubyserver/proxy.go
index 503d9ac8f..9cb8289e8 100644
--- a/internal/rubyserver/proxy.go
+++ b/internal/rubyserver/proxy.go
@@ -12,6 +12,16 @@ import (
 	"google.golang.org/grpc/metadata"
 )
 
+// ProxyHeaderWhitelist is the list of http/2 headers that will be
+// forwarded as-is to gitaly-ruby.
+var ProxyHeaderWhitelist = []string{"gitaly-servers"}
+
+const (
+	repoPathHeader     = "gitaly-repo-path"
+	glRepositoryHeader = "gitaly-gl-repository"
+	repoAltDirsHeader  = "gitaly-repo-alt-dirs"
+)
+
 // SetHeaders adds headers that tell gitaly-ruby the full path to the repository.
 func SetHeaders(ctx context.Context, repo *pb.Repository) (context.Context, error) {
 	repoPath, err := helper.GetPath(repo)
@@ -28,6 +38,15 @@ func SetHeaders(ctx context.Context, repo *pb.Repository) (context.Context, erro
 		glRepositoryHeader, repo.GlRepository,
 		repoAltDirsHeader, repoAltDirsCombined,
 	)
+
+	if inMD, ok := metadata.FromIncomingContext(ctx); ok {
+		for _, header := range ProxyHeaderWhitelist {
+			for _, v := range inMD[header] {
+				md = metadata.Join(md, metadata.Pairs(header, v))
+			}
+		}
+	}
+
 	newCtx := metadata.NewOutgoingContext(ctx, md)
 	return newCtx, nil
 }
diff --git a/internal/rubyserver/proxy_test.go b/internal/rubyserver/proxy_test.go
new file mode 100644
index 000000000..397cafe09
--- /dev/null
+++ b/internal/rubyserver/proxy_test.go
@@ -0,0 +1,46 @@
+package rubyserver
+
+import (
+	"testing"
+
+	"gitlab.com/gitlab-org/gitaly/internal/testhelper"
+
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/metadata"
+)
+
+func TestSetHeadersBlocksUnknownMetadata(t *testing.T) {
+	ctx, cancel := testhelper.Context()
+	defer cancel()
+
+	otherKey := "unknown-key"
+	otherValue := "test-value"
+	inCtx := metadata.NewIncomingContext(ctx, metadata.Pairs(otherKey, otherValue))
+
+	outCtx, err := SetHeaders(inCtx, testRepo)
+	require.NoError(t, err)
+
+	outMd, ok := metadata.FromOutgoingContext(outCtx)
+	require.True(t, ok, "outgoing context should have metadata")
+
+	_, ok = outMd[otherKey]
+	require.False(t, ok, "outgoing MD should not contain non-whitelisted key")
+}
+
+func TestSetHeadersPreservesWhitelistedMetadata(t *testing.T) {
+	ctx, cancel := testhelper.Context()
+	defer cancel()
+
+	key := "gitaly-servers"
+	require.Contains(t, ProxyHeaderWhitelist, key, "sanity check")
+	value := "test-value"
+	inCtx := metadata.NewIncomingContext(ctx, metadata.Pairs(key, value))
+
+	outCtx, err := SetHeaders(inCtx, testRepo)
+	require.NoError(t, err)
+
+	outMd, ok := metadata.FromOutgoingContext(outCtx)
+	require.True(t, ok, "outgoing context should have metadata")
+
+	require.Equal(t, []string{value}, outMd[key], "outgoing MD should contain whitelisted key")
+}
diff --git a/internal/rubyserver/rubyserver.go b/internal/rubyserver/rubyserver.go
index 9d0466115..3615ddeeb 100644
--- a/internal/rubyserver/rubyserver.go
+++ b/internal/rubyserver/rubyserver.go
@@ -23,12 +23,6 @@ import (
 	"google.golang.org/grpc"
 )
 
-const (
-	repoPathHeader     = "gitaly-repo-path"
-	glRepositoryHeader = "gitaly-gl-repository"
-	repoAltDirsHeader  = "gitaly-repo-alt-dirs"
-)
-
 var (
 	socketDir string
 
@@ -101,6 +95,7 @@ func Start() (*Server, error) {
 		"GITALY_RUBY_GIT_BIN_PATH=" + command.GitPath(),
 		fmt.Sprintf("GITALY_RUBY_WRITE_BUFFER_SIZE=%d", streamio.WriteBufferSize),
 		"GITALY_RUBY_GITLAB_SHELL_PATH=" + cfg.GitlabShell.Dir,
+		"GITALY_RUBY_GITALY_BIN_DIR=" + cfg.BinDir,
 	}
 
 	args := []string{"bundle", "exec", "bin/gitaly-ruby", fmt.Sprintf("%d", os.Getpid()), socketPath()}
diff --git a/internal/rubyserver/rubyserver_test.go b/internal/rubyserver/rubyserver_test.go
index c39782d3b..77943731e 100644
--- a/internal/rubyserver/rubyserver_test.go
+++ b/internal/rubyserver/rubyserver_test.go
@@ -23,8 +23,6 @@ func TestStopSafe(t *testing.T) {
 }
 
 func TestSetHeaders(t *testing.T) {
-	testRepo := testhelper.TestRepository()
-
 	testCases := []struct {
 		repo    *pb.Repository
 		errType codes.Code
diff --git a/internal/rubyserver/testhelper_test.go b/internal/rubyserver/testhelper_test.go
new file mode 100644
index 000000000..cc83a7923
--- /dev/null
+++ b/internal/rubyserver/testhelper_test.go
@@ -0,0 +1,22 @@
+package rubyserver
+
+import (
+	"os"
+	"testing"
+
+	"gitlab.com/gitlab-org/gitaly/internal/testhelper"
+)
+
+var (
+	testRepo = testhelper.TestRepository()
+)
+
+func TestMain(m *testing.M) {
+	os.Exit(testMain(m))
+}
+
+func testMain(m *testing.M) int {
+	defer testhelper.MustHaveNoChildProcess()
+
+	return m.Run()
+}
-- 
cgit v1.2.3