From 09cbe8ba653c26b2bcb5ee6ec09a2144817423d0 Mon Sep 17 00:00:00 2001
From: Remi Vichery <remi@alkira.com>
Date: Sun, 26 Mar 2023 11:11:57 -0700
Subject: [PATCH] fixup! Add identity token for all Azure cloud environments

---
 authority/provisioner/azure.go      | 13 +++++++++++--
 authority/provisioner/azure_test.go |  2 ++
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/authority/provisioner/azure.go b/authority/provisioner/azure.go
index 03bebd4c..fcfbab27 100644
--- a/authority/provisioner/azure.go
+++ b/authority/provisioner/azure.go
@@ -43,7 +43,9 @@ var azureXMSMirIDRegExp = regexp.MustCompile(`(?i)^/subscriptions/([^/]+)/resour
 // azureEnvironments is the list of all Azure environments.
 var azureEnvironments = map[string]string{
 	"AzurePublicCloud":       "https://management.azure.com/",
+	"AzureCloud":             "https://management.azure.com/",
 	"AzureUSGovernmentCloud": "https://management.usgovcloudapi.net/",
+	"AzureUSGovernment":      "https://management.usgovcloudapi.net/",
 	"AzureChinaCloud":        "https://management.chinacloudapi.cn/",
 	"AzureGermanCloud":       "https://management.microsoftazure.de/",
 }
@@ -118,6 +120,7 @@ type Azure struct {
 	oidcConfig             openIDConfiguration
 	keyStore               *keyStore
 	ctl                    *Controller
+	environment            string
 }
 
 // GetID returns the provisioner unique identifier.
@@ -184,12 +187,14 @@ func (p *Azure) GetIdentityToken(subject, caURL string) (string, error) {
 
 	// default to AzurePublicCloud to keep existing behavior
 	identityTokenResource := azureEnvironments["AzurePublicCloud"]
-	environment, err := p.getAzureEnvironment()
+
+	var err error
+	p.environment, err = p.getAzureEnvironment()
 	if err != nil {
 		return "", errors.Wrap(err, "error getting azure environment")
 	}
 
-	if resource, ok := azureEnvironments[environment]; ok {
+	if resource, ok := azureEnvironments[p.environment]; ok {
 		identityTokenResource = resource
 	}
 
@@ -479,6 +484,10 @@ func (p *Azure) assertConfig() {
 
 // getAzureEnvironment returns the Azure environment for the current instance
 func (p *Azure) getAzureEnvironment() (string, error) {
+	if p.environment != "" {
+		return p.environment, nil
+	}
+
 	req, err := http.NewRequest("GET", p.config.instanceComputeURL, http.NoBody)
 	if err != nil {
 		return "", errors.Wrap(err, "error creating request")
diff --git a/authority/provisioner/azure_test.go b/authority/provisioner/azure_test.go
index 89a9d644..51d46c5a 100644
--- a/authority/provisioner/azure_test.go
+++ b/authority/provisioner/azure_test.go
@@ -166,6 +166,8 @@ func TestAzure_GetIdentityToken(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
+			// reset environment between tests to avoid caching issues
+			p1.environment = ""
 			tt.azure.config.identityTokenURL = tt.identityTokenURL + "?want_resource=" + azureEnvironments[tt.wantEnvironment]
 			tt.azure.config.instanceComputeURL = tt.instanceComputeURL + "/" + tt.wantEnvironment
 			got, err := tt.azure.GetIdentityToken(tt.args.subject, tt.args.caURL)