diff options
Diffstat (limited to 'workhorse/internal/upstream')
-rw-r--r-- | workhorse/internal/upstream/upstream.go | 31 | ||||
-rw-r--r-- | workhorse/internal/upstream/upstream_test.go | 59 |
2 files changed, 80 insertions, 10 deletions
diff --git a/workhorse/internal/upstream/upstream.go b/workhorse/internal/upstream/upstream.go index c0678b1cb3e..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,34 +215,51 @@ 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, u.Version, geoProxyRoundTripper, proxypkg.WithCustomHeaders(geoProxyWorkhorseHeaders), + proxypkg.WithForcedTargetHostHeader(), ) 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 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) { |