From ba8c41a6ed2b2fe771a0aac4b1b6de774bffa374 Mon Sep 17 00:00:00 2001
From: Sam Coe <samcoe@users.noreply.github.com>
Date: Sat, 29 Jan 2022 09:37:24 +0200
Subject: [PATCH] Add GH_HOST to hosts list if it has corresponding auth token
 (#5087)

* Add GH_HOST to hosts list if it has coresponding auth token

* Keep default host at front of hosts list

* Always add GH_HOST to hosts list
---
 internal/config/from_env.go      | 32 ++++++++++++++++++++++----------
 internal/config/from_env_test.go | 15 +++++++++++++++
 pkg/set/string_set.go            |  1 +
 3 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/internal/config/from_env.go b/internal/config/from_env.go
index 27cf3c54b..3cc19879d 100644
--- a/internal/config/from_env.go
+++ b/internal/config/from_env.go
@@ -3,9 +3,11 @@ package config
 import (
 	"fmt"
 	"os"
+	"sort"
 	"strconv"
 
 	"github.com/cli/cli/v2/internal/ghinstance"
+	"github.com/cli/cli/v2/pkg/set"
 )
 
 const (
@@ -34,19 +36,29 @@ type envConfig struct {
 }
 
 func (c *envConfig) Hosts() ([]string, error) {
-	hasDefault := false
 	hosts, err := c.Config.Hosts()
-	for _, h := range hosts {
-		if h == ghinstance.Default() {
-			hasDefault = true
-		}
+	if err != nil {
+		return nil, err
 	}
-	token, _ := AuthTokenFromEnv(ghinstance.Default())
-	if (err != nil || !hasDefault) && token != "" {
-		hosts = append([]string{ghinstance.Default()}, hosts...)
-		return hosts, nil
+
+	hostSet := set.NewStringSet()
+	hostSet.AddValues(hosts)
+
+	// If GH_HOST is set then add it to list.
+	if host := os.Getenv(GH_HOST); host != "" {
+		hostSet.Add(host)
 	}
-	return hosts, err
+
+	// If there is a valid environment variable token for the
+	// default host then add default host to list.
+	if token, _ := AuthTokenFromEnv(ghinstance.Default()); token != "" {
+		hostSet.Add(ghinstance.Default())
+	}
+
+	s := hostSet.ToSlice()
+	// If default host is in list then move it to the front.
+	sort.SliceStable(s, func(i, j int) bool { return s[i] == ghinstance.Default() })
+	return s, nil
 }
 
 func (c *envConfig) DefaultHost() (string, error) {
diff --git a/internal/config/from_env_test.go b/internal/config/from_env_test.go
index bf81c7976..4bce09e85 100644
--- a/internal/config/from_env_test.go
+++ b/internal/config/from_env_test.go
@@ -44,6 +44,7 @@ func TestInheritEnv(t *testing.T) {
 	tests := []struct {
 		name                    string
 		baseConfig              string
+		GH_HOST                 string
 		GITHUB_TOKEN            string
 		GITHUB_ENTERPRISE_TOKEN string
 		GH_TOKEN                string
@@ -285,9 +286,23 @@ func TestInheritEnv(t *testing.T) {
 				writeable: false,
 			},
 		},
+		{
+			name:                "GH_HOST adds host entry when paired with environment token",
+			baseConfig:          ``,
+			GH_HOST:             "example.org",
+			GH_ENTERPRISE_TOKEN: "GH_ENTERPRISE_TOKEN",
+			hostname:            "example.org",
+			wants: wants{
+				hosts:     []string{"example.org"},
+				token:     "GH_ENTERPRISE_TOKEN",
+				source:    "GH_ENTERPRISE_TOKEN",
+				writeable: false,
+			},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
+			setenv(t, "GH_HOST", tt.GH_HOST)
 			setenv(t, "GITHUB_TOKEN", tt.GITHUB_TOKEN)
 			setenv(t, "GITHUB_ENTERPRISE_TOKEN", tt.GITHUB_ENTERPRISE_TOKEN)
 			setenv(t, "GH_TOKEN", tt.GH_TOKEN)
diff --git a/pkg/set/string_set.go b/pkg/set/string_set.go
index ea567c74b..0bd679a8d 100644
--- a/pkg/set/string_set.go
+++ b/pkg/set/string_set.go
@@ -10,6 +10,7 @@ type stringSet struct {
 func NewStringSet() *stringSet {
 	s := &stringSet{}
 	s.m = make(map[string]struct{})
+	s.v = []string{}
 	return s
 }