Add latest changes from gitlab-org/gitlab@14-7-stable-eev14.7.0-rc42

6 files changed, 76 insertions, 16 deletions

diff --git a/workhorse/.gitignore b/workhorse/.gitignore
index 7d339fef482..97a27630a9c 100644
--- a/workhorse/.gitignore
+++ b/workhorse/.gitignore
@@ -9,3 +9,4 @@ testdata/alt-public
 /_build
 coverage.html
 /*.toml
+/gitaly.pid
diff --git a/workhorse/Makefile b/workhorse/Makefile
index 3cf592b0cff..031fe581d28 100644
--- a/workhorse/Makefile
+++ b/workhorse/Makefile
@@ -65,7 +65,12 @@ install: $(EXE_ALL)
 .PHONY:	test
 test: prepare-tests
 	$(call message,$@)
-	go test -tags "$(BUILD_TAGS)" ./... ;\
+	@if [ -z "$${GITALY_ADDRESS+x}" ] ; then \
+		echo "To run gitaly integration tests set GITALY_ADDRESS=tcp://127.0.0.1:8075" ; \
+	else \
+		$(MAKE) run-gitaly ; \
+	fi
+	@go test -tags "$(BUILD_TAGS)" ./... ;\
 	status="$$?" ;\
 	if [ -f "$(GITALY_PID_FILE)" ] ; then \
 		echo "Clean up Gitaly server for workhorse integration test" ;\
@@ -93,27 +98,24 @@ clean-build:
 	rm -rf $(TARGET_DIR)
 
 .PHONY:	prepare-tests
-prepare-tests:	run-gitaly
 prepare-tests:	testdata/data/group/test.git $(EXE_ALL)
 prepare-tests:	testdata/scratch
 
 .PHONY: run-gitaly
-run-gitaly: gitaly.pid
+run-gitaly: $(GITALY_PID_FILE)
 
 $(GITALY_PID_FILE): gitaly.toml
-	@{ \
-	if [ -z "$${GITALY_ADDRESS+x}" ] ; then \
-		echo "To run gitaly integration tests set GITALY_ADDRESS=tcp://127.0.0.1:8075" ; \
-	else \
-		cd .. ; \
-		GITALY_TESTING_NO_GIT_HOOKS=1 GITALY_PID_FILE=workhorse/$(GITALY_PID_FILE) $(GITALY) workhorse/gitaly.toml ; \
-	fi \
-	} &
+	$(call message, "Starting gitaly")
+	cd ..; GITALY_TESTING_NO_GIT_HOOKS=1 GITALY_PID_FILE=workhorse/$(GITALY_PID_FILE) scripts/gitaly-test-spawn workhorse/gitaly.toml
 
 gitaly.toml: ../tmp/tests/gitaly/config.toml
 	sed -e 's/^socket_path.*$$/listen_addr = "0.0.0.0:8075"/;s/^\[auth\]$$//;s/^token.*$$//;s/^internal_socket_dir.*$$//' \
 		$< > $@
 
+../tmp/tests/gitaly/config.toml:
+	$(call message, "Building a complete test environment")
+	cd .. ; ./scripts/setup-test-env
+
 testdata/data/group/test.git:
 	$(call message,$@)
 	git clone --quiet --bare https://gitlab.com/gitlab-org/gitlab-test.git $@
diff --git a/workhorse/internal/proxy/proxy.go b/workhorse/internal/proxy/proxy.go
index 40be0fa96f6..be161c833a9 100644
--- a/workhorse/internal/proxy/proxy.go
+++ b/workhorse/internal/proxy/proxy.go
@@ -18,10 +18,17 @@ type Proxy struct {
 	Version                string
 	reverseProxy           *httputil.ReverseProxy
 	AllowResponseBuffering bool
+	customHeaders          map[string]string
 }
 
-func NewProxy(myURL *url.URL, version string, roundTripper http.RoundTripper) *Proxy {
-	p := Proxy{Version: version, AllowResponseBuffering: true}
+func WithCustomHeaders(customHeaders map[string]string) func(*Proxy) {
+	return func(proxy *Proxy) {
+		proxy.customHeaders = customHeaders
+	}
+}
+
+func NewProxy(myURL *url.URL, version string, roundTripper http.RoundTripper, options ...func(*Proxy)) *Proxy {
+	p := Proxy{Version: version, AllowResponseBuffering: true, customHeaders: make(map[string]string)}
 
 	if myURL == nil {
 		myURL = defaultTarget
@@ -31,6 +38,11 @@ func NewProxy(myURL *url.URL, version string, roundTripper http.RoundTripper) *P
 	u.Path = ""
 	p.reverseProxy = httputil.NewSingleHostReverseProxy(&u)
 	p.reverseProxy.Transport = roundTripper
+
+	for _, option := range options {
+		option(&p)
+	}
+
 	return &p
 }
 
@@ -43,6 +55,10 @@ func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	req.Header.Set("Gitlab-Workhorse", p.Version)
 	req.Header.Set("Gitlab-Workhorse-Proxy-Start", fmt.Sprintf("%d", time.Now().UnixNano()))
 
+	for k, v := range p.customHeaders {
+		req.Header.Set(k, v)
+	}
+
 	if p.AllowResponseBuffering {
 		helper.AllowResponseBuffering(w)
 	}
diff --git a/workhorse/internal/upstream/upstream.go b/workhorse/internal/upstream/upstream.go
index 065cae53e2b..c0678b1cb3e 100644
--- a/workhorse/internal/upstream/upstream.go
+++ b/workhorse/internal/upstream/upstream.go
@@ -37,6 +37,7 @@ var (
 		upload.RewrittenFieldsHeader,
 	}
 	geoProxyApiPollingInterval = 10 * time.Second
+	geoProxyWorkhorseHeaders   = map[string]string{"Gitlab-Workhorse-Geo-Proxy": "1"}
 )
 
 type upstream struct {
@@ -237,7 +238,12 @@ func (u *upstream) updateGeoProxyFields(geoProxyURL *url.URL) {
 	}
 
 	geoProxyRoundTripper := roundtripper.NewBackendRoundTripper(u.geoProxyBackend, "", u.ProxyHeadersTimeout, u.DevelopmentMode)
-	geoProxyUpstream := proxypkg.NewProxy(u.geoProxyBackend, u.Version, geoProxyRoundTripper)
+	geoProxyUpstream := proxypkg.NewProxy(
+		u.geoProxyBackend,
+		u.Version,
+		geoProxyRoundTripper,
+		proxypkg.WithCustomHeaders(geoProxyWorkhorseHeaders),
+	)
 	u.geoProxyCableRoute = u.wsRoute(`^/-/cable\z`, geoProxyUpstream)
 	u.geoProxyRoute = u.route("", "", geoProxyUpstream, withGeoProxy())
 }
diff --git a/workhorse/internal/upstream/upstream_test.go b/workhorse/internal/upstream/upstream_test.go
index 2031fa84ce2..80e59202b69 100644
--- a/workhorse/internal/upstream/upstream_test.go
+++ b/workhorse/internal/upstream/upstream_test.go
@@ -209,6 +209,23 @@ func TestGeoProxyFeatureEnablingAndDisabling(t *testing.T) {
 	runTestCases(t, ws, testCasesProxied)
 }
 
+func TestGeoProxySetsCustomHeader(t *testing.T) {
+	remoteServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		require.Equal(t, "1", r.Header.Get("Gitlab-Workhorse-Geo-Proxy"), "custom proxy header")
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer remoteServer.Close()
+
+	geoProxyEndpointResponseBody := fmt.Sprintf(`{"geo_proxy_url":"%v"}`, remoteServer.URL)
+	railsServer, deferredClose := startRailsServer("Local Rails server", &geoProxyEndpointResponseBody)
+	defer deferredClose()
+
+	ws, wsDeferredClose, _ := startWorkhorseServer(railsServer.URL, true)
+	defer wsDeferredClose()
+
+	http.Get(ws.URL)
+}
+
 func runTestCases(t *testing.T, ws *httptest.Server, testCases []testCase) {
 	t.Helper()
 	for _, tc := range testCases {
diff --git a/workhorse/proxy_test.go b/workhorse/proxy_test.go
index 48d7db428fd..754deea0032 100644
--- a/workhorse/proxy_test.go
+++ b/workhorse/proxy_test.go
@@ -22,12 +22,12 @@ import (
 
 const testVersion = "123"
 
-func newProxy(url string, rt http.RoundTripper) *proxy.Proxy {
+func newProxy(url string, rt http.RoundTripper, opts ...func(*proxy.Proxy)) *proxy.Proxy {
 	parsedURL := helper.URLMustParse(url)
 	if rt == nil {
 		rt = roundtripper.NewTestBackendRoundTripper(parsedURL)
 	}
-	return proxy.NewProxy(parsedURL, testVersion, rt)
+	return proxy.NewProxy(parsedURL, testVersion, rt, opts...)
 }
 
 func TestProxyRequest(t *testing.T) {
@@ -64,6 +64,24 @@ func TestProxyRequest(t *testing.T) {
 	require.Equal(t, "test", w.Header().Get("Custom-Response-Header"), "custom response header")
 }
 
+func TestProxyWithCustomHeaders(t *testing.T) {
+	ts := testhelper.TestServerWithHandler(regexp.MustCompile(`/url/path\z`), func(w http.ResponseWriter, r *http.Request) {
+		require.Equal(t, "value", r.Header.Get("Custom-Header"), "custom proxy header")
+		require.Equal(t, testVersion, r.Header.Get("Gitlab-Workhorse"), "version header")
+
+		_, err := w.Write([]byte(`ok`))
+		require.NoError(t, err, "write ok response")
+	})
+
+	httpRequest, err := http.NewRequest("POST", ts.URL+"/url/path", nil)
+	require.NoError(t, err)
+
+	w := httptest.NewRecorder()
+	testProxy := newProxy(ts.URL, nil, proxy.WithCustomHeaders(map[string]string{"Custom-Header": "value"}))
+	testProxy.ServeHTTP(w, httpRequest)
+	testhelper.RequireResponseBody(t, w, "ok")
+}
+
 func TestProxyError(t *testing.T) {
 	httpRequest, err := http.NewRequest("POST", "/url/path", bytes.NewBufferString("REQUEST"))
 	require.NoError(t, err)