Add latest changes from gitlab-org/gitlab@master

2 files changed, 79 insertions, 10 deletions

diff --git a/workhorse/internal/upstream/upstream.go b/workhorse/internal/upstream/upstream.go
index 6258016c10c..6d107fc28cd 100644
--- a/workhorse/internal/upstream/upstream.go
+++ b/workhorse/internal/upstream/upstream.go
@@ -37,7 +37,6 @@ var (
 		upload.RewrittenFieldsHeader,
 	}
 	geoProxyApiPollingInterval = 10 * time.Second
-	geoProxyWorkhorseHeaders   = map[string]string{"Gitlab-Workhorse-Geo-Proxy": "1"}
 )
 
 type upstream struct {
@@ -48,6 +47,7 @@ type upstream struct {
 	CableRoundTripper     http.RoundTripper
 	APIClient             *apipkg.API
 	geoProxyBackend       *url.URL
+	geoProxyExtraData     string
 	geoLocalRoutes        []routeEntry
 	geoProxyCableRoute    routeEntry
 	geoProxyRoute         routeEntry
@@ -215,28 +215,44 @@ func (u *upstream) pollGeoProxyAPI() {
 
 // Calls /api/v4/geo/proxy and sets up routes
 func (u *upstream) callGeoProxyAPI() {
-	geoProxyURL, err := u.APIClient.GetGeoProxyURL()
+	geoProxyData, err := u.APIClient.GetGeoProxyData()
 	if err != nil {
 		log.WithError(err).WithFields(log.Fields{"geoProxyBackend": u.geoProxyBackend}).Error("Geo Proxy: Unable to determine Geo Proxy URL. Fallback on cached value.")
 		return
 	}
 
-	if u.geoProxyBackend.String() != geoProxyURL.String() {
-		log.WithFields(log.Fields{"oldGeoProxyURL": u.geoProxyBackend, "newGeoProxyURL": geoProxyURL}).Info("Geo Proxy: URL changed")
-		u.updateGeoProxyFields(geoProxyURL)
+	hasProxyDataChanged := false
+	if u.geoProxyBackend.String() != geoProxyData.GeoProxyURL.String() {
+		log.WithFields(log.Fields{"oldGeoProxyURL": u.geoProxyBackend, "newGeoProxyURL": geoProxyData.GeoProxyURL}).Info("Geo Proxy: URL changed")
+		hasProxyDataChanged = true
+	}
+
+	if u.geoProxyExtraData != geoProxyData.GeoProxyExtraData {
+		// extra data is usually a JWT, thus not explicitly logging it
+		log.Info("Geo Proxy: signed data changed")
+		hasProxyDataChanged = true
+	}
+
+	if hasProxyDataChanged {
+		u.updateGeoProxyFieldsFromData(geoProxyData)
 	}
 }
 
-func (u *upstream) updateGeoProxyFields(geoProxyURL *url.URL) {
+func (u *upstream) updateGeoProxyFieldsFromData(geoProxyData *apipkg.GeoProxyData) {
 	u.mu.Lock()
 	defer u.mu.Unlock()
 
-	u.geoProxyBackend = geoProxyURL
+	u.geoProxyBackend = geoProxyData.GeoProxyURL
+	u.geoProxyExtraData = geoProxyData.GeoProxyExtraData
 
 	if u.geoProxyBackend.String() == "" {
 		return
 	}
 
+	geoProxyWorkhorseHeaders := map[string]string{
+		"Gitlab-Workhorse-Geo-Proxy":            "1",
+		"Gitlab-Workhorse-Geo-Proxy-Extra-Data": u.geoProxyExtraData,
+	}
 	geoProxyRoundTripper := roundtripper.NewBackendRoundTripper(u.geoProxyBackend, "", u.ProxyHeadersTimeout, u.DevelopmentMode)
 	geoProxyUpstream := proxypkg.NewProxy(
 		u.geoProxyBackend,
diff --git a/workhorse/internal/upstream/upstream_test.go b/workhorse/internal/upstream/upstream_test.go
index 80e59202b69..8f054f5ccef 100644
--- a/workhorse/internal/upstream/upstream_test.go
+++ b/workhorse/internal/upstream/upstream_test.go
@@ -209,21 +209,74 @@ func TestGeoProxyFeatureEnablingAndDisabling(t *testing.T) {
 	runTestCases(t, ws, testCasesProxied)
 }
 
-func TestGeoProxySetsCustomHeader(t *testing.T) {
+func TestGeoProxyUpdatesExtraDataWhenChanged(t *testing.T) {
+	var expectedGeoProxyExtraData string
+
 	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")
+		require.Equal(t, expectedGeoProxyExtraData, r.Header.Get("Gitlab-Workhorse-Geo-Proxy-Extra-Data"), "custom extra data header")
 		w.WriteHeader(http.StatusOK)
 	}))
 	defer remoteServer.Close()
 
-	geoProxyEndpointResponseBody := fmt.Sprintf(`{"geo_proxy_url":"%v"}`, remoteServer.URL)
+	geoProxyEndpointExtraData1 := fmt.Sprintf(`{"geo_proxy_url":"%v","geo_proxy_extra_data":"data1"}`, remoteServer.URL)
+	geoProxyEndpointExtraData2 := fmt.Sprintf(`{"geo_proxy_url":"%v","geo_proxy_extra_data":"data2"}`, remoteServer.URL)
+	geoProxyEndpointExtraData3 := fmt.Sprintf(`{"geo_proxy_url":"%v"}`, remoteServer.URL)
+	geoProxyEndpointResponseBody := geoProxyEndpointExtraData1
+	expectedGeoProxyExtraData = "data1"
+
 	railsServer, deferredClose := startRailsServer("Local Rails server", &geoProxyEndpointResponseBody)
 	defer deferredClose()
 
-	ws, wsDeferredClose, _ := startWorkhorseServer(railsServer.URL, true)
+	ws, wsDeferredClose, waitForNextApiPoll := startWorkhorseServer(railsServer.URL, true)
 	defer wsDeferredClose()
 
 	http.Get(ws.URL)
+
+	// Verify that the expected header changes after next updated poll.
+	geoProxyEndpointResponseBody = geoProxyEndpointExtraData2
+	expectedGeoProxyExtraData = "data2"
+	waitForNextApiPoll()
+
+	http.Get(ws.URL)
+
+	// Validate that non-existing extra data results in empty header
+	geoProxyEndpointResponseBody = geoProxyEndpointExtraData3
+	expectedGeoProxyExtraData = ""
+	waitForNextApiPoll()
+
+	http.Get(ws.URL)
+}
+
+func TestGeoProxySetsCustomHeader(t *testing.T) {
+	testCases := []struct {
+		desc      string
+		json      string
+		extraData string
+	}{
+		{"no extra data", `{"geo_proxy_url":"%v"}`, ""},
+		{"with extra data", `{"geo_proxy_url":"%v","geo_proxy_extra_data":"extra-geo-data"}`, "extra-geo-data"},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(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")
+				require.Equal(t, tc.extraData, r.Header.Get("Gitlab-Workhorse-Geo-Proxy-Extra-Data"), "custom proxy extra data header")
+				w.WriteHeader(http.StatusOK)
+			}))
+			defer remoteServer.Close()
+
+			geoProxyEndpointResponseBody := fmt.Sprintf(tc.json, 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) {