Add code review remarks

2 files changed, 7 insertions, 15 deletions

diff --git a/go/internal/gitlabnet/client.go b/go/internal/gitlabnet/client.go
index 26c24d4..b668fbd 100644
--- a/go/internal/gitlabnet/client.go
+++ b/go/internal/gitlabnet/client.go
@@ -7,7 +7,6 @@ import (
 	"fmt"
 	"io"
 	"net/http"
-	"os"
 	"strings"
 
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
@@ -110,8 +109,7 @@ func (c *GitlabClient) DoRequest(method, path string, data interface{}) (*http.R
 	request.Header.Set(secretHeaderName, encodedSecret)
 
 	request.Header.Add("Content-Type", "application/json")
-	ip := ipAddr()
-	request.Header.Add("X_FORWARDED_FOR", ip)
+	request.Header.Add("X_FORWARDED_FOR", c.config.IpAddr())
 
 	request.Close = true
 
@@ -127,14 +125,6 @@ func (c *GitlabClient) DoRequest(method, path string, data interface{}) (*http.R
 	return response, nil
 }
 
-func ipAddr() string {
-	address := os.Getenv("SSH_CONNECTION")
-	if address != "" {
-		return strings.Fields(address)[0]
-	}
-	return address
-}
-
 func ParseJSON(hr *http.Response, response interface{}) error {
 	if err := json.NewDecoder(hr.Body).Decode(response); err != nil {
 		return ParsingError
diff --git a/go/internal/gitlabnet/client_test.go b/go/internal/gitlabnet/client_test.go
index debe618..07ac7b8 100644
--- a/go/internal/gitlabnet/client_test.go
+++ b/go/internal/gitlabnet/client_test.go
@@ -55,7 +55,7 @@ func TestClients(t *testing.T) {
 			Path: "/api/v4/internal/with_ip",
 			Handler: func(w http.ResponseWriter, r *http.Request) {
 				header := r.Header.Get("X_FORWARDED_FOR")
-				require.Equal(t, header, "127.0.0.1")
+				require.Equal(t, "127.0.0.1", header)
 			},
 		},
 		{
@@ -229,21 +229,23 @@ func testAuthenticationHeader(t *testing.T, client *GitlabClient) {
 
 func testXForwardedForHeader(t *testing.T, client *GitlabClient) {
 	t.Run("X-Forwarded-For for GET", func(t *testing.T) {
-		os.Setenv("SSH_CONNECTION", "127.0.0.1 0")
+		err := os.Setenv("SSH_CONNECTION", "127.0.0.1 0")
+		require.Nil(t, err)
+
 		response, err := client.Get("/with_ip")
-		defer response.Body.Close()
 
 		require.NoError(t, err)
 		require.NotNil(t, response)
+		response.Body.Close()
 	})
 
 	t.Run("X-Forwarded-For for POST", func(t *testing.T) {
 		data := map[string]string{"key": "value"}
 		os.Setenv("SSH_CONNECTION", "127.0.0.1 0")
 		response, err := client.Post("/with_ip", data)
-		defer response.Body.Close()
 
 		require.NoError(t, err)
 		require.NotNil(t, response)
+		response.Body.Close()
 	})
 }