Add latest changes from gitlab-org/gitlab@master

2 files changed, 24 insertions, 1 deletions

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 {