From 1a35ce38ade724c7d71daf01859e3552306e86f6 Mon Sep 17 00:00:00 2001
From: Meredith Lancaster <malancas@github.com>
Date: Thu, 4 Apr 2024 08:21:27 -0600
Subject: [PATCH] check for enterprise host

Signed-off-by: Meredith Lancaster <malancas@github.com>
---
 pkg/cmd/attestation/auth/host.go      | 21 +++++++++-
 pkg/cmd/attestation/auth/host_test.go | 56 +++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 2 deletions(-)
 create mode 100644 pkg/cmd/attestation/auth/host_test.go

diff --git a/pkg/cmd/attestation/auth/host.go b/pkg/cmd/attestation/auth/host.go
index 998dcb7f5..1e5206813 100644
--- a/pkg/cmd/attestation/auth/host.go
+++ b/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
diff --git a/pkg/cmd/attestation/auth/host_test.go b/pkg/cmd/attestation/auth/host_test.go
new file mode 100644
index 000000000..1192e1d9a
--- /dev/null
+++ b/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)
+		}
+	}
+}