4 files changed, 59 insertions, 9 deletions

diff --git a/workhorse/internal/upstream/handlers.go b/workhorse/internal/upstream/handlers.go
index 5974170e172..85fee0bf7e2 100644
--- a/workhorse/internal/upstream/handlers.go
+++ b/workhorse/internal/upstream/handlers.go
@@ -6,7 +6,7 @@ import (
 	"io"
 	"net/http"
 
-	"gitlab.com/gitlab-org/gitlab/workhorse/internal/helper"
+	"gitlab.com/gitlab-org/gitlab/workhorse/internal/helper/fail"
 )
 
 func contentEncodingHandler(h http.Handler) http.Handler {
@@ -26,7 +26,7 @@ func contentEncodingHandler(h http.Handler) http.Handler {
 		}
 
 		if err != nil {
-			helper.Fail500(w, r, fmt.Errorf("contentEncodingHandler: %v", err))
+			fail.Request(w, r, fmt.Errorf("contentEncodingHandler: %v", err))
 			return
 		}
 		defer body.Close()
diff --git a/workhorse/internal/upstream/routes.go b/workhorse/internal/upstream/routes.go
index c47053ad682..982f3a5b5f8 100644
--- a/workhorse/internal/upstream/routes.go
+++ b/workhorse/internal/upstream/routes.go
@@ -425,7 +425,7 @@ func configureRoutes(u *upstream) {
 func denyWebsocket(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if websocket.IsWebSocketUpgrade(r) {
-			helper.HTTPError(w, r, "websocket upgrade not allowed", http.StatusBadRequest)
+			httpError(w, r, "websocket upgrade not allowed", http.StatusBadRequest)
 			return
 		}
 
diff --git a/workhorse/internal/upstream/upstream.go b/workhorse/internal/upstream/upstream.go
index 248f190e316..34fe300192f 100644
--- a/workhorse/internal/upstream/upstream.go
+++ b/workhorse/internal/upstream/upstream.go
@@ -16,6 +16,7 @@ import (
 	"net/url"
 	"strings"
 
+	"github.com/sebest/xff"
 	"github.com/sirupsen/logrus"
 
 	"gitlab.com/gitlab-org/labkit/correlation"
@@ -24,6 +25,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab/workhorse/internal/builds"
 	"gitlab.com/gitlab-org/gitlab/workhorse/internal/config"
 	"gitlab.com/gitlab-org/gitlab/workhorse/internal/helper"
+	"gitlab.com/gitlab-org/gitlab/workhorse/internal/helper/nginx"
 	proxypkg "gitlab.com/gitlab-org/gitlab/workhorse/internal/proxy"
 	"gitlab.com/gitlab-org/gitlab/workhorse/internal/rejectmethods"
 	"gitlab.com/gitlab-org/gitlab/workhorse/internal/upload"
@@ -124,19 +126,19 @@ func (u *upstream) configureURLPrefix() {
 }
 
 func (u *upstream) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	helper.FixRemoteAddr(r)
+	fixRemoteAddr(r)
 
-	helper.DisableResponseBuffering(w)
+	nginx.DisableResponseBuffering(w)
 
 	// Drop RequestURI == "*" (FIXME: why?)
 	if r.RequestURI == "*" {
-		helper.HTTPError(w, r, "Connection upgrade not allowed", http.StatusBadRequest)
+		httpError(w, r, "Connection upgrade not allowed", http.StatusBadRequest)
 		return
 	}
 
 	// Disallow connect
 	if r.Method == "CONNECT" {
-		helper.HTTPError(w, r, "CONNECT not allowed", http.StatusBadRequest)
+		httpError(w, r, "CONNECT not allowed", http.StatusBadRequest)
 		return
 	}
 
@@ -144,7 +146,7 @@ func (u *upstream) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	URIPath := urlprefix.CleanURIPath(r.URL.EscapedPath())
 	prefix := u.URLPrefix
 	if !prefix.Match(URIPath) {
-		helper.HTTPError(w, r, fmt.Sprintf("Not found %q", URIPath), http.StatusNotFound)
+		httpError(w, r, fmt.Sprintf("Not found %q", URIPath), http.StatusNotFound)
 		return
 	}
 
@@ -155,7 +157,7 @@ func (u *upstream) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	if route == nil {
 		// The protocol spec in git/Documentation/technical/http-protocol.txt
 		// says we must return 403 if no matching service is found.
-		helper.HTTPError(w, r, "Forbidden", http.StatusForbidden)
+		httpError(w, r, "Forbidden", http.StatusForbidden)
 		return
 	}
 
@@ -275,3 +277,21 @@ func (u *upstream) updateGeoProxyFieldsFromData(geoProxyData *apipkg.GeoProxyDat
 	u.geoProxyCableRoute = u.wsRoute(`^/-/cable\z`, geoProxyUpstream)
 	u.geoProxyRoute = u.route("", "", geoProxyUpstream, withGeoProxy())
 }
+
+func httpError(w http.ResponseWriter, r *http.Request, error string, code int) {
+	if r.ProtoAtLeast(1, 1) {
+		// Force client to disconnect if we render request error
+		w.Header().Set("Connection", "close")
+	}
+
+	http.Error(w, error, code)
+}
+
+func fixRemoteAddr(r *http.Request) {
+	// Unix domain sockets have a remote addr of @. This will make the
+	// xff package lookup the X-Forwarded-For address if available.
+	if r.RemoteAddr == "@" {
+		r.RemoteAddr = "127.0.0.1:0"
+	}
+	r.RemoteAddr = xff.GetRemoteAddr(r)
+}
diff --git a/workhorse/internal/upstream/upstream_test.go b/workhorse/internal/upstream/upstream_test.go
index 7ab3e67116f..705e40c74d5 100644
--- a/workhorse/internal/upstream/upstream_test.go
+++ b/workhorse/internal/upstream/upstream_test.go
@@ -435,3 +435,33 @@ func startWorkhorseServer(railsServerURL string, enableGeoProxyFeature bool) (*h
 
 	return ws, ws.Close, waitForNextApiPoll
 }
+
+func TestFixRemoteAddr(t *testing.T) {
+	testCases := []struct {
+		initial   string
+		forwarded string
+		expected  string
+	}{
+		{initial: "@", forwarded: "", expected: "127.0.0.1:0"},
+		{initial: "@", forwarded: "18.245.0.1", expected: "18.245.0.1:0"},
+		{initial: "@", forwarded: "127.0.0.1", expected: "127.0.0.1:0"},
+		{initial: "@", forwarded: "192.168.0.1", expected: "127.0.0.1:0"},
+		{initial: "192.168.1.1:0", forwarded: "", expected: "192.168.1.1:0"},
+		{initial: "192.168.1.1:0", forwarded: "18.245.0.1", expected: "18.245.0.1:0"},
+	}
+
+	for _, tc := range testCases {
+		req, err := http.NewRequest("POST", "unix:///tmp/test.socket/info/refs", nil)
+		require.NoError(t, err)
+
+		req.RemoteAddr = tc.initial
+
+		if tc.forwarded != "" {
+			req.Header.Add("X-Forwarded-For", tc.forwarded)
+		}
+
+		fixRemoteAddr(req)
+
+		require.Equal(t, tc.expected, req.RemoteAddr)
+	}
+}