check for enterprise host

Signed-off-by: Meredith Lancaster <malancas@github.com>

pkg/cmd/attestation/auth/host.go

@@ -2,15 +2,32 @@ package auth
 
 import (
 	"errors"
+	"strings"
 
 	"github.com/cli/go-gh/v2/pkg/auth"
 )
 
-var ErrUnsupportedHost = errors.New("The GH_HOST environment variable is set to a custom GitHub host. gh attestation does not currently support custom GitHub Enterprise hosts")
+var ErrUnsupportedHost = errors.New("An unsupported host was detected. Note that gh attestation does not currently support GHES")
+
+const (
+	github    = "github.com"
+	localhost = "github.localhost"
+	// tenancyHost is the domain name of a tenancy GitHub instance
+	tenancyHost = "ghe.com"
+)
+
+func isEnterprise(host string) bool {
+	return host != github && host != localhost && !isTenancy(host)
+}
+
+func isTenancy(host string) bool {
+	return strings.HasSuffix(host, "."+tenancyHost)
+}
 
 func IsHostSupported() error {
 	host, _ := auth.DefaultHost()
-	if host != "github.com" {
+
+	if isEnterprise(host) {
 		return ErrUnsupportedHost
 	}
 	return nil

pkg/cmd/attestation/auth/host_test.go

@@ -0,0 +1,56 @@
+package auth
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestIsHostSupported(t *testing.T) {
+	testcases := []struct {
+		name        string
+		expectedErr bool
+		host        string
+	}{
+		{
+			name:        "Default github.com host",
+			expectedErr: false,
+			host:        "github.com",
+		},
+		{
+			name:        "Localhost",
+			expectedErr: false,
+			host:        "github.localhost",
+		},
+		{
+			name:        "No host set",
+			expectedErr: false,
+			host:        "",
+		},
+		{
+			name:        "GHE tenant host",
+			expectedErr: false,
+			host:        "some-tenant.ghe.com",
+		},
+		{
+			name:        "Unsupported host",
+			expectedErr: true,
+			host:        "my-unsupported-host.github.com",
+		},
+	}
+
+	for _, tc := range testcases {
+		err := os.Setenv("GH_HOST", tc.host)
+		require.NoError(t, err)
+
+		err = IsHostSupported()
+		if tc.expectedErr {
+			assert.Error(t, err)
+			assert.ErrorIs(t, err, ErrUnsupportedHost)
+		} else {
+			assert.NoError(t, err)
+		}
+	}
+}