From 49906996bd99c5b2dfa96a0ff9ed4756c649dfa6 Mon Sep 17 00:00:00 2001
From: Scott Murray <scottalexandermurray@gmail.com>
Date: Thu, 5 May 2022 03:45:17 +1200
Subject: [PATCH] only check availability for IMDS (#685)

Co-authored-by: Joel Hendrix <jhendrix@microsoft.com>
---
 autorest/adal/token.go      | 13 +++++++++++++
 autorest/adal/token_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/autorest/adal/token.go b/autorest/adal/token.go
index 0964a5881..07efaa402 100644
--- a/autorest/adal/token.go
+++ b/autorest/adal/token.go
@@ -1363,12 +1363,25 @@ func NewMultiTenantServicePrincipalTokenFromCertificate(multiTenantCfg MultiTena
 
 // MSIAvailable returns true if the MSI endpoint is available for authentication.
 func MSIAvailable(ctx context.Context, s Sender) bool {
+	msiType, _, err := getMSIType()
+
+	if err != nil {
+		return false
+	}
+
+	if msiType != msiTypeIMDS {
+		return true
+	}
+
 	if s == nil {
 		s = sender()
 	}
+
 	resp, err := getMSIEndpoint(ctx, s)
+
 	if err == nil {
 		resp.Body.Close()
 	}
+
 	return err == nil
 }
diff --git a/autorest/adal/token_test.go b/autorest/adal/token_test.go
index aa9291135..264ed749d 100644
--- a/autorest/adal/token_test.go
+++ b/autorest/adal/token_test.go
@@ -1396,6 +1396,32 @@ func TestMSIAvailableSuccess(t *testing.T) {
 	}
 }
 
+func TestMSIAvailableAppService(t *testing.T) {
+	os.Setenv("MSI_ENDPOINT", "http://localhost")
+	os.Setenv("MSI_SECRET", "super")
+	defer func() {
+		os.Unsetenv("MSI_ENDPOINT")
+		os.Unsetenv("MSI_SECRET")
+	}()
+	c := mocks.NewSender()
+	c.AppendResponse(mocks.NewResponse())
+	available := MSIAvailable(context.Background(), c)
+
+	if !available {
+		t.Fatal("expected MSI to be available")
+	}
+}
+
+func TestMSIAvailableIMDS(t *testing.T) {
+	c := mocks.NewSender()
+	c.AppendResponse(mocks.NewResponse())
+	available := MSIAvailable(context.Background(), c)
+
+	if !available {
+		t.Fatal("expected MSI to be available")
+	}
+}
+
 func TestMSIAvailableSlow(t *testing.T) {
 	c := mocks.NewSender()
 	// introduce a long response delay to simulate the endpoint not being available