From 9bd19df04095b75517b0858207e941e6b27e4a58 Mon Sep 17 00:00:00 2001
From: Feiyu Shi <feshi@microsoft.com>
Date: Sun, 24 May 2020 16:53:26 -0700
Subject: [PATCH] validate if user assigned identity exists

---
 builder/azure/arm/azure_client.go |  9 +++++++++
 builder/azure/arm/builder.go      | 13 +++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/builder/azure/arm/azure_client.go b/builder/azure/arm/azure_client.go
index 2fdfafd8f..0a2c1165d 100644
--- a/builder/azure/arm/azure_client.go
+++ b/builder/azure/arm/azure_client.go
@@ -14,6 +14,7 @@ import (
 	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2018-04-01/compute"
 	newCompute "github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-03-01/compute"
 	"github.com/Azure/azure-sdk-for-go/services/keyvault/mgmt/2018-02-14/keyvault"
+	"github.com/Azure/azure-sdk-for-go/services/msi/mgmt/2018-11-30/msi"
 	"github.com/Azure/azure-sdk-for-go/services/network/mgmt/2018-01-01/network"
 	"github.com/Azure/azure-sdk-for-go/services/resources/mgmt/2018-02-01/resources"
 	armStorage "github.com/Azure/azure-sdk-for-go/services/storage/mgmt/2017-10-01/storage"
@@ -47,6 +48,7 @@ type AzureClient struct {
 	compute.SnapshotsClient
 	newCompute.GalleryImageVersionsClient
 	newCompute.GalleryImagesClient
+	msi.UserAssignedIdentitiesClient
 
 	InspectorMaxLength int
 	Template           *CaptureTemplate
@@ -240,6 +242,13 @@ func NewAzureClient(subscriptionID, resourceGroupName, storageAccountName string
 	azureClient.GalleryImagesClient.UserAgent = fmt.Sprintf("%s %s", useragent.String(), azureClient.GalleryImagesClient.UserAgent)
 	azureClient.GalleryImagesClient.Client.PollingDuration = PollingDuration
 
+	azureClient.UserAssignedIdentitiesClient = msi.NewUserAssignedIdentitiesClientWithBaseURI(cloud.ResourceManagerEndpoint, subscriptionID)
+	azureClient.UserAssignedIdentitiesClient.Authorizer = autorest.NewBearerAuthorizer(servicePrincipalToken)
+	azureClient.UserAssignedIdentitiesClient.RequestInspector = withInspection(maxlen)
+	azureClient.UserAssignedIdentitiesClient.ResponseInspector = byConcatDecorators(byInspecting(maxlen), errorCapture(azureClient))
+	azureClient.UserAssignedIdentitiesClient.UserAgent = fmt.Sprintf("%s %s", useragent.String(), azureClient.UserAssignedIdentitiesClient.UserAgent)
+	azureClient.UserAssignedIdentitiesClient.Client.PollingDuration = PollingDuration
+
 	keyVaultURL, err := url.Parse(cloud.KeyVaultEndpoint)
 	if err != nil {
 		return nil, err
diff --git a/builder/azure/arm/builder.go b/builder/azure/arm/builder.go
index cda50d861..911e59157 100644
--- a/builder/azure/arm/builder.go
+++ b/builder/azure/arm/builder.go
@@ -16,6 +16,7 @@ import (
 	"github.com/dgrijalva/jwt-go"
 	"github.com/hashicorp/hcl/v2/hcldec"
 	packerAzureCommon "github.com/hashicorp/packer/builder/azure/common"
+	"github.com/hashicorp/packer/builder/azure/common/client"
 	"github.com/hashicorp/packer/builder/azure/common/constants"
 	"github.com/hashicorp/packer/builder/azure/common/lin"
 	packerCommon "github.com/hashicorp/packer/common"
@@ -110,6 +111,18 @@ func (b *Builder) Run(ctx context.Context, ui packer.Ui, hook packer.Hook) (pack
 		return nil, fmt.Errorf("could not determine the ObjectID for the user, which is required for Windows builds")
 	}
 
+	if len(b.config.UserAssignedManagedIdentities) != 0 {
+		for _, rid := range b.config.UserAssignedManagedIdentities {
+			r, err := client.ParseResourceID(rid)
+			if err != nil {
+				return nil, err
+			}
+			if _, err = azureClient.UserAssignedIdentitiesClient.Get(ctx, r.ResourceGroup, r.ResourceName.String()); err != nil {
+				return nil, fmt.Errorf("Cannot locate user assigned managed identity %s", rid)
+			}
+		}
+	}
+
 	if b.config.isManagedImage() {
 		_, err := azureClient.GroupsClient.Get(ctx, b.config.ManagedImageResourceGroupName)
 		if err != nil {