Pass original IP from PROXY requests to internal API calls

2 files changed, 30 insertions, 0 deletions

diff --git a/client/client_test.go b/client/client_test.go
index 66ce2d8..06036b6 100644
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -76,6 +76,7 @@ func TestClients(t *testing.T) {
 			testErrorMessage(t, client)
 			testAuthenticationHeader(t, client)
 			testJWTAuthenticationHeader(t, client)
+			testXForwardedForHeader(t, client)
 		})
 	}
 }
@@ -221,6 +222,21 @@ func testJWTAuthenticationHeader(t *testing.T, client *GitlabNetClient) {
 	})
 }
 
+func testXForwardedForHeader(t *testing.T, client *GitlabNetClient) {
+	t.Run("X-Forwarded-For Header inserted if original address in context", func(t *testing.T) {
+		ctx := context.WithValue(context.Background(), OriginalRemoteIPContextKey{}, "196.7.0.238")
+		response, err := client.Get(ctx, "/x_forwarded_for")
+		require.NoError(t, err)
+		require.NotNil(t, response)
+
+		defer response.Body.Close()
+
+		responseBody, err := io.ReadAll(response.Body)
+		require.NoError(t, err)
+		require.Equal(t, "196.7.0.238", string(responseBody))
+	})
+}
+
 func buildRequests(t *testing.T, relativeURLRoot string) []testserver.TestRequestHandler {
 	requests := []testserver.TestRequestHandler{
 		{
@@ -257,6 +273,12 @@ func buildRequests(t *testing.T, relativeURLRoot string) []testserver.TestReques
 			},
 		},
 		{
+			Path: "/api/v4/internal/x_forwarded_for",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				fmt.Fprint(w, r.Header.Get("X-Forwarded-For"))
+			},
+		},
+		{
 			Path: "/api/v4/internal/error",
 			Handler: func(w http.ResponseWriter, r *http.Request) {
 				w.Header().Set("Content-Type", "application/json")
diff --git a/client/gitlabnet.go b/client/gitlabnet.go
index 31131d2..c34f148 100644
--- a/client/gitlabnet.go
+++ b/client/gitlabnet.go
@@ -41,6 +41,9 @@ type ApiError struct {
 	Msg string
 }
 
+// To use as the key in a Context to set an X-Forwarded-For header in a request
+type OriginalRemoteIPContextKey struct{}
+
 func (e *ApiError) Error() string {
 	return e.Msg
 }
@@ -150,6 +153,11 @@ func (c *GitlabNetClient) DoRequest(ctx context.Context, method, path string, da
 	}
 	request.Header.Set(apiSecretHeaderName, tokenString)
 
+	originalRemoteIP, ok := ctx.Value(OriginalRemoteIPContextKey{}).(string)
+	if ok {
+		request.Header.Add("X-Forwarded-For", originalRemoteIP)
+	}
+
 	request.Header.Add("Content-Type", "application/json")
 	request.Header.Add("User-Agent", c.userAgent)
 	request.Close = true