Add cr remarks

8 files changed, 53 insertions, 23 deletions

diff --git a/go/internal/config/config.go b/go/internal/config/config.go
index 36a25c4..d958a82 100644
--- a/go/internal/config/config.go
+++ b/go/internal/config/config.go
@@ -6,7 +6,6 @@ import (
 	"os"
 	"path"
 	"path/filepath"
-	"strings"
 
 	yaml "gopkg.in/yaml.v2"
 )
@@ -36,7 +35,6 @@ type Config struct {
 	Secret         string             `yaml:"secret"`
 	HttpSettings   HttpSettingsConfig `yaml:"http_settings"`
 	HttpClient     *HttpClient
-	IPAddr         string
 }
 
 func New() (*Config, error) {
@@ -55,7 +53,6 @@ func NewFromDir(dir string) (*Config, error) {
 func newFromFile(filename string) (*Config, error) {
 	cfg := &Config{
 		RootDir: path.Dir(filename),
-		IPAddr:  getIPAddr(),
 	}
 
 	configBytes, err := ioutil.ReadFile(filename)
@@ -126,11 +123,3 @@ func parseSecret(cfg *Config) error {
 
 	return nil
 }
-
-func getIPAddr() string {
-	address := os.Getenv("SSH_CONNECTION")
-	if address != "" {
-		return strings.Fields(address)[0]
-	}
-	return address
-}
diff --git a/go/internal/config/config_test.go b/go/internal/config/config_test.go
index 1dfee28..e31ff70 100644
--- a/go/internal/config/config_test.go
+++ b/go/internal/config/config_test.go
@@ -2,7 +2,6 @@ package config
 
 import (
 	"fmt"
-	"os"
 	"path"
 	"testing"
 
@@ -111,10 +110,3 @@ func TestParseConfig(t *testing.T) {
 		})
 	}
 }
-
-func TestGetIPAddr(t *testing.T) {
-	err := os.Setenv("SSH_CONNECTION", "127.0.0.1 0")
-
-	require.Nil(t, err)
-	require.Equal(t, getIPAddr(), "127.0.0.1")
-}
diff --git a/go/internal/gitlabnet/client.go b/go/internal/gitlabnet/client.go
index 3aa30ad..e61b58d 100644
--- a/go/internal/gitlabnet/client.go
+++ b/go/internal/gitlabnet/client.go
@@ -10,6 +10,7 @@ import (
 	"strings"
 
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/sshenv"
 )
 
 const (
@@ -109,7 +110,10 @@ func (c *GitlabClient) DoRequest(method, path string, data interface{}) (*http.R
 	request.Header.Set(secretHeaderName, encodedSecret)
 
 	request.Header.Add("Content-Type", "application/json")
-	request.Header.Add("X_FORWARDED_FOR", c.config.IPAddr)
+	ipAddr := sshenv.LocalAddr()
+	if ipAddr != "" {
+		request.Header.Add("X-Forwarded-For", ipAddr)
+	}
 
 	request.Close = true
 
diff --git a/go/internal/gitlabnet/client_test.go b/go/internal/gitlabnet/client_test.go
index 3f25d3c..e813e23 100644
--- a/go/internal/gitlabnet/client_test.go
+++ b/go/internal/gitlabnet/client_test.go
@@ -53,7 +53,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")
+				header := r.Header.Get("X-Forwarded-For")
 				require.Equal(t, "127.0.0.1", header)
 			},
 		},
@@ -107,7 +107,6 @@ func TestClients(t *testing.T) {
 
 			tc.config.GitlabUrl = url
 			tc.config.Secret = "sssh, it's a secret"
-			tc.config.IPAddr = "127.0.0.1"
 			client, err := GetClient(tc.config)
 			require.NoError(t, err)
 
@@ -228,6 +227,10 @@ 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) {
+		cleanup, err := testhelper.Setenv("SSH_CONNECTION", "127.0.0.1 0")
+		require.NoError(t, err)
+		defer cleanup()
+
 		response, err := client.Get("/with_ip")
 
 		require.NoError(t, err)
@@ -237,6 +240,10 @@ func testXForwardedForHeader(t *testing.T, client *GitlabClient) {
 
 	t.Run("X-Forwarded-For for POST", func(t *testing.T) {
 		data := map[string]string{"key": "value"}
+		cleanup, err := testhelper.Setenv("SSH_CONNECTION", "127.0.0.1 0")
+		require.NoError(t, err)
+		defer cleanup()
+
 		response, err := client.Post("/with_ip", data)
 
 		require.NoError(t, err)
diff --git a/go/internal/sshenv/sshenv.go b/go/internal/sshenv/sshenv.go
new file mode 100644
index 0000000..c16e262
--- /dev/null
+++ b/go/internal/sshenv/sshenv.go
@@ -0,0 +1,15 @@
+package sshenv
+
+import (
+	"os"
+	"strings"
+)
+
+func LocalAddr() string {
+	address := os.Getenv("SSH_CONNECTION")
+
+	if address != "" {
+		return strings.Fields(address)[0]
+	}
+	return address
+}
diff --git a/go/internal/sshenv/sshenv_test.go b/go/internal/sshenv/sshenv_test.go
new file mode 100644
index 0000000..2691df8
--- /dev/null
+++ b/go/internal/sshenv/sshenv_test.go
@@ -0,0 +1,17 @@
+package sshenv
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/testhelper"
+)
+
+func TestLocalAddr(t *testing.T) {
+	cleanup, err := testhelper.Setenv("SSH_CONNECTION", "127.0.0.1 0")
+	require.NoError(t, err)
+	defer cleanup()
+
+	require.Nil(t, err)
+	require.Equal(t, LocalAddr(), "127.0.0.1")
+}
diff --git a/go/internal/testhelper/testhelper.go b/go/internal/testhelper/testhelper.go
index 5c900aa..a925c79 100644
--- a/go/internal/testhelper/testhelper.go
+++ b/go/internal/testhelper/testhelper.go
@@ -85,3 +85,9 @@ func getTestDataDir() (string, error) {
 
 	return path.Join(path.Dir(currentFile), "testdata"), nil
 }
+
+func Setenv(key, value string) (func(), error) {
+	oldValue := os.Getenv(key)
+	err := os.Setenv(key, value)
+	return func() { os.Setenv(key, oldValue) }, err
+}
diff --git a/support/go-test b/support/go-test
index 0c043f6..47fd012 100755
--- a/support/go-test
+++ b/support/go-test
@@ -5,7 +5,7 @@ include GoBuild
 
 def main
   ensure_build_dir_exists
-  run!(GO_ENV, %w[go test ./...], chdir: GO_DIR)
+  run!(GO_ENV, %w[go test -v ./...], chdir: GO_DIR)
   puts 'OK'
 end