From c2ecdd98c6a5987e7af27acd1064e12cbee48a27 Mon Sep 17 00:00:00 2001
From: Arjen Schwarz <arjen.schwarz@bulletproof.net>
Date: Tue, 12 Dec 2017 20:12:53 +1100
Subject: [PATCH] 5691: Invalid image URLs make Azure builder crash

---
 builder/azure/arm/step_delete_os_disk.go      | 12 ++++--
 builder/azure/arm/step_delete_os_disk_test.go | 40 +++++++++++++++++++
 builder/azure/arm/step_deploy_template.go     |  4 ++
 .../azure/arm/step_deploy_template_test.go    | 25 ++++++++++++
 4 files changed, 77 insertions(+), 4 deletions(-)

diff --git a/builder/azure/arm/step_delete_os_disk.go b/builder/azure/arm/step_delete_os_disk.go
index af5e2fa9b..878c45cc6 100644
--- a/builder/azure/arm/step_delete_os_disk.go
+++ b/builder/azure/arm/step_delete_os_disk.go
@@ -1,6 +1,7 @@
 package arm
 
 import (
+	"errors"
 	"fmt"
 	"net/url"
 	"strings"
@@ -80,11 +81,14 @@ func (s *StepDeleteOSDisk) Run(state multistep.StateBag) multistep.StepAction {
 	}
 
 	xs := strings.Split(u.Path, "/")
+	if len(xs) < 3 {
+		err = errors.New("Failed to parse OS Disk's VHD URI!")
+	} else {
+		var storageAccountName = xs[1]
+		var blobName = strings.Join(xs[2:], "/")
 
-	var storageAccountName = xs[1]
-	var blobName = strings.Join(xs[2:], "/")
-
-	err = s.delete(storageAccountName, blobName)
+		err = s.delete(storageAccountName, blobName)
+	}
 	return processStepResult(err, s.error, state)
 }
 
diff --git a/builder/azure/arm/step_delete_os_disk_test.go b/builder/azure/arm/step_delete_os_disk_test.go
index 7fdd7468e..5a962dd8a 100644
--- a/builder/azure/arm/step_delete_os_disk_test.go
+++ b/builder/azure/arm/step_delete_os_disk_test.go
@@ -104,6 +104,46 @@ func TestStepDeleteOSDiskShouldHandleComplexStorageContainerNames(t *testing.T)
 	}
 }
 
+func TestStepDeleteOSDiskShouldFailIfVHDNameCannotBeURLParsed(t *testing.T) {
+	var testSubject = &StepDeleteOSDisk{
+		delete:        func(string, string) error { return nil },
+		say:           func(message string) {},
+		error:         func(e error) {},
+		deleteManaged: func(string, string) error { return nil },
+	}
+
+	// Invalid URL per https://golang.org/src/net/url/url_test.go
+	stateBag := DeleteTestStateBagStepDeleteOSDisk("http://[fe80::1%en0]/")
+
+	var result = testSubject.Run(stateBag)
+	if result != multistep.ActionHalt {
+		t.Fatalf("Expected the step to return 'ActionHalt', but got '%v'.", result)
+	}
+
+	if _, ok := stateBag.GetOk(constants.Error); ok == false {
+		t.Fatalf("Expected the step to not stateBag['%s'], but it was.", constants.Error)
+	}
+}
+func TestStepDeleteOSDiskShouldFailIfVHDNameIsTooShort(t *testing.T) {
+	var testSubject = &StepDeleteOSDisk{
+		delete:        func(string, string) error { return nil },
+		say:           func(message string) {},
+		error:         func(e error) {},
+		deleteManaged: func(string, string) error { return nil },
+	}
+
+	stateBag := DeleteTestStateBagStepDeleteOSDisk("storage.blob.core.windows.net/abc")
+
+	var result = testSubject.Run(stateBag)
+	if result != multistep.ActionHalt {
+		t.Fatalf("Expected the step to return 'ActionHalt', but got '%d'.", result)
+	}
+
+	if _, ok := stateBag.GetOk(constants.Error); ok == false {
+		t.Fatalf("Expected the step to not stateBag['%s'], but it was.", constants.Error)
+	}
+}
+
 func TestStepDeleteOSDiskShouldPassIfManagedDiskInTempResourceGroup(t *testing.T) {
 	var testSubject = &StepDeleteOSDisk{
 		delete: func(string, string) error { return nil },
diff --git a/builder/azure/arm/step_deploy_template.go b/builder/azure/arm/step_deploy_template.go
index b67dd5c1d..5978a8811 100644
--- a/builder/azure/arm/step_deploy_template.go
+++ b/builder/azure/arm/step_deploy_template.go
@@ -1,6 +1,7 @@
 package arm
 
 import (
+	"errors"
 	"fmt"
 	"net/url"
 	"strings"
@@ -140,6 +141,9 @@ func (s *StepDeployTemplate) deleteImage(imageType string, imageName string, res
 		return err
 	}
 	xs := strings.Split(u.Path, "/")
+	if len(xs) < 3 {
+		return errors.New("Unable to parse path of image " + imageName)
+	}
 	var storageAccountName = xs[1]
 	var blobName = strings.Join(xs[2:], "/")
 
diff --git a/builder/azure/arm/step_deploy_template_test.go b/builder/azure/arm/step_deploy_template_test.go
index d812e0ffa..4250e9e0b 100644
--- a/builder/azure/arm/step_deploy_template_test.go
+++ b/builder/azure/arm/step_deploy_template_test.go
@@ -82,6 +82,31 @@ func TestStepDeployTemplateShouldTakeStepArgumentsFromStateBag(t *testing.T) {
 	}
 }
 
+func TestStepDeployTemplateDeleteImageShouldFailWhenImageUrlCannotBeParsed(t *testing.T) {
+	var testSubject = &StepDeployTemplate{
+		say:   func(message string) {},
+		error: func(e error) {},
+		name:  "--deployment-name--",
+	}
+	// Invalid URL per https://golang.org/src/net/url/url_test.go
+	err := testSubject.deleteImage("image", "http://[fe80::1%en0]/", "Unit Test: ResourceGroupName")
+	if err == nil {
+		t.Fatal("Expected a failure because of the failed image name")
+	}
+}
+
+func TestStepDeployTemplateDeleteImageShouldFailWithInvalidImage(t *testing.T) {
+	var testSubject = &StepDeployTemplate{
+		say:   func(message string) {},
+		error: func(e error) {},
+		name:  "--deployment-name--",
+	}
+	err := testSubject.deleteImage("image", "storage.blob.core.windows.net/abc", "Unit Test: ResourceGroupName")
+	if err == nil {
+		t.Fatal("Expected a failure because of the failed image name")
+	}
+}
+
 func createTestStateBagStepDeployTemplate() multistep.StateBag {
 	stateBag := new(multistep.BasicStateBag)