operator: support getting MP node image

edgelesssys · Jan 23, 2024 · 9c71c4d · 9c71c4d
1 parent a0a9281
commit 9c71c4d
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 1 deletion.
diff --git a/operators/constellation-node-operator/internal/cloud/azure/client/nodeimage.go b/operators/constellation-node-operator/internal/cloud/azure/client/nodeimage.go
@@ -12,6 +12,7 @@ import (
 	"strings"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v5"
+	"github.com/edgelesssys/constellation/v2/internal/mpimage"
 )
 
 // GetNodeImage returns the image name of the node.
@@ -27,12 +28,34 @@ func (c *Client) GetNodeImage(ctx context.Context, providerID string) (string, e
 	if resp.Properties == nil ||
 		resp.Properties.StorageProfile == nil ||
 		resp.Properties.StorageProfile.ImageReference == nil ||
-		resp.Properties.StorageProfile.ImageReference.ID == nil && resp.Properties.StorageProfile.ImageReference.CommunityGalleryImageID == nil {
+		resp.Properties.StorageProfile.ImageReference.ID == nil &&
+			resp.Properties.StorageProfile.ImageReference.CommunityGalleryImageID == nil &&
+			resp.Properties.StorageProfile.ImageReference.Publisher == nil &&
+			resp.Properties.StorageProfile.ImageReference.Offer == nil &&
+			resp.Properties.StorageProfile.ImageReference.SKU == nil &&
+			resp.Properties.StorageProfile.ImageReference.Version == nil {
 		return "", fmt.Errorf("node %q does not have valid image reference", providerID)
 	}
+
+	// Marketplace Image is used, format it to an URI and return it.
+	if resp.Properties.StorageProfile.ImageReference.Publisher != nil &&
+		resp.Properties.StorageProfile.ImageReference.Offer != nil &&
+		resp.Properties.StorageProfile.ImageReference.SKU != nil &&
+		resp.Properties.StorageProfile.ImageReference.Version != nil {
+		return mpimage.AzureMarketplaceImage{
+			Publisher: *resp.Properties.StorageProfile.ImageReference.Publisher,
+			Offer:     *resp.Properties.StorageProfile.ImageReference.Offer,
+			SKU:       *resp.Properties.StorageProfile.ImageReference.SKU,
+			Version:   *resp.Properties.StorageProfile.ImageReference.Version,
+		}.URI(), nil
+	}
+
+	// Image ID is set, return it.
 	if resp.Properties.StorageProfile.ImageReference.ID != nil {
 		return *resp.Properties.StorageProfile.ImageReference.ID, nil
 	}
+
+	// Community Gallery image ID is set otherwise.
 	return *resp.Properties.StorageProfile.ImageReference.CommunityGalleryImageID, nil
 }
 

diff --git a/operators/constellation-node-operator/internal/cloud/azure/client/nodeimage_test.go b/operators/constellation-node-operator/internal/cloud/azure/client/nodeimage_test.go
@@ -53,6 +53,22 @@ func TestGetNodeImage(t *testing.T) {
 			},
 			wantImage: "/CommunityGalleries/gallery-name/Images/image-name/Versions/1.2.3",
 		},
+		"getting marketplace image works": {
+			providerID: "azure:///subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/scale-set-name/virtualMachines/instance-id",
+			vm: armcompute.VirtualMachineScaleSetVM{
+				Properties: &armcompute.VirtualMachineScaleSetVMProperties{
+					StorageProfile: &armcompute.StorageProfile{
+						ImageReference: &armcompute.ImageReference{
+							Publisher: to.Ptr("edgelesssystems"),
+							Offer:     to.Ptr("constellation"),
+							SKU:       to.Ptr("constellation"),
+							Version:   to.Ptr("2.14.2"),
+						},
+					},
+				},
+			},
+			wantImage: "constellation-marketplace-image://Azure?offer=constellation&publisher=edgelesssystems&sku=constellation&version=2.14.2",
+		},
 		"splitting providerID fails": {
 			providerID: "invalid",
 			wantErr:    true,