[node operator] self-initialize resources

Signed-off-by: Malte Poll <mp@edgeless.systems>

operators/constellation-node-operator/internal/azure/client/api.go

@@ -30,6 +30,8 @@ type scaleSetsAPI interface {
 	BeginDeleteInstances(ctx context.Context, resourceGroupName string, vmScaleSetName string, vmInstanceIDs armcomputev2.VirtualMachineScaleSetVMInstanceRequiredIDs,
 		options *armcomputev2.VirtualMachineScaleSetsClientBeginDeleteInstancesOptions,
 	) (*runtime.Poller[armcomputev2.VirtualMachineScaleSetsClientDeleteInstancesResponse], error)
+	NewListPager(resourceGroupName string, options *armcomputev2.VirtualMachineScaleSetsClientListOptions,
+	) *runtime.Pager[armcomputev2.VirtualMachineScaleSetsClientListResponse]
 }
 
 type capacityPoller interface {

operators/constellation-node-operator/internal/azure/client/client.go

@@ -9,6 +9,7 @@ import (
 
 // Client is a client for the Azure Cloud.
 type Client struct {
+	config cloudConfig
 	scaleSetsAPI
 	virtualMachineScaleSetVMsAPI
 	capacityPollerGenerator func(resourceGroup, scaleSet string, wantedCapacity int64) capacityPoller
@@ -37,6 +38,7 @@ func NewFromDefault(configPath string) (*Client, error) {
 	}
 
 	return &Client{
+		config:                       *config,
 		scaleSetsAPI:                 scaleSetAPI,
 		virtualMachineScaleSetVMsAPI: virtualMachineScaleSetVMsAPI,
 		capacityPollerGenerator: func(resourceGroup, scaleSet string, wantedCapacity int64) capacityPoller {

operators/constellation-node-operator/internal/azure/client/client_test.go

@@ -16,6 +16,7 @@ type stubScaleSetsAPI struct {
 	deleteResponse armcomputev2.VirtualMachineScaleSetsClientDeleteInstancesResponse
 	deleteErr      error
 	resultErr      error
+	pager          *stubVMSSPager
 }
 
 func (a *stubScaleSetsAPI) Get(ctx context.Context, resourceGroupName string, vmScaleSetName string,
@@ -54,12 +55,20 @@ func (a *stubScaleSetsAPI) BeginDeleteInstances(ctx context.Context, resourceGro
 	return poller, a.deleteErr
 }
 
+func (a *stubScaleSetsAPI) NewListPager(resourceGroupName string, options *armcomputev2.VirtualMachineScaleSetsClientListOptions,
+) *runtime.Pager[armcomputev2.VirtualMachineScaleSetsClientListResponse] {
+	return runtime.NewPager(runtime.PagingHandler[armcomputev2.VirtualMachineScaleSetsClientListResponse]{
+		More:    a.pager.moreFunc(),
+		Fetcher: a.pager.fetcherFunc(),
+	})
+}
+
 type stubvirtualMachineScaleSetVMsAPI struct {
 	scaleSetVM      armcomputev2.VirtualMachineScaleSetVMsClientGetResponse
 	getErr          error
 	instanceView    armcomputev2.VirtualMachineScaleSetVMsClientGetInstanceViewResponse
 	instanceViewErr error
-	pager           *stubPager
+	pager           *stubVMSSVMPager
 }
 
 func (a *stubvirtualMachineScaleSetVMsAPI) Get(ctx context.Context, resourceGroupName string, vmScaleSetName string, instanceID string,
@@ -102,19 +111,19 @@ func (p *stubPoller[T]) Result(ctx context.Context, out *T) error {
 	return p.resultErr
 }
 
-type stubPager struct {
+type stubVMSSVMPager struct {
 	list     []armcomputev2.VirtualMachineScaleSetVM
 	fetchErr error
 	more     bool
 }
 
-func (p *stubPager) moreFunc() func(armcomputev2.VirtualMachineScaleSetVMsClientListResponse) bool {
+func (p *stubVMSSVMPager) moreFunc() func(armcomputev2.VirtualMachineScaleSetVMsClientListResponse) bool {
 	return func(armcomputev2.VirtualMachineScaleSetVMsClientListResponse) bool {
 		return p.more
 	}
 }
 
-func (p *stubPager) fetcherFunc() func(context.Context, *armcomputev2.VirtualMachineScaleSetVMsClientListResponse) (armcomputev2.VirtualMachineScaleSetVMsClientListResponse, error) {
+func (p *stubVMSSVMPager) fetcherFunc() func(context.Context, *armcomputev2.VirtualMachineScaleSetVMsClientListResponse) (armcomputev2.VirtualMachineScaleSetVMsClientListResponse, error) {
 	return func(context.Context, *armcomputev2.VirtualMachineScaleSetVMsClientListResponse) (armcomputev2.VirtualMachineScaleSetVMsClientListResponse, error) {
 		page := make([]*armcomputev2.VirtualMachineScaleSetVM, len(p.list))
 		for i := range p.list {
@@ -127,3 +136,29 @@ func (p *stubPager) fetcherFunc() func(context.Context, *armcomputev2.VirtualMac
 		}, p.fetchErr
 	}
 }
+
+type stubVMSSPager struct {
+	list     []armcomputev2.VirtualMachineScaleSet
+	fetchErr error
+	more     bool
+}
+
+func (p *stubVMSSPager) moreFunc() func(armcomputev2.VirtualMachineScaleSetsClientListResponse) bool {
+	return func(armcomputev2.VirtualMachineScaleSetsClientListResponse) bool {
+		return p.more
+	}
+}
+
+func (p *stubVMSSPager) fetcherFunc() func(context.Context, *armcomputev2.VirtualMachineScaleSetsClientListResponse) (armcomputev2.VirtualMachineScaleSetsClientListResponse, error) {
+	return func(context.Context, *armcomputev2.VirtualMachineScaleSetsClientListResponse) (armcomputev2.VirtualMachineScaleSetsClientListResponse, error) {
+		page := make([]*armcomputev2.VirtualMachineScaleSet, len(p.list))
+		for i := range p.list {
+			page[i] = &p.list[i]
+		}
+		return armcomputev2.VirtualMachineScaleSetsClientListResponse{
+			VirtualMachineScaleSetListResult: armcomputev2.VirtualMachineScaleSetListResult{
+				Value: page,
+			},
+		}, p.fetchErr
+	}
+}

operators/constellation-node-operator/internal/azure/client/config.go

@@ -13,6 +13,7 @@ import (
 type cloudConfig struct {
 	TenantID       string `json:"tenantId,omitempty"`
 	SubscriptionID string `json:"subscriptionId,omitempty"`
+	ResourceGroup  string `json:"resourceGroup,omitempty"`
 }
 
 // loadConfig loads the cloud config from the given path.

operators/constellation-node-operator/internal/azure/client/instanceview_test.go

@@ -11,7 +11,7 @@ import (
 )
 
 // this state is included in most VMs but not needed
-// to determine the node state as every provisioned VM also has a power state
+// to determine the node state as every provisioned VM also has a power state.
 const provisioningStateSucceeded = "ProvisioningState/succeeded"
 
 func TestNodeStateFromStatuses(t *testing.T) {

operators/constellation-node-operator/internal/azure/client/nodeimage_test.go

@@ -199,7 +199,7 @@ func TestCreateNode(t *testing.T) {
 			assert := assert.New(t)
 			require := require.New(t)
 
-			pager := &stubPager{
+			pager := &stubVMSSVMPager{
 				list:     tc.preexistingVMs,
 				fetchErr: tc.fetchErr,
 			}
@@ -325,7 +325,7 @@ func TestCapacityPollingHandler(t *testing.T) {
 	assert.Error(handler.Poll(context.Background()))
 
 	// let Poll finish
-	handler.scaleSetsAPI.(*stubScaleSetsAPI).scaleSet.SKU = &armcomputev2.SKU{Capacity: to.Ptr(int64(wantCapacity))}
+	handler.scaleSetsAPI.(*stubScaleSetsAPI).scaleSet.SKU = &armcomputev2.SKU{Capacity: to.Ptr(wantCapacity)}
 	assert.NoError(handler.Poll(context.Background()))
 	assert.True(handler.Done())
 	assert.NoError(handler.Result(context.Background(), &gotCapacity))

operators/constellation-node-operator/internal/azure/client/scaleset.go

@@ -0,0 +1,26 @@
+package client
+
+import (
+	"context"
+	"fmt"
+)
+
+// getScaleSets retrieves the IDs of all scale sets of a resource group.
+func (c *Client) getScaleSets(ctx context.Context) ([]string, error) {
+	pager := c.scaleSetsAPI.NewListPager(c.config.ResourceGroup, nil)
+	var scaleSets []string
+
+	for pager.More() {
+		page, err := pager.NextPage(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("paging scale sets: %w", err)
+		}
+		for _, scaleSet := range page.Value {
+			if scaleSet == nil || scaleSet.ID == nil {
+				continue
+			}
+			scaleSets = append(scaleSets, *scaleSet.ID)
+		}
+	}
+	return scaleSets, nil
+}

operators/constellation-node-operator/internal/azure/client/scaleset_test.go

@@ -0,0 +1,56 @@
+package client
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
+	armcomputev2 "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v2"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGetScaleSets(t *testing.T) {
+	testCases := map[string]struct {
+		scaleSet      armcomputev2.VirtualMachineScaleSet
+		fetchPageErr  error
+		wantScaleSets []string
+		wantErr       bool
+	}{
+		"fetching scale sets works": {
+			scaleSet: armcomputev2.VirtualMachineScaleSet{
+				ID: to.Ptr("/subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/scale-set-name"),
+			},
+			wantScaleSets: []string{"/subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/scale-set-name"},
+		},
+		"fetching scale sets fails": {
+			fetchPageErr: errors.New("fetch page error"),
+			wantErr:      true,
+		},
+		"scale set is invalid": {},
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			assert := assert.New(t)
+			require := require.New(t)
+
+			client := Client{
+				scaleSetsAPI: &stubScaleSetsAPI{
+					pager: &stubVMSSPager{
+						list:     []armcomputev2.VirtualMachineScaleSet{tc.scaleSet},
+						fetchErr: tc.fetchPageErr,
+					},
+				},
+			}
+			gotScaleSets, err := client.getScaleSets(context.Background())
+			if tc.wantErr {
+				assert.Error(err)
+				return
+			}
+			require.NoError(err)
+			assert.ElementsMatch(tc.wantScaleSets, gotScaleSets)
+		})
+	}
+}

operators/constellation-node-operator/internal/azure/client/scalinggroup.go

@@ -3,6 +3,7 @@ package client
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v2"
 )
@@ -52,3 +53,32 @@ func (c *Client) SetScalingGroupImage(ctx context.Context, scalingGroupID, image
 	}
 	return nil
 }
+
+// GetScalingGroupName retrieves the name of a scaling group.
+func (c *Client) GetScalingGroupName(ctx context.Context, scalingGroupID string) (string, error) {
+	_, _, scaleSet, err := splitVMSSID(scalingGroupID)
+	if err != nil {
+		return "", fmt.Errorf("getting scaling group name: %w", err)
+	}
+	return strings.ToLower(scaleSet), nil
+}
+
+// ListScalingGroups retrieves a list of scaling groups for the cluster.
+func (c *Client) ListScalingGroups(ctx context.Context, uid string) (controlPlaneGroupIDs []string, workerGroupIDs []string, err error) {
+	scaleSetIDs, err := c.getScaleSets(ctx)
+	if err != nil {
+		return nil, nil, fmt.Errorf("listing scaling groups: %w", err)
+	}
+	for _, scaleSetID := range scaleSetIDs {
+		_, _, scaleSet, err := splitVMSSID(scaleSetID)
+		if err != nil {
+			return nil, nil, fmt.Errorf("getting scaling group name: %w", err)
+		}
+		if scaleSet == "constellation-scale-set-controlplanes-"+uid {
+			controlPlaneGroupIDs = append(controlPlaneGroupIDs, scaleSetID)
+		} else if strings.HasPrefix(scaleSet, "constellation-scale-set-workers-"+uid) {
+			workerGroupIDs = append(workerGroupIDs, scaleSetID)
+		}
+	}
+	return controlPlaneGroupIDs, workerGroupIDs, nil
+}

operators/constellation-node-operator/internal/azure/client/scalinggroup_test.go

@@ -123,3 +123,97 @@ func TestSetScalingGroupImage(t *testing.T) {
 		})
 	}
 }
+
+func TestGetScalingGroupName(t *testing.T) {
+	testCases := map[string]struct {
+		scalingGroupID string
+		wantName       string
+		wantErr        bool
+	}{
+		"getting name works": {
+			scalingGroupID: "/subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/scale-set-name",
+			wantName:       "scale-set-name",
+		},
+		"uppercase name is lowercased": {
+			scalingGroupID: "/subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/SCALE-SET-NAME",
+			wantName:       "scale-set-name",
+		},
+		"splitting scalingGroupID fails": {
+			scalingGroupID: "invalid",
+			wantErr:        true,
+		},
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			assert := assert.New(t)
+			require := require.New(t)
+
+			client := Client{}
+			gotName, err := client.GetScalingGroupName(context.Background(), tc.scalingGroupID)
+			if tc.wantErr {
+				assert.Error(err)
+				return
+			}
+			require.NoError(err)
+			assert.Equal(tc.wantName, gotName)
+		})
+	}
+}
+
+func TestListScalingGroups(t *testing.T) {
+	testCases := map[string]struct {
+		scaleSet          armcomputev2.VirtualMachineScaleSet
+		fetchPageErr      error
+		wantControlPlanes []string
+		wantWorkers       []string
+		wantErr           bool
+	}{
+		"listing control-plane works": {
+			scaleSet: armcomputev2.VirtualMachineScaleSet{
+				ID: to.Ptr("/subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/constellation-scale-set-controlplanes-uid"),
+			},
+			wantControlPlanes: []string{"/subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/constellation-scale-set-controlplanes-uid"},
+		},
+		"listing worker works": {
+			scaleSet: armcomputev2.VirtualMachineScaleSet{
+				ID: to.Ptr("/subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/constellation-scale-set-workers-uid"),
+			},
+			wantWorkers: []string{"/subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/constellation-scale-set-workers-uid"},
+		},
+		"listing other works": {
+			scaleSet: armcomputev2.VirtualMachineScaleSet{
+				ID: to.Ptr("/subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/other"),
+			},
+		},
+		"fetching scale sets fails": {
+			fetchPageErr: errors.New("fetch page error"),
+			wantErr:      true,
+		},
+		"scale set is invalid": {},
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			assert := assert.New(t)
+			require := require.New(t)
+
+			client := Client{
+				scaleSetsAPI: &stubScaleSetsAPI{
+					pager: &stubVMSSPager{
+						list:     []armcomputev2.VirtualMachineScaleSet{tc.scaleSet},
+						fetchErr: tc.fetchPageErr,
+					},
+				},
+			}
+			gotControlPlanes, gotWorkers, err := client.ListScalingGroups(context.Background(), "uid")
+			if tc.wantErr {
+				assert.Error(err)
+				return
+			}
+			require.NoError(err)
+			assert.ElementsMatch(tc.wantControlPlanes, gotControlPlanes)
+			assert.ElementsMatch(tc.wantWorkers, gotWorkers)
+		})
+	}
+}

operators/constellation-node-operator/internal/azure/client/vmss.go

@@ -9,7 +9,7 @@ import (
 var vmssIDRegexp = regexp.MustCompile(`^/subscriptions/([^/]+)/resourceGroups/([^/]+)/providers/Microsoft.Compute/virtualMachineScaleSets/([^/]+)$`)
 
 // joinVMSSID joins scale set parameters to generate a virtual machine scale set (VMSS) ID.
-// Format: /subscriptions/<subscription>/resourceGroups/<resource-group>/providers/Microsoft.Compute/virtualMachineScaleSets/<scale-set>
+// Format: /subscriptions/<subscription>/resourceGroups/<resource-group>/providers/Microsoft.Compute/virtualMachineScaleSets/<scale-set> .
 func joinVMSSID(subscriptionID, resourceGroup, scaleSet string) string {
 	return fmt.Sprintf("/subscriptions/%s/resourceGroups/%s/providers/Microsoft.Compute/virtualMachineScaleSets/%s", subscriptionID, resourceGroup, scaleSet)
 }