From 3dc5172cf005604c85ec7717f11c14178bfdcd19 Mon Sep 17 00:00:00 2001
From: catalinaperalta <catalinaperaltah@hotmail.com>
Date: Mon, 12 Jul 2021 09:16:09 -0700
Subject: [PATCH] azidentity: Add environment variable for resource ID (#14872)

Fixes #14742
---
 sdk/azidentity/managed_identity_credential.go |  6 +++-
 .../managed_identity_credential_test.go       | 30 +++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/sdk/azidentity/managed_identity_credential.go b/sdk/azidentity/managed_identity_credential.go
index 48880536dc31..cf64e837ed9c 100644
--- a/sdk/azidentity/managed_identity_credential.go
+++ b/sdk/azidentity/managed_identity_credential.go
@@ -72,7 +72,11 @@ func NewManagedIdentityCredential(id string, options *ManagedIdentityCredentialO
 	client.msiType = msiType
 	// check if no clientID is specified then check if it exists in an environment variable
 	if len(id) == 0 {
-		id = os.Getenv("AZURE_CLIENT_ID")
+		if options.ID == ResourceID {
+			id = os.Getenv("AZURE_RESOURCE_ID")
+		} else {
+			id = os.Getenv("AZURE_CLIENT_ID")
+		}
 	}
 	return &ManagedIdentityCredential{id: id, client: client}, nil
 }
diff --git a/sdk/azidentity/managed_identity_credential_test.go b/sdk/azidentity/managed_identity_credential_test.go
index 7cd3cf659d81..718b4c95f457 100644
--- a/sdk/azidentity/managed_identity_credential_test.go
+++ b/sdk/azidentity/managed_identity_credential_test.go
@@ -662,3 +662,33 @@ func TestManagedIdentityCredential_CreateAccessTokenExpiresOnFail(t *testing.T)
 		t.Fatalf("expected to receive an error but received none")
 	}
 }
+
+func TestManagedIdentityCredential_ResourceID_envVar(t *testing.T) {
+	// setting a dummy value for IDENTITY_ENDPOINT in order to be able to get a ManagedIdentityCredential type
+	_ = os.Setenv("IDENTITY_ENDPOINT", "somevalue")
+	_ = os.Setenv("IDENTITY_HEADER", "header")
+	_ = os.Setenv("AZURE_CLIENT_ID", "client_id")
+	_ = os.Setenv("AZURE_RESOURCE_ID", "resource_id")
+	defer clearEnvVars("IDENTITY_ENDPOINT", "IDENTITY_HEADER", "AZURE_CLIENT_ID", "AZURE_RESOURCE_ID")
+	cred, err := NewManagedIdentityCredential("", &ManagedIdentityCredentialOptions{ID: ResourceID})
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if cred.id != "resource_id" {
+		t.Fatal("unexpected id value stored")
+	}
+	cred, err = NewManagedIdentityCredential("", nil)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if cred.id != "client_id" {
+		t.Fatal("unexpected id value stored")
+	}
+	cred, err = NewManagedIdentityCredential("", &ManagedIdentityCredentialOptions{ID: ClientID})
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if cred.id != "client_id" {
+		t.Fatal("unexpected id value stored")
+	}
+}