Set correct value of X-Forwarded-For header in PreAuthorize request

internal/api/api.go

@@ -157,6 +157,8 @@ func (api *API) newRequest(r *http.Request, body io.Reader, suffix string) (*htt
 	// configurations (Passenger) to solve auth request routing problems.
 	authReq.Header.Set("Gitlab-Workhorse", api.Version)
+	helper.SetForwardedFor(&authReq.Header, r)
+
 	tokenString, err := secret.JWTTokenString(secret.DefaultClaims)
 	if err != nil {
 		return nil, fmt.Errorf("newRequest: sign JWT: %v", err)

internal/helper/helpers.go

@@ -3,10 +3,12 @@ package helper
 import (
 	"errors"
 	"log"
+	"net"
 	"net/http"
 	"net/url"
 	"os"
 	"os/exec"
+	"strings"
 	"syscall"
 )
@@ -142,3 +144,19 @@ func DisableResponseBuffering(w http.ResponseWriter) {
 func AllowResponseBuffering(w http.ResponseWriter) {
 	w.Header().Del(NginxResponseBufferHeader)
 }
+
+func SetForwardedFor(newHeaders *http.Header, originalRequest *http.Request) {
+	if clientIP, _, err := net.SplitHostPort(originalRequest.RemoteAddr); err == nil {
+		var header string
+
+		// If we aren't the first proxy retain prior
+		// X-Forwarded-For information as a comma+space
+		// separated list and fold multiple headers into one.
+		if prior, ok := originalRequest.Header["X-Forwarded-For"]; ok {
+			header = strings.Join(prior, ", ") + ", " + clientIP
+		} else {
+			header = clientIP
+		}
+		newHeaders.Set("X-Forwarded-For", header)
+	}
+}

internal/helper/helpers_test.go

+package helper
+
+import (
+	"net/http"
+	"testing"
+)
+
+func TestSetForwardedForGeneratesHeader(t *testing.T) {
+	testCases := []struct {
+		remoteAddr           string
+		previousForwardedFor []string
+		expected             string
+	}{
+		{
+			"8.8.8.8:3000",
+			nil,
+			"8.8.8.8",
+		},
+		{
+			"8.8.8.8:3000",
+			[]string{"138.124.33.63, 151.146.211.237"},
+			"138.124.33.63, 151.146.211.237, 8.8.8.8",
+		},
+		{
+			"8.8.8.8:3000",
+			[]string{"8.154.76.107", "115.206.118.179"},
+			"8.154.76.107, 115.206.118.179, 8.8.8.8",
+		},
+	}
+	for _, tc := range testCases {
+		headers := http.Header{}
+		originalRequest := http.Request{
+			RemoteAddr: tc.remoteAddr,
+		}
+
+		if tc.previousForwardedFor != nil {
+			originalRequest.Header = http.Header{
+				"X-Forwarded-For": tc.previousForwardedFor,
+			}
+		}
+
+		SetForwardedFor(&headers, &originalRequest)
+
+		result := headers.Get("X-Forwarded-For")
+		if result != tc.expected {
+			t.Fatalf("Expected %v, got %v", tc.expected, result)
+		}
+	}
+}

internal/terminal/terminal.go

@@ -2,9 +2,7 @@ package terminal
 import (
 	"log"
-	"net"
 	"net/http"
-	"strings"
 	"time"
 	"github.com/gorilla/websocket"
@@ -102,15 +100,7 @@ func pingLoop(conn Connection) {
 func connectToServer(terminal *api.TerminalSettings, r *http.Request) (Connection, error) {
 	terminal = terminal.Clone()
-	// Pass along X-Forwarded-For, appending request.RemoteAddr, to the server
-	// we're connecting to.
-	if ip, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
-		if chains, ok := r.Header["X-Forwarded-For"]; ok {
-			terminal.Header.Set("X-Forwarded-For", strings.Join(chains, ", ")+", "+ip)
-		} else {
-			terminal.Header.Set("X-Forwarded-For", ip)
-		}
-	}
+	helper.SetForwardedFor(&terminal.Header, r)
 	conn, _, err := terminal.Dial()
 	if err != nil {