Add datadisks to disk set when using SIG as source

builder/azure/chroot/diskset.go

@@ -4,7 +4,7 @@ import "github.com/hashicorp/packer/builder/azure/common/client"
 
 // Diskset represents all of the disks or snapshots associated with an image.
 // It maps lun to resource ids. The OS disk is stored with lun=-1.
-type Diskset map[int]client.Resource
+type Diskset map[int32]client.Resource
 
 // OS return the OS disk resource ID or nil if it is not assigned
 func (ds Diskset) OS() *client.Resource {
@@ -15,7 +15,7 @@ func (ds Diskset) OS() *client.Resource {
 }
 
 // Data return the data disk resource ID or nil if it is not assigned
-func (ds Diskset) Data(lun int) *client.Resource {
+func (ds Diskset) Data(lun int32) *client.Resource {
 	if r, ok := ds[lun]; ok {
 		return &r
 	}

builder/azure/chroot/diskset_test.go

@@ -10,7 +10,7 @@ func diskset(ids ...string) Diskset {
 		if err != nil {
 			panic(err)
 		}
-		diskset[i-1] = r
+		diskset[int32(i-1)] = r
 	}
 	return diskset
 }

builder/azure/chroot/step_create_new_diskset.go

@@ -22,6 +22,8 @@ type StepCreateNewDiskset struct {
 	OSDiskSizeGB             int32  // optional, ignored if 0
 	OSDiskStorageAccountType string // from compute.DiskStorageAccountTypes
 
+	DataDiskIDPrefix string
+
 	disks Diskset
 
 	HyperVGeneration string // For OS disk
@@ -43,10 +45,10 @@ func (s *StepCreateNewDiskset) Run(ctx context.Context, state multistep.StateBag
 	azcli := state.Get("azureclient").(client.AzureClientSet)
 	ui := state.Get("ui").(packer.Ui)
 
-	s.disks = make(map[int]client.Resource)
+	s.disks = make(Diskset)
 
 	errorMessage := func(format string, params ...interface{}) multistep.StepAction {
-		err := fmt.Errorf("StepCreateNewDisk.Run: error: "+format, params...)
+		err := fmt.Errorf("StepCreateNewDiskset.Run: error: "+format, params...)
 		state.Put("error", err)
 		ui.Error(err.Error())
 		return multistep.ActionHalt
@@ -63,7 +65,7 @@ func (s *StepCreateNewDiskset) Run(ctx context.Context, state multistep.StateBag
 	}
 
 	// transform step config to disk model
-	disk := s.getOSDiskDefinition(azcli)
+	disk := s.getOSDiskDefinition(azcli.SubscriptionID())
 
 	// Initiate disk creation
 	f, err := azcli.DisksClient().CreateOrUpdate(ctx, osDisk.ResourceGroup, osDisk.ResourceName.String(), disk)
@@ -74,21 +76,76 @@ func (s *StepCreateNewDiskset) Run(ctx context.Context, state multistep.StateBag
 	state.Put(stateBagKey_Diskset, s.disks) // update the statebag
 	ui.Say(fmt.Sprintf("Creating disk %q", s.OSDiskID))
 
+	type Future struct {
+		client.Resource
+		compute.DisksCreateOrUpdateFuture
+	}
+	futures := []Future{{osDisk, f}}
+
+	if s.SourceImageResourceID != "" {
+		datadiskSuffix := 0 // initialize
+
+		// retrieve image to see if there are any datadisks
+		imageID, err := client.ParseResourceID(s.SourceImageResourceID)
+		if err != nil {
+			return errorMessage("could not parse source image id %q: %v", s.SourceImageResourceID, err)
+		}
+		if !strings.EqualFold(imageID.Provider+"/"+imageID.ResourceType.String(),
+			"Microsoft.Compute/galleries/images/versions") {
+			return errorMessage("source image id is not a shared image version %q, expected type 'Microsoft.Compute/galleries/images/versions'", imageID)
+		}
+		image, err := azcli.GalleryImageVersionsClient().Get(ctx,
+			imageID.ResourceGroup,
+			imageID.ResourceName[0], imageID.ResourceName[1], imageID.ResourceName[2], "")
+		if err != nil {
+			return errorMessage("error retrieving source image %q: %v", imageID, err)
+		}
+		if image.GalleryImageVersionProperties != nil &&
+			image.GalleryImageVersionProperties.StorageProfile != nil &&
+			image.GalleryImageVersionProperties.StorageProfile.DataDiskImages != nil {
+			for i, ddi := range *image.GalleryImageVersionProperties.StorageProfile.DataDiskImages {
+				if ddi.Lun == nil {
+					return errorMessage("unexpected: lun is null for data disk # %d", i)
+				}
+				datadiskID, err := client.ParseResourceID(fmt.Sprintf("%s%d", s.DataDiskIDPrefix, datadiskSuffix))
+				datadiskSuffix++
+				if err != nil {
+					return errorMessage("unable to construct resource id for datadisk: %v", err)
+				}
+
+				disk := s.getDatadiskDefinitionFromImage(*ddi.Lun)
+				// Initiate disk creation
+				f, err := azcli.DisksClient().CreateOrUpdate(ctx, datadiskID.ResourceGroup, datadiskID.ResourceName.String(), disk)
+				if err != nil {
+					return errorMessage("Failed to initiate resource creation: %q", datadiskID)
+				}
+				s.disks[*ddi.Lun] = datadiskID          // save the resoure we just create in our disk set
+				state.Put(stateBagKey_Diskset, s.disks) // update the statebag
+				ui.Say(fmt.Sprintf("Creating disk %q", datadiskID))
+
+				futures = append(futures, Future{datadiskID, f})
+			}
+		}
+	}
+
+	ui.Say("Waiting for disks to be created.")
+
 	// Wait for completion
-	{
+	for _, f := range futures {
 		cli := azcli.PollClient() // quick polling for quick operations
 		cli.PollingDelay = time.Second
 		err = f.WaitForCompletionRef(ctx, cli)
 		if err != nil {
 			return errorMessage(
-				"error creating new disk '%s': %v", s.OSDiskID, err)
+				"error creating new disk '%s': %v", f.Resource, err)
 		}
+		ui.Say(fmt.Sprintf("Disk %q created", f.Resource))
 	}
 
 	return multistep.ActionContinue
 }
 
-func (s *StepCreateNewDiskset) getOSDiskDefinition(azcli client.AzureClientSet) compute.Disk {
+func (s StepCreateNewDiskset) getOSDiskDefinition(subscriptionID string) compute.Disk {
 	disk := compute.Disk{
 		Location: to.StringPtr(s.Location),
 		DiskProperties: &compute.DiskProperties{
@@ -117,7 +174,7 @@ func (s *StepCreateNewDiskset) getOSDiskDefinition(azcli client.AzureClientSet)
 		disk.CreationData.ImageReference = &compute.ImageDiskReference{
 			ID: to.StringPtr(fmt.Sprintf(
 				"/subscriptions/%s/providers/Microsoft.Compute/locations/%s/publishers/%s/artifacttypes/vmimage/offers/%s/skus/%s/versions/%s",
-				azcli.SubscriptionID(), s.Location,
+				subscriptionID, s.Location,
 				s.SourcePlatformImage.Publisher, s.SourcePlatformImage.Offer, s.SourcePlatformImage.Sku, s.SourcePlatformImage.Version)),
 		}
 	case s.SourceOSDiskResourceID != "":
@@ -134,6 +191,28 @@ func (s *StepCreateNewDiskset) getOSDiskDefinition(azcli client.AzureClientSet)
 	return disk
 }
 
+func (s StepCreateNewDiskset) getDatadiskDefinitionFromImage(lun int32) compute.Disk {
+	disk := compute.Disk{
+		Location: to.StringPtr(s.Location),
+		DiskProperties: &compute.DiskProperties{
+			CreationData: &compute.CreationData{},
+		},
+	}
+
+	disk.CreationData.CreateOption = compute.FromImage
+	disk.CreationData.GalleryImageReference = &compute.ImageDiskReference{
+		ID:  to.StringPtr(s.SourceImageResourceID),
+		Lun: to.Int32Ptr(lun),
+	}
+
+	if s.OSDiskStorageAccountType != "" {
+		disk.Sku = &compute.DiskSku{
+			Name: compute.DiskStorageAccountTypes(s.OSDiskStorageAccountType),
+		}
+	}
+	return disk
+}
+
 func (s *StepCreateNewDiskset) Cleanup(state multistep.StateBag) {
 	if !s.SkipCleanup {
 		azcli := state.Get("azureclient").(client.AzureClientSet)
@@ -154,7 +233,7 @@ func (s *StepCreateNewDiskset) Cleanup(state multistep.StateBag) {
 				err = f.WaitForCompletionRef(context.TODO(), azcli.PollClient())
 			}
 			if err != nil {
-				log.Printf("StepCreateNewDisk.Cleanup: error: %+v", err)
+				log.Printf("StepCreateNewDiskset.Cleanup: error: %+v", err)
 				ui.Error(fmt.Sprintf("error deleting disk '%s': %v.", d, err))
 			}
 		}

builder/azure/chroot/step_create_new_diskset_test.go

@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"reflect"
 	"regexp"
+	"strings"
 	"testing"
 
 	"github.com/hashicorp/packer/builder/azure/common/client"
@@ -17,33 +18,24 @@ import (
 )
 
 func TestStepCreateNewDisk_Run(t *testing.T) {
-	type fields struct {
-		ResourceID             string
-		DiskSizeGB             int32
-		DiskStorageAccountType string
-		HyperVGeneration       string
-		Location               string
-		PlatformImage          *client.PlatformImage
-		SourceDiskResourceID   string
-
-		expectedPutDiskBody string
-	}
 	tests := []struct {
-		name   string
+		name                  string
-		fields fields
+		fields                StepCreateNewDiskset
-		want   multistep.StepAction
+		expectedPutDiskBodies []string
+		want                  multistep.StepAction
+		verifyDiskset         *Diskset
 	}{
 		{
 			name: "from disk",
-			fields: fields{
+			fields: StepCreateNewDiskset{
-				ResourceID:             "/subscriptions/SubscriptionID/resourcegroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName",
+				OSDiskID:                 "/subscriptions/SubscriptionID/resourcegroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName",
-				DiskSizeGB:             42,
+				OSDiskSizeGB:             42,
-				DiskStorageAccountType: string(compute.PremiumLRS),
+				OSDiskStorageAccountType: string(compute.PremiumLRS),
-				HyperVGeneration:       string(compute.V1),
+				HyperVGeneration:         string(compute.V1),
-				Location:               "westus",
+				Location:                 "westus",
-				SourceDiskResourceID:   "SourceDisk",
+				SourceOSDiskResourceID:   "SourceDisk",
-
+			},
-				expectedPutDiskBody: `
+			expectedPutDiskBodies: []string{`
 				{
 					"location": "westus",
 					"properties": {
@@ -58,25 +50,25 @@ func TestStepCreateNewDisk_Run(t *testing.T) {
 					"sku": {
 						"name": "Premium_LRS"
 					}
-				}`,
+				}`},
-			},
+			want:          multistep.ActionContinue,
-			want: multistep.ActionContinue,
+			verifyDiskset: &Diskset{-1: resource("/subscriptions/SubscriptionID/resourceGroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName")},
 		},
 		{
-			name: "from image",
+			name: "from platform image",
-			fields: fields{
+			fields: StepCreateNewDiskset{
-				ResourceID:             "/subscriptions/SubscriptionID/resourcegroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName",
+				OSDiskID:                 "/subscriptions/SubscriptionID/resourcegroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName",
-				DiskStorageAccountType: string(compute.StandardLRS),
+				OSDiskStorageAccountType: string(compute.StandardLRS),
-				HyperVGeneration:       string(compute.V1),
+				HyperVGeneration:         string(compute.V1),
-				Location:               "westus",
+				Location:                 "westus",
-				PlatformImage: &client.PlatformImage{
+				SourcePlatformImage: &client.PlatformImage{
 					Publisher: "Microsoft",
 					Offer:     "Windows",
 					Sku:       "2016-DataCenter",
 					Version:   "2016.1.4",
 				},
-
+			},
-				expectedPutDiskBody: `
+			expectedPutDiskBodies: []string{`
 				{
 					"location": "westus",
 					"properties": {
@@ -92,33 +84,107 @@ func TestStepCreateNewDisk_Run(t *testing.T) {
 					"sku": {
 						"name": "Standard_LRS"
 					}
-				}`,
+				}`},
+			want:          multistep.ActionContinue,
+			verifyDiskset: &Diskset{-1: resource("/subscriptions/SubscriptionID/resourceGroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName")},
+		},
+		{
+			name: "from shared image",
+			fields: StepCreateNewDiskset{
+				OSDiskID:                 "/subscriptions/SubscriptionID/resourcegroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName",
+				OSDiskStorageAccountType: string(compute.StandardLRS),
+				DataDiskIDPrefix:         "/subscriptions/SubscriptionID/resourcegroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryDataDisk-",
+				HyperVGeneration:         string(compute.V1),
+				Location:                 "westus",
+				SourceImageResourceID:    "/subscriptions/SubscriptionID/resourcegroups/imagegroup/providers/Microsoft.Compute/galleries/MyGallery/images/MyImage/versions/1.2.3",
 			},
+
+			expectedPutDiskBodies: []string{`
+				{
+					"location": "westus",
+					"properties": {
+						"osType": "Linux",
+						"hyperVGeneration": "V1",
+						"creationData": {
+							"createOption":"FromImage",
+							"galleryImageReference": {
+								"id":"/subscriptions/SubscriptionID/resourcegroups/imagegroup/providers/Microsoft.Compute/galleries/MyGallery/images/MyImage/versions/1.2.3"
+							}
+						}
+					},
+					"sku": {
+						"name": "Standard_LRS"
+					}
+				}`, `
+				{
+					"location": "westus",
+					"properties": {
+						"creationData": {
+							"createOption":"FromImage",
+							"galleryImageReference": {
+								"id": "/subscriptions/SubscriptionID/resourcegroups/imagegroup/providers/Microsoft.Compute/galleries/MyGallery/images/MyImage/versions/1.2.3",
+								"lun": 5
+							}
+						}
+					},
+					"sku": {
+						"name": "Standard_LRS"
+					}
+				}`, `
+				{
+					"location": "westus",
+					"properties": {
+						"creationData": {
+							"createOption":"FromImage",
+							"galleryImageReference": {
+								"id": "/subscriptions/SubscriptionID/resourcegroups/imagegroup/providers/Microsoft.Compute/galleries/MyGallery/images/MyImage/versions/1.2.3",
+								"lun": 9
+							}
+						}
+					},
+					"sku": {
+						"name": "Standard_LRS"
+					}
+				}`, `
+				{
+					"location": "westus",
+					"properties": {
+						"creationData": {
+							"createOption":"FromImage",
+							"galleryImageReference": {
+								"id": "/subscriptions/SubscriptionID/resourcegroups/imagegroup/providers/Microsoft.Compute/galleries/MyGallery/images/MyImage/versions/1.2.3",
+								"lun": 3
+							}
+						}
+					},
+					"sku": {
+						"name": "Standard_LRS"
+					}
+				}`},
 			want: multistep.ActionContinue,
+			verifyDiskset: &Diskset{
+				-1: resource("/subscriptions/SubscriptionID/resourceGroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName"),
+				3:  resource("/subscriptions/SubscriptionID/resourceGroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryDataDisk-2"),
+				5:  resource("/subscriptions/SubscriptionID/resourceGroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryDataDisk-0"),
+				9:  resource("/subscriptions/SubscriptionID/resourceGroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryDataDisk-1"),
+			},
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			s := StepCreateNewDiskset{
+			s := tt.fields
-				OSDiskID:                 tt.fields.ResourceID,
-				OSDiskSizeGB:             tt.fields.DiskSizeGB,
-				OSDiskStorageAccountType: tt.fields.DiskStorageAccountType,
-				HyperVGeneration:         tt.fields.HyperVGeneration,
-				Location:                 tt.fields.Location,
-				SourcePlatformImage:      tt.fields.PlatformImage,
-				SourceOSDiskResourceID:   tt.fields.SourceDiskResourceID,
-			}
 
-			expectedPutDiskBody := regexp.MustCompile(`[\s\n]`).ReplaceAllString(tt.fields.expectedPutDiskBody, "")
+			bodyCount := 0
-
+			m := compute.NewDisksClient("SubscriptionID")
-			m := compute.NewDisksClient("subscriptionId")
 			m.Sender = autorest.SenderFunc(func(r *http.Request) (*http.Response, error) {
 				if r.Method != "PUT" {
 					t.Fatal("Expected only a PUT disk call")
 				}
 				b, _ := ioutil.ReadAll(r.Body)
+				expectedPutDiskBody := regexp.MustCompile(`[\s\n]`).ReplaceAllString(tt.expectedPutDiskBodies[bodyCount], "")
+				bodyCount++
 				if string(b) != expectedPutDiskBody {
-					t.Fatalf("expected body to be %q, but got %q", expectedPutDiskBody, string(b))
+					t.Fatalf("expected body #%d to be %q, but got %q", bodyCount, expectedPutDiskBody, string(b))
 				}
 				return &http.Response{
 					Request:    r,
@@ -126,16 +192,55 @@ func TestStepCreateNewDisk_Run(t *testing.T) {
 				}, nil
 			})
 
+			giv := compute.NewGalleryImageVersionsClient("SubscriptionID")
+			giv.Sender = autorest.SenderFunc(func(r *http.Request) (*http.Response, error) {
+				if r.Method == "GET" &&
+					regexp.MustCompile(`(?i)/versions/1\.2\.3$`).MatchString(r.URL.Path) {
+					return &http.Response{
+						Request: r,
+						Body: ioutil.NopCloser(strings.NewReader(`{
+							"properties": { "storageProfile": {
+								"dataDiskImages":[
+									{ "lun": 5 },
+									{ "lun": 9 },
+									{ "lun": 3 }
+								]
+							} }
+						}`)),
+						StatusCode: 200,
+					}, nil
+				}
+				return &http.Response{
+					Request:    r,
+					Status:     "Unexpected request",
+					StatusCode: 500,
+				}, nil
+			})
+
 			state := new(multistep.BasicStateBag)
 			state.Put("azureclient", &client.AzureClientSetMock{
-				SubscriptionIDMock: "SubscriptionID",
+				SubscriptionIDMock:             "SubscriptionID",
-				DisksClientMock:    m,
+				DisksClientMock:                m,
+				GalleryImageVersionsClientMock: giv,
 			})
 			state.Put("ui", packer.TestUi(t))
 
 			if got := s.Run(context.TODO(), state); !reflect.DeepEqual(got, tt.want) {
 				t.Errorf("StepCreateNewDisk.Run() = %v, want %v", got, tt.want)
 			}
+
+			ds := state.Get(stateBagKey_Diskset)
+			if tt.verifyDiskset != nil && !reflect.DeepEqual(*tt.verifyDiskset, ds) {
+				t.Errorf("Error verifying diskset after Run(), got %v, want %v", ds, *&tt.verifyDiskset)
+			}
 		})
 	}
 }
+
+func resource(id string) client.Resource {
+	v, err := client.ParseResourceID(id)
+	if err != nil {
+		panic(err)
+	}
+	return v
+}