operator: fix get-handling of Azure marketplace images (#2846)

* operator: support getting MP images * operator: support getting MP node image * operator: refactorings
edgelesssys · Jan 24, 2024 · e07ea4b · e07ea4b
1 parent da26dae
commit e07ea4b
Show file tree

Hide file tree

Showing 4 changed files with 78 additions and 4 deletions.
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,13 +28,32 @@ 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)
 	}
+
+	// Image ID is set, return it.
 	if resp.Properties.StorageProfile.ImageReference.ID != nil {
 		return *resp.Properties.StorageProfile.ImageReference.ID, nil
 	}
-	return *resp.Properties.StorageProfile.ImageReference.CommunityGalleryImageID, nil
+
+	// Community Gallery image ID is set, return it.
+	if resp.Properties.StorageProfile.ImageReference.CommunityGalleryImageID != nil {
+		return *resp.Properties.StorageProfile.ImageReference.CommunityGalleryImageID, nil
+	}
+
+	// Last possible option: Marketplace Image is used, format it to an URI and return it.
+	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
 }
 
 // GetScalingGroupID returns the scaling group ID of the node.

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,

diff --git a/operators/constellation-node-operator/internal/cloud/azure/client/scalinggroup.go b/operators/constellation-node-operator/internal/cloud/azure/client/scalinggroup.go
@@ -29,17 +29,37 @@ func (c *Client) GetScalingGroupImage(ctx context.Context, scalingGroupID string
 	if err != nil {
 		return "", err
 	}
+
 	if res.Properties == nil ||
 		res.Properties.VirtualMachineProfile == nil ||
 		res.Properties.VirtualMachineProfile.StorageProfile == nil ||
 		res.Properties.VirtualMachineProfile.StorageProfile.ImageReference == nil ||
-		res.Properties.VirtualMachineProfile.StorageProfile.ImageReference.ID == nil && res.Properties.VirtualMachineProfile.StorageProfile.ImageReference.CommunityGalleryImageID == nil {
+		res.Properties.VirtualMachineProfile.StorageProfile.ImageReference.ID == nil &&
+			res.Properties.VirtualMachineProfile.StorageProfile.ImageReference.CommunityGalleryImageID == nil &&
+			(res.Properties.VirtualMachineProfile.StorageProfile.ImageReference.Publisher == nil ||
+				res.Properties.VirtualMachineProfile.StorageProfile.ImageReference.Offer == nil ||
+				res.Properties.VirtualMachineProfile.StorageProfile.ImageReference.SKU == nil ||
+				res.Properties.VirtualMachineProfile.StorageProfile.ImageReference.Version == nil) {
 		return "", fmt.Errorf("scalet set %q does not have valid image reference", scalingGroupID)
 	}
+
+	// Image ID is set, return it.
 	if res.Properties.VirtualMachineProfile.StorageProfile.ImageReference.ID != nil {
 		return *res.Properties.VirtualMachineProfile.StorageProfile.ImageReference.ID, nil
 	}
-	return *res.Properties.VirtualMachineProfile.StorageProfile.ImageReference.CommunityGalleryImageID, nil
+
+	// Community Gallery Image ID is set, return it.
+	if res.Properties.VirtualMachineProfile.StorageProfile.ImageReference.CommunityGalleryImageID != nil {
+		return *res.Properties.VirtualMachineProfile.StorageProfile.ImageReference.CommunityGalleryImageID, nil
+	}
+
+	// Last possible option: Marketplace Image is used, format it to an URI and return it.
+	return mpimage.AzureMarketplaceImage{
+		Publisher: *res.Properties.VirtualMachineProfile.StorageProfile.ImageReference.Publisher,
+		Offer:     *res.Properties.VirtualMachineProfile.StorageProfile.ImageReference.Offer,
+		SKU:       *res.Properties.VirtualMachineProfile.StorageProfile.ImageReference.SKU,
+		Version:   *res.Properties.VirtualMachineProfile.StorageProfile.ImageReference.Version,
+	}.URI(), nil
 }
 
 // SetScalingGroupImage sets the image URI of the scaling group.

diff --git a/operators/constellation-node-operator/internal/cloud/azure/client/scalinggroup_test.go b/operators/constellation-node-operator/internal/cloud/azure/client/scalinggroup_test.go
@@ -57,6 +57,24 @@ func TestGetScalingGroupImage(t *testing.T) {
 			},
 			wantImage: "/communityGalleries/gallery-name/Images/image-name/Versions/1.2.3",
 		},
+		"getting marketplace image works": {
+			scalingGroupID: "/subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/scale-set-name",
+			scaleSet: armcompute.VirtualMachineScaleSet{
+				Properties: &armcompute.VirtualMachineScaleSetProperties{
+					VirtualMachineProfile: &armcompute.VirtualMachineScaleSetVMProfile{
+						StorageProfile: &armcompute.VirtualMachineScaleSetStorageProfile{
+							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 scalingGroupID fails": {
 			scalingGroupID: "invalid",
 			wantErr:        true,