From da807847a237d566894ce434bbe49af476a83cf8 Mon Sep 17 00:00:00 2001
From: Paul Meyer <paulmey@microsoft.com>
Date: Mon, 3 Jun 2019 23:06:19 +0000
Subject: [PATCH] Allow skipping OSDisk cleanup

---
 builder/azure/chroot/builder.go              |  3 ++
 builder/azure/chroot/step_create_new_disk.go | 30 +++++++++++---------
 2 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/builder/azure/chroot/builder.go b/builder/azure/chroot/builder.go
index 75e195add72..674e03ed039 100644
--- a/builder/azure/chroot/builder.go
+++ b/builder/azure/chroot/builder.go
@@ -42,6 +42,7 @@ type Config struct {
 	OSDiskSizeGB             int32  `mapstructure:"os_disk_size_gb"`
 	OSDiskStorageAccountType string `mapstructure:"os_disk_storage_account_type"`
 	OSDiskCacheType          string `mapstructure:"os_disk_cache_type"`
+	OSDiskSkipCleanup        bool   `mapstructure:"os_disk_skip_cleanup"`
 
 	ImageResourceID       string `mapstructure:"image_resource_id"`
 	ImageOSState          string `mapstructure:"image_os_state"`
@@ -319,6 +320,8 @@ func (b *Builder) Run(ctx context.Context, ui packer.Ui, hook packer.Hook) (pack
 					HyperVGeneration:       b.config.ImageHyperVGeneration,
 					Location:               info.Location,
 					PlatformImage:          pi,
+
+					SkipCleanup: b.config.OSDiskSkipCleanup,
 				})
 		} else {
 			panic("Unknown image source: " + b.config.Source)
diff --git a/builder/azure/chroot/step_create_new_disk.go b/builder/azure/chroot/step_create_new_disk.go
index 407e71bde7a..2fd9928e518 100644
--- a/builder/azure/chroot/step_create_new_disk.go
+++ b/builder/azure/chroot/step_create_new_disk.go
@@ -21,6 +21,8 @@ type StepCreateNewDisk struct {
 	HyperVGeneration                        string
 	Location                                string
 	PlatformImage                           *client.PlatformImage
+
+	SkipCleanup bool
 }
 
 func (s StepCreateNewDisk) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction {
@@ -76,21 +78,23 @@ func (s StepCreateNewDisk) Run(ctx context.Context, state multistep.StateBag) mu
 }
 
 func (s StepCreateNewDisk) Cleanup(state multistep.StateBag) {
-	azcli := state.Get("azureclient").(client.AzureClientSet)
-	ui := state.Get("ui").(packer.Ui)
-	diskResourceID := state.Get("os_disk_resource_id").(string)
+	if !s.SkipCleanup {
+		azcli := state.Get("azureclient").(client.AzureClientSet)
+		ui := state.Get("ui").(packer.Ui)
+		diskResourceID := state.Get("os_disk_resource_id").(string)
 
-	ui.Say(fmt.Sprintf("Waiting for disk %q detach to complete", diskResourceID))
-	err := NewDiskAttacher(azcli).WaitForDetach(context.Background(), diskResourceID)
+		ui.Say(fmt.Sprintf("Waiting for disk %q detach to complete", diskResourceID))
+		err := NewDiskAttacher(azcli).WaitForDetach(context.Background(), diskResourceID)
 
-	ui.Say(fmt.Sprintf("Deleting disk %q", diskResourceID))
+		ui.Say(fmt.Sprintf("Deleting disk %q", diskResourceID))
 
-	f, err := azcli.DisksClient().Delete(context.TODO(), s.ResourceGroup, s.DiskName)
-	if err == nil {
-		err = f.WaitForCompletionRef(context.TODO(), azcli.PollClient())
-	}
-	if err != nil {
-		log.Printf("StepCreateNewDisk.Cleanup: error: %+v", err)
-		ui.Error(fmt.Sprintf("error deleting new disk '%s': %v.", diskResourceID, err))
+		f, err := azcli.DisksClient().Delete(context.TODO(), s.ResourceGroup, s.DiskName)
+		if err == nil {
+			err = f.WaitForCompletionRef(context.TODO(), azcli.PollClient())
+		}
+		if err != nil {
+			log.Printf("StepCreateNewDisk.Cleanup: error: %+v", err)
+			ui.Error(fmt.Sprintf("error deleting new disk '%s': %v.", diskResourceID, err))
+		}
 	}
 }