Get instance role from tags on Azure

internal/cloud/azure/imds.go

@@ -10,9 +10,12 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"io"
 	"net/http"
 	"time"
+
+	"github.com/edgelesssys/constellation/v2/internal/role"
 )
 
 // subset of azure imds API: https://docs.microsoft.com/en-us/azure/virtual-machines/windows/instance-metadata-service?tabs=linux
@@ -87,17 +90,29 @@ func (c *imdsClient) UID(ctx context.Context) (string, error) {
 		}
 	}
 
-	if len(c.cache.Compute.Tags) == 0 {
-		return "", errors.New("unable to get uid")
-	}
-
 	for _, tag := range c.cache.Compute.Tags {
-		if tag.Name == "uid" {
+		if tag.Name == "constellation-uid" {
 			return tag.Value, nil
 		}
 	}
 
-	return "", errors.New("unable to get uid from metadata tags")
+	return "", fmt.Errorf("unable to get uid from metadata tags %v", c.cache.Compute.Tags)
+}
+
+func (c *imdsClient) Role(ctx context.Context) (role.Role, error) {
+	if c.timeForUpdate() || len(c.cache.Compute.Tags) == 0 {
+		if err := c.update(ctx); err != nil {
+			return role.Unknown, err
+		}
+	}
+
+	for _, tag := range c.cache.Compute.Tags {
+		if tag.Name == "role" {
+			return role.FromString(tag.Value), nil
+		}
+	}
+
+	return role.Unknown, fmt.Errorf("unable to get role from metadata tags %v", c.cache.Compute.Tags)
 }
 
 // timeForUpdate checks whether an update is needed due to cache age.

internal/cloud/azure/imds_test.go

@@ -15,12 +15,16 @@ import (
 	"net/http/httptest"
 	"testing"
 
+	"github.com/edgelesssys/constellation/v2/internal/role"
 	"github.com/stretchr/testify/assert"
 	"google.golang.org/grpc/test/bufconn"
 )
 
 func TestIMDSClient(t *testing.T) {
-	uidTags := []metadataTag{{Name: "uid", Value: "uid"}}
+	uidTags := []metadataTag{
+		{Name: "constellation-uid", Value: "uid"},
+		{Name: "role", Value: "worker"},
+	}
 	response := metadataResponse{
 		Compute: metadataResponseCompute{
 			ResourceID:    "resource-id",
@@ -44,6 +48,14 @@ func TestIMDSClient(t *testing.T) {
 		Compute: metadataResponseCompute{
 			ResourceID:    "resource-id",
 			ResourceGroup: "resource-group",
+			Tags:          []metadataTag{{Name: "role", Value: "worker"}},
+		},
+	}
+	responseWithoutRole := metadataResponse{
+		Compute: metadataResponseCompute{
+			ResourceID:    "resource-id",
+			ResourceGroup: "resource-group",
+			Tags:          []metadataTag{{Name: "constellation-uid", Value: "uid"}},
 		},
 	}
 
@@ -55,30 +67,43 @@ func TestIMDSClient(t *testing.T) {
 		wantResourceGroup    string
 		wantUIDErr           bool
 		wantUID              string
+		wantRoleErr          bool
+		wantRole             role.Role
 	}{
 		"metadata response parsed": {
 			server:            newHTTPBufconnServerWithMetadataResponse(response),
 			wantProviderID:    "resource-id",
 			wantResourceGroup: "resource-group",
 			wantUID:           "uid",
+			wantRole:          role.Worker,
 		},
 		"metadata response without resource ID": {
 			server:            newHTTPBufconnServerWithMetadataResponse(responseWithoutID),
 			wantProviderIDErr: true,
 			wantResourceGroup: "resource-group",
 			wantUID:           "uid",
+			wantRole:          role.Worker,
 		},
 		"metadata response without UID tag": {
 			server:            newHTTPBufconnServerWithMetadataResponse(responseWithoutUID),
 			wantProviderID:    "resource-id",
 			wantResourceGroup: "resource-group",
 			wantUIDErr:        true,
+			wantRole:          role.Worker,
+		},
+		"metadata response without role tag": {
+			server:            newHTTPBufconnServerWithMetadataResponse(responseWithoutRole),
+			wantProviderID:    "resource-id",
+			wantResourceGroup: "resource-group",
+			wantUID:           "uid",
+			wantRoleErr:       true,
 		},
 		"metadata response without resource group": {
 			server:               newHTTPBufconnServerWithMetadataResponse(responseWithoutGroup),
 			wantProviderID:       "resource-id",
 			wantResourceGroupErr: true,
 			wantUID:              "uid",
+			wantRole:             role.Worker,
 		},
 		"invalid imds response detected": {
 			server: newHTTPBufconnServer(func(writer http.ResponseWriter, request *http.Request) {
@@ -87,6 +112,7 @@ func TestIMDSClient(t *testing.T) {
 			wantProviderIDErr:    true,
 			wantResourceGroupErr: true,
 			wantUIDErr:           true,
+			wantRoleErr:          true,
 		},
 	}
 
@@ -131,6 +157,14 @@ func TestIMDSClient(t *testing.T) {
 				assert.NoError(err)
 				assert.Equal(tc.wantUID, uid)
 			}
+
+			role, err := iClient.Role(ctx)
+			if tc.wantRoleErr {
+				assert.Error(err)
+			} else {
+				assert.NoError(err)
+				assert.Equal(tc.wantRole, role)
+			}
 		})
 	}
 }

internal/cloud/azure/metadata.go

@@ -300,7 +300,13 @@ func (m *Metadata) getAppInsights(ctx context.Context) (*armapplicationinsights.
 			if component == nil || component.Tags == nil {
 				continue
 			}
-			if *component.Tags["uid"] == uid {
+
+			tag, ok := component.Tags["constellation-uid"]
+			if !ok || tag == nil {
+				continue
+			}
+
+			if *tag == uid {
 				return component, nil
 			}
 		}

internal/cloud/azure/metadata_test.go

@@ -15,6 +15,7 @@ import (
 	armcomputev2 "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v2"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/network/armnetwork"
 	"github.com/edgelesssys/constellation/v2/internal/cloud/metadata"
+	"github.com/edgelesssys/constellation/v2/internal/role"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -24,6 +25,7 @@ func TestList(t *testing.T) {
 		{
 			Name:       "scale-set-name-instance-id",
 			ProviderID: "azure:///subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/scale-set-name/virtualMachines/instance-id",
+			Role:       role.Worker,
 			VPCIP:      "192.0.2.0",
 			SSHKeys:    map[string][]string{"user": {"key-data"}},
 		},
@@ -87,6 +89,7 @@ func TestSelf(t *testing.T) {
 	wantScaleSetInstance := metadata.InstanceMetadata{
 		Name:       "scale-set-name-instance-id",
 		ProviderID: "azure:///subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/scale-set-name/virtualMachines/instance-id",
+		Role:       role.Worker,
 		VPCIP:      "192.0.2.0",
 		SSHKeys:    map[string][]string{"user": {"key-data"}},
 	}
@@ -650,6 +653,10 @@ func newScaleSetsStub() *stubScaleSetsAPI {
 		pager: &stubVirtualMachineScaleSetsClientListPager{
 			list: []armcomputev2.VirtualMachineScaleSet{{
 				Name: to.Ptr("scale-set-name"),
+				Tags: map[string]*string{
+					"constellation-uid": to.Ptr("uid"),
+					"role":              to.Ptr("worker"),
+				},
 			}},
 		},
 	}
@@ -683,9 +690,14 @@ func newVirtualMachineScaleSetsVMsStub() *stubVirtualMachineScaleSetVMsAPI {
 					},
 				},
 			},
+			Tags: map[string]*string{
+				"constellation-uid": to.Ptr("uid"),
+				"role":              to.Ptr("worker"),
+			},
 		},
 		pager: &stubVirtualMachineScaleSetVMPager{
-			list: []armcomputev2.VirtualMachineScaleSetVM{{
+			list: []armcomputev2.VirtualMachineScaleSetVM{
+				{
 					Name:       to.Ptr("scale-set-name_instance-id"),
 					InstanceID: to.Ptr("instance-id"),
 					ID:         to.Ptr("/subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/scale-set-name/virtualMachines/instance-id"),
@@ -711,7 +723,12 @@ func newVirtualMachineScaleSetsVMsStub() *stubVirtualMachineScaleSetVMsAPI {
 							},
 						},
 					},
-			}},
+					Tags: map[string]*string{
+						"constellation-uid": to.Ptr("uid"),
+						"role":              to.Ptr("worker"),
+					},
+				},
+			},
 		},
 	}
 }

internal/cloud/azure/scaleset.go

@@ -10,7 +10,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"regexp"
 	"strings"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
@@ -21,11 +20,6 @@ import (
 	"github.com/edgelesssys/constellation/v2/internal/role"
 )
 
-var (
-	controlPlaneScaleSetRegexp = regexp.MustCompile(`constellation-scale-set-controlplanes-[0-9a-zA-Z]+$`)
-	workerScaleSetRegexp       = regexp.MustCompile(`constellation-scale-set-workers-[0-9a-zA-Z]+$`)
-)
-
 // getScaleSetVM tries to get an azure vm belonging to a scale set.
 func (m *Metadata) getScaleSetVM(ctx context.Context, providerID string) (metadata.InstanceMetadata, error) {
 	_, resourceGroup, scaleSet, instanceID, err := azureshared.ScaleSetInformationFromProviderID(providerID)
@@ -45,7 +39,7 @@ func (m *Metadata) getScaleSetVM(ctx context.Context, providerID string) (metada
 		return metadata.InstanceMetadata{}, err
 	}
 
-	return convertScaleSetVMToCoreInstance(scaleSet, vmResp.VirtualMachineScaleSetVM, networkInterfaces, publicIPAddress)
+	return convertScaleSetVMToCoreInstance(vmResp.VirtualMachineScaleSetVM, networkInterfaces, publicIPAddress)
 }
 
 // listScaleSetVMs lists all scale set VMs in the current resource group.
@@ -75,7 +69,7 @@ func (m *Metadata) listScaleSetVMs(ctx context.Context, resourceGroup string) ([
 					if err != nil {
 						return nil, err
 					}
-					instance, err := convertScaleSetVMToCoreInstance(*scaleSet.Name, *vm, interfaces, "")
+					instance, err := convertScaleSetVMToCoreInstance(*vm, interfaces, "")
 					if err != nil {
 						return nil, err
 					}
@@ -88,7 +82,9 @@ func (m *Metadata) listScaleSetVMs(ctx context.Context, resourceGroup string) ([
 }
 
 // convertScaleSetVMToCoreInstance converts an azure scale set virtual machine with interface configurations into a core.Instance.
-func convertScaleSetVMToCoreInstance(scaleSet string, vm armcomputev2.VirtualMachineScaleSetVM, networkInterfaces []armnetwork.Interface, publicIPAddress string) (metadata.InstanceMetadata, error) {
+func convertScaleSetVMToCoreInstance(vm armcomputev2.VirtualMachineScaleSetVM, networkInterfaces []armnetwork.Interface,
+	publicIPAddress string,
+) (metadata.InstanceMetadata, error) {
 	if vm.ID == nil {
 		return metadata.InstanceMetadata{}, errors.New("retrieving instance from armcompute API client returned no instance ID")
 	}
@@ -101,10 +97,15 @@ func convertScaleSetVMToCoreInstance(scaleSet string, vm armcomputev2.VirtualMac
 	} else {
 		sshKeys = extractSSHKeys(*vm.Properties.OSProfile.LinuxConfiguration.SSH)
 	}
+
+	if vm.Tags == nil {
+		return metadata.InstanceMetadata{}, errors.New("retrieving instance from armcompute API client returned no tags")
+	}
+
 	return metadata.InstanceMetadata{
 		Name:       *vm.Properties.OSProfile.ComputerName,
 		ProviderID: "azure://" + *vm.ID,
-		Role:       extractScaleSetVMRole(scaleSet),
+		Role:       extractScaleSetVMRole(vm.Tags),
 		VPCIP:      extractVPCIP(networkInterfaces),
 		PublicIP:   publicIPAddress,
 		SSHKeys:    sshKeys,
@@ -112,14 +113,18 @@ func convertScaleSetVMToCoreInstance(scaleSet string, vm armcomputev2.VirtualMac
 }
 
 // extractScaleSetVMRole extracts the constellation role of a scale set using its name.
-func extractScaleSetVMRole(scaleSet string) role.Role {
+func extractScaleSetVMRole(tags map[string]*string) role.Role {
-	if controlPlaneScaleSetRegexp.MatchString(scaleSet) {
+	if tags == nil {
-		return role.ControlPlane
-	}
-	if workerScaleSetRegexp.MatchString(scaleSet) {
-		return role.Worker
-	}
 		return role.Unknown
+	}
+	roleStr, ok := tags["role"]
+	if !ok {
+		return role.Unknown
+	}
+	if roleStr == nil {
+		return role.Unknown
+	}
+	return role.FromString(*roleStr)
 }
 
 // ImageReferenceFromImage sets the `ID` or `CommunityGalleryImageID` field

internal/cloud/azure/scaleset_test.go

@@ -24,6 +24,7 @@ func TestGetScaleSetVM(t *testing.T) {
 	wantInstance := metadata.InstanceMetadata{
 		Name:       "scale-set-name-instance-id",
 		ProviderID: "azure:///subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/scale-set-name/virtualMachines/instance-id",
+		Role:       role.Worker,
 		VPCIP:      "192.0.2.0",
 		SSHKeys:    map[string][]string{"user": {"key-data"}},
 	}
@@ -83,6 +84,7 @@ func TestListScaleSetVMs(t *testing.T) {
 		{
 			Name:       "scale-set-name-instance-id",
 			ProviderID: "azure:///subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/scale-set-name/virtualMachines/instance-id",
+			Role:       role.Worker,
 			VPCIP:      "192.0.2.0",
 			SSHKeys:    map[string][]string{"user": {"key-data"}},
 		},
@@ -203,7 +205,7 @@ func TestConvertScaleSetVMToCoreInstance(t *testing.T) {
 			assert := assert.New(t)
 			require := require.New(t)
 
-			instance, err := convertScaleSetVMToCoreInstance("scale-set", tc.inVM, tc.inInterface, tc.inPublicIP)
+			instance, err := convertScaleSetVMToCoreInstance(tc.inVM, tc.inInterface, tc.inPublicIP)
 
 			if tc.wantErr {
 				assert.Error(err)
@@ -217,19 +219,31 @@ func TestConvertScaleSetVMToCoreInstance(t *testing.T) {
 
 func TestExtractScaleSetVMRole(t *testing.T) {
 	testCases := map[string]struct {
-		scaleSet string
+		tags     map[string]*string
 		wantRole role.Role
 	}{
-		"bootstrapper role": {
+		"control-plane role": {
-			scaleSet: "constellation-scale-set-controlplanes-abcd123",
+			tags:     map[string]*string{"role": to.Ptr("control-plane")},
 			wantRole: role.ControlPlane,
 		},
-		"node role": {
+		"worker role": {
-			scaleSet: "constellation-scale-set-workers-abcd123",
+			tags:     map[string]*string{"role": to.Ptr("worker")},
 			wantRole: role.Worker,
 		},
 		"unknown role": {
-			scaleSet: "unknown",
+			tags:     map[string]*string{"role": to.Ptr("foo")},
+			wantRole: role.Unknown,
+		},
+		"no role": {
+			tags:     map[string]*string{},
+			wantRole: role.Unknown,
+		},
+		"nil role": {
+			tags:     map[string]*string{"role": nil},
+			wantRole: role.Unknown,
+		},
+		"nil tags": {
+			tags:     nil,
 			wantRole: role.Unknown,
 		},
 	}
@@ -238,7 +252,7 @@ func TestExtractScaleSetVMRole(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			assert := assert.New(t)
 
-			role := extractScaleSetVMRole(tc.scaleSet)
+			role := extractScaleSetVMRole(tc.tags)
 
 			assert.Equal(tc.wantRole, role)
 		})
@@ -266,7 +280,7 @@ func newListContainingNilScaleSetVirtualMachinesStub() *stubVirtualMachineScaleS
 					ID:         to.Ptr("/subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/scale-set-name/virtualMachines/instance-id"),
 					InstanceID: to.Ptr("instance-id"),
 					Tags: map[string]*string{
-						"tag-key": to.Ptr("tag-value"),
+						"role": to.Ptr("worker"),
 					},
 					Properties: &armcomputev2.VirtualMachineScaleSetVMProperties{
 						NetworkProfile: &armcomputev2.NetworkProfile{