Add support for Virtual Machine Scale Sets in the Boundary Azure Plugin (#22)

* first pass at pulling in VMSS data

* infer vmss set id and external name from vmss data as none of these items are present by default

* updated scale set name formatting to align with azure portal and json outputs

* Update plugin/plugin.go

Updated following code review

Co-authored-by: Johan Brandhorst-Satzkorn <johan.brandhorst@gmail.com>

* Update plugin/plugin.go

Updated following code review

Co-authored-by: Johan Brandhorst-Satzkorn <johan.brandhorst@gmail.com>

* complete refactoring of initial implementation following code review, taking into accounts comments and common sense updates

* updated README.md to reflect changes to support VMSS

* fixed typo in README.md

* WIP: initial updates following code review

* refactored resource functions

* removed instrumentation statements

* resolve ifname comments

* address formatting issues in accordance with style recommendations

---------

Co-authored-by: Johan Brandhorst-Satzkorn <johan.brandhorst@gmail.com>

plugin/client.go

@@ -65,6 +65,56 @@ func getVirtualMachinesClient(opt ...Option) (*compute.VirtualMachinesClient, er
 	return &client, nil
 }
 
+func getVirtualMachineScaleSetsClient(opt ...Option) (*compute.VirtualMachineScaleSetsClient, error) {
+	opts, err := getOpts(opt...)
+	if err != nil {
+		return nil, fmt.Errorf("error parsing ops: %w", err)
+	}
+	if opts.withSubscriptionId == "" {
+		return nil, errors.New("subscription id is empty")
+	}
+	if opts.withAuthorizer == nil {
+		return nil, errors.New("authorizer is nil")
+	}
+	var client compute.VirtualMachineScaleSetsClient
+	switch opts.withBaseUrl {
+	case "":
+		client = compute.NewVirtualMachineScaleSetsClient(opts.withSubscriptionId)
+	default:
+		client = compute.NewVirtualMachineScaleSetsClientWithBaseURI(opts.withBaseUrl, opts.withSubscriptionId)
+	}
+	client.Authorizer = opts.withAuthorizer
+	if err := client.AddToUserAgent(opts.withUserAgent); err != nil {
+		return nil, fmt.Errorf("error adding identifier to user-agent: %w", err)
+	}
+	return &client, nil
+}
+
+func getVirtualMachineScaleSetVMsClient(opt ...Option) (*compute.VirtualMachineScaleSetVMsClient, error) {
+	opts, err := getOpts(opt...)
+	if err != nil {
+		return nil, fmt.Errorf("error parsing ops: %w", err)
+	}
+	if opts.withSubscriptionId == "" {
+		return nil, errors.New("subscription id is empty")
+	}
+	if opts.withAuthorizer == nil {
+		return nil, errors.New("authorizer is nil")
+	}
+	var client compute.VirtualMachineScaleSetVMsClient
+	switch opts.withBaseUrl {
+	case "":
+		client = compute.NewVirtualMachineScaleSetVMsClient(opts.withSubscriptionId)
+	default:
+		client = compute.NewVirtualMachineScaleSetVMsClientWithBaseURI(opts.withBaseUrl, opts.withSubscriptionId)
+	}
+	client.Authorizer = opts.withAuthorizer
+	if err := client.AddToUserAgent(opts.withUserAgent); err != nil {
+		return nil, fmt.Errorf("error adding identifier to user-agent: %w", err)
+	}
+	return &client, nil
+}
+
 func getNetworkInterfacesClient(opt ...Option) (*network.InterfacesClient, error) {
 	opts, err := getOpts(opt...)
 	if err != nil {

plugin/const.go

@@ -4,22 +4,30 @@
 package plugin
 
 const (
-	constBaseUrl                   = "base_url"
+	// Azure service constants
+	constMsComputeService = "Microsoft.Compute"
+	constMsNetworkService = "Microsoft.Network"
+	constSubscriptions    = "subscriptions"
+	constResourceGroups   = "resourceGroups"
+	constProviders        = "providers"
+
+	// Resource type constants
+	constVirtualMachinesResource         = "virtualMachines"
+	constVirtualMachineScaleSetsResource = "virtualMachineScaleSets"
+	constNetworkInterfacesResource       = "networkInterfaces"
+	constPublicIpAddressesResource       = "publicIPAddresses"
+
+	// Field constants
 	constFilter                    = "filter"
-	constSubscriptionId            = "subscription_id"
-	constClientId                  = "client_id"
+	constSecretId                  = "secret_id"
 	constSecretValue               = "secret_value"
 	constCredsLastRotatedTime      = "creds_last_rotated_time"
-	constSecretId                  = "secret_id"
-	constTenantId                  = "tenant_id"
-	constDefaultFilter             = "resourceType eq 'Microsoft.Compute/virtualMachines'"
-	constMsComputeService          = "Microsoft.Compute"
-	constVirtualMachinesResource   = "virtualMachines"
-	constMsNetworkService          = "Microsoft.Network"
-	constNetworkInterfacesResource = "networkInterfaces"
-	constPublicIpAddressesResource = "publicIPAddresses"
-	constSubscriptions             = "subscriptions"
-	constResourceGroups            = "resourceGroups"
-	constProviders                 = "providers"
 	constDisableCredentialRotation = "disable_credential_rotation"
+	constSubscriptionId            = "subscription_id"
+	constTenantId                  = "tenant_id"
+	constClientId                  = "client_id"
+	constBaseUrl                   = "base_url"
+
+	// Worker pool configuration
+	constDefaultWorkers = 1000
 )

plugin/network_processor.go

@@ -0,0 +1,277 @@
+package plugin
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/Azure/azure-sdk-for-go/profiles/latest/compute/mgmt/compute"
+	"github.com/Azure/azure-sdk-for-go/profiles/latest/network/mgmt/network"
+)
+
+type networkInfo struct {
+	IpAddresses []string
+}
+
+// Common interface for IP configuration processing
+type networkProcessor interface {
+	processIPConfigurations(ctx context.Context, ipConfigs []network.InterfaceIPConfiguration, ifName string) (*networkInfo, error)
+	getPublicIPAddress(ctx context.Context, publicIPRef *network.PublicIPAddress, ipConf *network.InterfaceIPConfiguration, ifName string) (string, error)
+	getNetworkInterface(ctx context.Context, interfaceID string) (*network.Interface, string, error)
+}
+
+// Standard VM network processor
+type vmNetworkProcessor struct {
+	clients       *azureClients
+	resourceGroup string
+}
+
+// VMSS network processor
+type vmssNetworkProcessor struct {
+	clients       *azureClients
+	resourceGroup string
+	vmssName      string
+	instanceID    string
+}
+
+func processNetworkProfile(
+	ctx context.Context,
+	interfaces []compute.NetworkInterfaceReference,
+	processor networkProcessor,
+	targetInfo *networkInfo,
+) error {
+
+	for _, ifaceRef := range interfaces {
+		if ifaceRef.ID == nil {
+			return fmt.Errorf("nil ID for network interface")
+		}
+
+		iface, ifName, err := processor.getNetworkInterface(ctx, *ifaceRef.ID)
+		if err != nil {
+			return err
+		}
+
+		if iface.InterfacePropertiesFormat == nil || iface.InterfacePropertiesFormat.IPConfigurations == nil {
+			continue
+		}
+
+		networkInfo, err := processor.processIPConfigurations(ctx, *iface.InterfacePropertiesFormat.IPConfigurations, ifName)
+		if err != nil {
+			return err
+		}
+		targetInfo.IpAddresses = append(targetInfo.IpAddresses, networkInfo.IpAddresses...)
+	}
+	return nil
+}
+
+func (p *vmNetworkProcessor) getNetworkInterface(ctx context.Context, interfaceID string) (*network.Interface, string, error) {
+	ifResGroup, ifName, err := splitId(interfaceID, constMsNetworkService, constNetworkInterfacesResource)
+	if err != nil {
+		return nil, "", fmt.Errorf("error splitting network interface id %q: %w", interfaceID, err)
+	}
+
+	iface, err := p.clients.ifClient.Get(ctx, ifResGroup, ifName, "")
+	if err != nil {
+		return nil, "", fmt.Errorf("error fetching network interface: %w", err)
+	}
+
+	return &iface, ifName, nil
+}
+
+func (p *vmNetworkProcessor) processIPConfigurations(
+	ctx context.Context,
+	ipConfigs []network.InterfaceIPConfiguration,
+	ifName string,
+) (
+	*networkInfo,
+	error,
+) {
+	tempInfo := &networkInfo{}
+	for _, ipconf := range ipConfigs {
+		// Process private IP
+		if ipconf.PrivateIPAddress != nil {
+			tempInfo.IpAddresses = append(tempInfo.IpAddresses, *ipconf.PrivateIPAddress)
+		}
+
+		// Process public IP if available
+		if ipconf.PublicIPAddress != nil {
+			ipAddress, err := p.getPublicIPAddress(ctx, ipconf.PublicIPAddress, &ipconf, ifName)
+			if err != nil {
+				return nil, err
+			}
+			tempInfo.IpAddresses = append(tempInfo.IpAddresses, ipAddress)
+		}
+	}
+	return tempInfo, nil
+}
+
+func (p *vmNetworkProcessor) getPublicIPAddress(
+	ctx context.Context,
+	publicIP *network.PublicIPAddress,
+	ipConf *network.InterfaceIPConfiguration,
+	ifName string,
+) (
+	string,
+	error,
+) {
+	if publicIP.ID == nil {
+		return "", fmt.Errorf("ip configuration has public IP address info but nil id")
+	}
+
+	ipResGroup, ipName, err := splitId(*publicIP.ID, constMsNetworkService, constPublicIpAddressesResource)
+	if err != nil {
+		return "", fmt.Errorf("error splitting public ip address id: %w", err)
+	}
+
+	pipInfo, err := p.clients.pipClient.Get(ctx, ipResGroup, ipName, "")
+	if err != nil {
+		return "", fmt.Errorf("error fetching public IP information: %w", err)
+	}
+
+	if pipInfo.PublicIPAddressPropertiesFormat == nil {
+		return "", fmt.Errorf("nil public ip address properties format for public ip %q", *publicIP.ID)
+	}
+
+	if pipInfo.PublicIPAddressPropertiesFormat.IPAddress == nil {
+		return "", fmt.Errorf("nil IP address for public ip %q", *publicIP.ID)
+	}
+
+	return *pipInfo.PublicIPAddressPropertiesFormat.IPAddress, nil
+}
+
+func (p *vmssNetworkProcessor) getNetworkInterface(ctx context.Context, interfaceID string) (*network.Interface, string, error) {
+	ifName := extractResourceSuffix(interfaceID)
+
+	iface, err := p.clients.ifClient.GetVirtualMachineScaleSetNetworkInterface(
+		ctx,
+		p.resourceGroup,
+		p.vmssName,
+		p.instanceID,
+		ifName,
+		"")
+	if err != nil {
+		return nil, "", fmt.Errorf("error fetching VMSS network interface: %w", err)
+	}
+
+	return &iface, ifName, nil
+}
+
+func (p *vmssNetworkProcessor) processIPConfigurations(
+	ctx context.Context,
+	ipConfigs []network.InterfaceIPConfiguration,
+	ifName string,
+) (
+	*networkInfo,
+	error,
+) {
+	tempInfo := &networkInfo{}
+	for _, ipconf := range ipConfigs {
+		if ipconf.PrivateIPAddress != nil {
+			tempInfo.IpAddresses = append(tempInfo.IpAddresses, *ipconf.PrivateIPAddress)
+		}
+
+		if ipconf.PublicIPAddress != nil {
+			ipAddress, err := p.getPublicIPAddress(ctx, ipconf.PublicIPAddress, &ipconf, ifName)
+			if err != nil {
+				return nil, err
+			}
+			tempInfo.IpAddresses = append(tempInfo.IpAddresses, ipAddress)
+		}
+	}
+	return tempInfo, nil
+}
+
+func (p *vmssNetworkProcessor) getPublicIPAddress(
+	ctx context.Context,
+	publicIP *network.PublicIPAddress,
+	ipConf *network.InterfaceIPConfiguration,
+	ifName string,
+) (
+	string,
+	error,
+) {
+	if publicIP == nil || publicIP.ID == nil || ipConf.Name == nil {
+		return "", fmt.Errorf("invalid public IP configuration")
+	}
+
+	pipName := extractResourceSuffix(*publicIP.ID)
+	pipInfo, err := p.clients.pipClient.GetVirtualMachineScaleSetPublicIPAddress(
+		ctx,
+		p.resourceGroup,
+		p.vmssName,
+		p.instanceID,
+		ifName,
+		*ipConf.Name,
+		pipName,
+		"")
+	if err != nil {
+		return "", fmt.Errorf("error fetching public IP information: %w", err)
+	}
+
+	if pipInfo.PublicIPAddressPropertiesFormat == nil {
+		return "", fmt.Errorf("nil public ip address properties format")
+	}
+
+	if pipInfo.PublicIPAddressPropertiesFormat.IPAddress == nil {
+		return "", fmt.Errorf("nil IP address for public ip %q", *publicIP.ID)
+	}
+
+	return *pipInfo.PublicIPAddressPropertiesFormat.IPAddress, nil
+}
+
+func processVMNetworkInterfaces(
+	ctx context.Context,
+	vm compute.VirtualMachine,
+	clients *azureClients,
+	resourceGroup string,
+	netInfo *networkInfo,
+) error {
+
+	if vm.VirtualMachineProperties == nil || vm.VirtualMachineProperties.NetworkProfile == nil {
+		return fmt.Errorf("error fetching network profile")
+	}
+
+	processor := &vmNetworkProcessor{
+		clients:       clients,
+		resourceGroup: resourceGroup,
+	}
+
+	return processNetworkProfile(ctx, *vm.VirtualMachineProperties.NetworkProfile.NetworkInterfaces, processor, netInfo)
+}
+
+func processVMSSNetworkInterfaces(
+	ctx context.Context,
+	vmssvm compute.VirtualMachineScaleSetVM,
+	resourceGroup, vmssName string,
+	clients *azureClients,
+	netInfo *networkInfo,
+) error {
+
+	if vmssvm.VirtualMachineScaleSetVMProperties == nil ||
+		vmssvm.VirtualMachineScaleSetVMProperties.NetworkProfile == nil {
+		return fmt.Errorf("error fetching network profile")
+	}
+
+	if vmssvm.InstanceID == nil {
+		return fmt.Errorf("instance ID is nil")
+	}
+
+	processor := &vmssNetworkProcessor{
+		clients:       clients,
+		resourceGroup: resourceGroup,
+		vmssName:      vmssName,
+		instanceID:    *vmssvm.InstanceID,
+	}
+
+	return processNetworkProfile(ctx,
+		*vmssvm.VirtualMachineScaleSetVMProperties.NetworkProfile.NetworkInterfaces, processor, netInfo)
+}
+
+// extractResourceSuffix extracts the resource name from the end of an Azure resource ID
+func extractResourceSuffix(resourceID string) string {
+	parts := strings.Split(resourceID, "/")
+	if len(parts) > 0 {
+		return parts[len(parts)-1]
+	}
+	return ""
+}

plugin/plugin.go

@@ -1,6 +1,8 @@
 // Copyright (c) HashiCorp, Inc.
 // SPDX-License-Identifier: MPL-2.0
 
+// Package plugin implements the Azure host plugin for Boundary,
+// providing functionality to manage and list Azure hosts.
 package plugin
 
 import (
@@ -8,17 +10,20 @@ import (
 	"errors"
 	"fmt"
 	"strings"
+	"sync"
 	"time"
 
+	"github.com/Azure/azure-sdk-for-go/profiles/latest/compute/mgmt/compute"
 	"github.com/Azure/azure-sdk-for-go/profiles/latest/network/mgmt/network"
 	"github.com/Azure/azure-sdk-for-go/profiles/latest/resources/mgmt/resources"
-	"github.com/hashicorp/boundary/sdk/pbs/controller/api/resources/hostcatalogs"
-	"github.com/hashicorp/boundary/sdk/pbs/controller/api/resources/hostsets"
 	"github.com/mitchellh/mapstructure"
+	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/types/known/structpb"
 
+	"github.com/hashicorp/boundary/sdk/pbs/controller/api/resources/hostcatalogs"
+	"github.com/hashicorp/boundary/sdk/pbs/controller/api/resources/hostsets"
 	pb "github.com/hashicorp/boundary/sdk/pbs/plugin"
 )
 
@@ -28,37 +33,19 @@ type AzurePlugin struct {
 }
 
 // Ensure that we are implementing HostPluginServiceServer
-var (
-	_ pb.HostPluginServiceServer = (*AzurePlugin)(nil)
-)
+var _ pb.HostPluginServiceServer = (*AzurePlugin)(nil)
 
 type SetAttributes struct {
 	Filter string `mapstructure:"filter"`
 }
 
-func rotateCredFromCallback(ctx context.Context, catalog *hostcatalogs.HostCatalog) (*pb.HostCatalogPersisted, error) {
-	authzInfo, err := getAuthorizationInfo(catalog)
-	if err != nil {
-		return nil, fmt.Errorf("error getting hauth config when creating catalog: %w", err)
-	}
-	newPass, err := rotateCredential(ctx, authzInfo)
-	if err != nil {
-		return nil, fmt.Errorf("error rotating credentials when creating catalog: %w", err)
-	}
-	if newPass == nil {
-		return nil, errors.New("new credential back from rotate is nil")
-	}
-	newSecrets, err := structpb.NewStruct(map[string]interface{}{
-		constSecretId:             *newPass.KeyId,
-		constSecretValue:          *newPass.SecretText,
-		constCredsLastRotatedTime: time.Now().Format(time.RFC3339Nano),
-	})
-	if err != nil {
-		return nil, fmt.Errorf("error formatting new credential data as struct")
-	}
-	return &pb.HostCatalogPersisted{
-		Secrets: newSecrets,
-	}, err
+type azureClients struct {
+	resClient    *resources.Client
+	vmClient     *compute.VirtualMachinesClient
+	vmssClient   *compute.VirtualMachineScaleSetsClient
+	vmssvmClient *compute.VirtualMachineScaleSetVMsClient
+	ifClient     *network.InterfacesClient
+	pipClient    *network.PublicIPAddressesClient
 }
 
 func (p *AzurePlugin) OnCreateCatalog(_ context.Context, req *pb.OnCreateCatalogRequest) (*pb.OnCreateCatalogResponse, error) {
@@ -94,13 +81,10 @@ func (p *AzurePlugin) OnUpdateCatalog(_ context.Context, req *pb.OnUpdateCatalog
 	}
 	currentCatalog := req.GetCurrentCatalog()
 	if currentCatalog == nil {
-		// Should never happen
 		return nil, status.Error(codes.FailedPrecondition, "current catalog is nil")
 	}
 	secrets := req.GetNewCatalog().GetSecrets()
 	if secrets == nil {
-		// If new secrets weren't passed in, don't rotate what we have on
-		// update.
 		return &pb.OnUpdateCatalogResponse{}, nil
 	}
 	if err := validateSecrets(secrets); err != nil {
@@ -142,10 +126,161 @@ func (p *AzurePlugin) OnUpdateSet(_ context.Context, req *pb.OnUpdateSetRequest)
 
 func (p *AzurePlugin) ListHosts(ctx context.Context, req *pb.ListHostsRequest) (*pb.ListHostsResponse, error) {
 	if len(req.GetSets()) == 0 {
-		// Nothing to fetch
 		return &pb.ListHostsResponse{}, nil
 	}
-	catalog := req.GetCatalog()
+
+	// Initialize Azure clients and authorization
+	clients, err := p.initializeAzureResources(req.GetCatalog(), req.GetPersisted().GetSecrets())
+	if err != nil {
+		return nil, err
+	}
+
+	// Find matching resources and map them to sets
+	vmResources, vmssResources, resourceToSetMap, err := p.findMatchingResources(ctx, req.GetSets(), clients.resClient)
+	if err != nil {
+		return nil, err
+	}
+
+	// Create mutex and errgroup for parallel processing
+	mu := &sync.Mutex{}
+	vmToNetworkMap := make(map[string]networkInfo)
+	g, ctx := errgroup.WithContext(ctx)
+
+	// Process standard VMs in parallel if they exist
+	if len(vmResources) > 0 {
+		g.Go(func() error {
+			vmNetworkMap, err := p.processStandardVMs(ctx, vmResources, clients)
+			if err != nil {
+				return fmt.Errorf("processing VMs: %w", err)
+			}
+
+			mu.Lock()
+			defer mu.Unlock()
+			for k, v := range vmNetworkMap {
+				vmToNetworkMap[k] = v
+			}
+			return nil
+		})
+	}
+
+	// Process VMSS in parallel if they exist
+	if len(vmssResources) > 0 {
+		g.Go(func() error {
+			vmssNetworkMap, err := p.processVMScaleSets(ctx, vmssResources, clients)
+			if err != nil {
+				return fmt.Errorf("processing VMSS: %w", err)
+			}
+
+			mu.Lock()
+			defer mu.Unlock()
+			for k, v := range vmssNetworkMap {
+				vmToNetworkMap[k] = v
+			}
+			return nil
+		})
+	}
+
+	// Wait for all goroutines to complete and check for errors
+	if err := g.Wait(); err != nil {
+		return nil, err
+	}
+
+	response := buildHostsResponse(vmToNetworkMap, resourceToSetMap)
+	return response, nil
+}
+
+func rotateCredFromCallback(ctx context.Context, catalog *hostcatalogs.HostCatalog) (*pb.HostCatalogPersisted, error) {
+	authzInfo, err := getAuthorizationInfo(catalog)
+	if err != nil {
+		return nil, fmt.Errorf("error getting auth config when creating catalog: %w", err)
+	}
+	newPass, err := rotateCredential(ctx, authzInfo)
+	if err != nil {
+		return nil, fmt.Errorf("error rotating credentials when creating catalog: %w", err)
+	}
+	if newPass == nil {
+		return nil, errors.New("new credential back from rotate is nil")
+	}
+	newSecrets, err := structpb.NewStruct(map[string]interface{}{
+		constSecretId:             *newPass.KeyId,
+		constSecretValue:          *newPass.SecretText,
+		constCredsLastRotatedTime: time.Now().Format(time.RFC3339Nano),
+	})
+	if err != nil {
+		return nil, fmt.Errorf("error formatting new credential data as struct")
+	}
+	return &pb.HostCatalogPersisted{
+		Secrets: newSecrets,
+	}, err
+}
+
+func buildHostsResponse(
+	vmToNetworkMap map[string]networkInfo,
+	resourceToSetMap map[string][]string) *pb.ListHostsResponse {
+
+	ret := &pb.ListHostsResponse{}
+
+	for resourceID, networkInfo := range vmToNetworkMap {
+		host, err := createHostFromResource(resourceID, networkInfo, resourceToSetMap)
+		if err != nil {
+			fmt.Printf("Error creating host from resource: %v\n", err)
+			continue
+		}
+		ret.Hosts = append(ret.Hosts, host)
+	}
+
+	return ret
+}
+
+func createHostFromResource(resourceID string, networkInfo networkInfo,
+	resourceToSetMap map[string][]string) (*pb.ListHostsResponseHost, error) {
+
+	// Extract the resource type from the ID
+	resourceType, err := extractResourceType(resourceID)
+	if err != nil {
+		return nil, fmt.Errorf("invalid resource ID: %w", err)
+	}
+
+	switch resourceType {
+	case constVirtualMachineScaleSetsResource:
+		externalName, err := getExternalNameforVMSSInstance(resourceID)
+		if err != nil {
+			externalName = "" // Use empty name if there's an error
+		}
+
+		setId, err := getSetForVMSSInstance(resourceID)
+		if err != nil {
+			return nil, fmt.Errorf("could not determine set for resource ID: %w", err)
+		}
+		return &pb.ListHostsResponseHost{
+			ExternalId:   resourceID,
+			ExternalName: externalName,
+			IpAddresses:  networkInfo.IpAddresses,
+			SetIds:       resourceToSetMap[setId],
+		}, nil
+
+	case constVirtualMachinesResource:
+		_, externalName, err := splitId(resourceID, constMsComputeService, constVirtualMachinesResource)
+		if err != nil {
+			externalName = "" // Use empty name if there's an error
+		}
+
+		return &pb.ListHostsResponseHost{
+			ExternalId:   resourceID,
+			ExternalName: externalName,
+			IpAddresses:  networkInfo.IpAddresses,
+			SetIds:       resourceToSetMap[resourceID],
+		}, nil
+
+	default:
+		return nil, fmt.Errorf("unsupported resource type: %s", resourceType) // Return explicit error
+	}
+}
+
+func (p *AzurePlugin) initializeAzureResources(
+	catalog *hostcatalogs.HostCatalog,
+	secrets *structpb.Struct) (*azureClients, error) {
+
 	if catalog == nil {
 		return nil, status.Error(codes.FailedPrecondition, "catalog is nil")
 	}
@@ -154,8 +289,8 @@ func (p *AzurePlugin) ListHosts(ctx context.Context, req *pb.ListHostsRequest) (
 		return nil, status.Error(codes.FailedPrecondition, "catalog has no attributes")
 	}
 
-	// Create an authorizer
-	catalog.Secrets = req.GetPersisted().GetSecrets()
+	// Create authorizer
+	catalog.Secrets = secrets
 	authzInfo, err := getAuthorizationInfo(catalog)
 	if err != nil {
 		return nil, fmt.Errorf("error fetching authorizationInfo: %w", err)
@@ -163,6 +298,7 @@ func (p *AzurePlugin) ListHosts(ctx context.Context, req *pb.ListHostsRequest) (
 	if authzInfo == nil {
 		return nil, status.Error(codes.FailedPrecondition, "authorization info is nil")
 	}
+
 	authorizer, err := authzInfo.autorestAuthorizer()
 	if err != nil {
 		return nil, fmt.Errorf("error fetching autorest authorizer: %w", err)
@@ -171,16 +307,13 @@ func (p *AzurePlugin) ListHosts(ctx context.Context, req *pb.ListHostsRequest) (
 		return nil, errors.New("fetched authorizer is nil")
 	}
 
-	commonOpts := make([]Option, 0, 3)
-	commonOpts = append(commonOpts,
+	// Set up common client options
+	commonOpts := []Option{
 		WithSubscriptionId(authzInfo.AuthParams.SubscriptionId),
 		WithAuthorizer(authorizer),
-	)
+	}
 
-	// Getting authorizer will fail if the attr fields are not populated, so
-	// it's safe to access directly at this point.
-	//
-	// TODO: plumb base URL everywhere, not just here
+	// Add base URL if specified
 	catalogAttrFields := catalogAttrs.GetFields()
 	if catalogAttrFields[constBaseUrl] != nil {
 		if baseUrl := catalogAttrFields[constBaseUrl].GetStringValue(); baseUrl != "" {
@@ -188,229 +321,137 @@ func (p *AzurePlugin) ListHosts(ctx context.Context, req *pb.ListHostsRequest) (
 		}
 	}
 
-	// Create clients
-	resClient, err := getResourcesClient(commonOpts...)
-	if err != nil {
-		return nil, fmt.Errorf("error fetching resources client: %w", err)
-	}
-	if resClient == nil {
-		return nil, errors.New("resource client is nil")
-	}
-	vmClient, err := getVirtualMachinesClient(commonOpts...)
-	if err != nil {
-		return nil, fmt.Errorf("error fetching virtual machines client: %w", err)
-	}
-	if vmClient == nil {
-		return nil, errors.New("virtual machines client is nil")
-	}
-	ifClient, err := getNetworkInterfacesClient(commonOpts...)
-	if err != nil {
-		return nil, fmt.Errorf("error fetching network interfaces client: %w", err)
-	}
-	if ifClient == nil {
-		return nil, errors.New("network interfaces client is nil")
-	}
-	pipClient, err := getPublicIpAddressesClient(commonOpts...)
-	if err != nil {
-		return nil, fmt.Errorf("error fetching public ip address client: %w", err)
-	}
-	if pipClient == nil {
-		return nil, errors.New("public ip address client is nil")
+	// Initialize clients
+	clients := &azureClients{}
+	var initErr error
+
+	// Initialize individual clients with error handling
+	if clients.resClient, initErr = getResourcesClient(commonOpts...); initErr != nil {
+		return nil, fmt.Errorf("error fetching resources client: %w", initErr)
 	}
 
-	// This next section finds resource IDs that match to the filter in each
-	// host set, and stores a mapping of, for each resource ID, which host sets
-	// it belongs to
-	resourceToSetMap := make(map[string][]string, len(req.GetSets())*10)
-	var resourceInfos []resources.GenericResourceExpanded
-	for _, set := range req.GetSets() {
-		if set == nil {
-			return nil, errors.New("set is nil")
-		}
-		setAttrs := set.GetAttributes()
-		if setAttrs == nil {
-			return nil, fmt.Errorf("set %s has no attributes", set.GetId())
-		}
-		setAttrFields := setAttrs.GetFields()
-		if len(setAttrFields) == 0 {
-			return nil, fmt.Errorf("set %s attributes has no filter", set.GetId())
-		}
-
-		filter := constDefaultFilter
-		if setAttrFields[constFilter] != nil {
-			if filter = setAttrFields[constFilter].GetStringValue(); filter == "" {
-				return nil, fmt.Errorf("set %s filter is empty", set.GetId())
-			}
-		}
-
-		// List values matching the filter; ask for provisioning state information
-		// to ensure we are seeing only fully provisioned machines
-		{
-			iter, err := resClient.ListComplete(ctx, filter, "", nil)
-			if err != nil {
-				return nil, fmt.Errorf("error listing resources: %w", err)
-			}
-			for iter.NotDone() {
-				val := iter.Value()
-				if err := iter.NextWithContext(ctx); err != nil {
-					return nil, fmt.Errorf("error iterating resources: %w", err)
-				}
-				if val.ID == nil { // something went wrong, likely iterator has advanced beyond the end
-					continue
-				}
-				if val.Type == nil || !strings.EqualFold(*val.Type, "Microsoft.Compute/virtualMachines") {
-					// no point continuing if we can't validate that it's a VM
-					continue
-				}
-				resourceInfos = append(resourceInfos, val)
-				resourceToSetMap[*val.ID] = append(resourceToSetMap[*val.ID], set.GetId())
-			}
-		}
+	if clients.vmClient, initErr = getVirtualMachinesClient(commonOpts...); initErr != nil {
+		return nil, fmt.Errorf("error fetching virtual machines client: %w", initErr)
 	}
 
-	// At this point we have a list of distinct resources and a mapping between
-	// them and the set(s) they belong to. Iterate through them and fetch
-	// information.
-	//
-	// Listing returns IDs which are basically URLs that have key-value
-	// pairs but not the individual components. Get requires the individual
-	// components, and won't accept the ID...
-	//
-	// Anyways, now we find VM details for those resources
-	vmToIfaceMap := make(map[string][]network.Interface)
-	{
-		for _, res := range resourceInfos {
-			resourceGroup, name, err := splitId(*res.ID, constMsComputeService, constVirtualMachinesResource)
-			if err != nil {
-				return nil, fmt.Errorf("error splitting vm id %q: %w", *res.ID, err)
-			}
-			vm, err := vmClient.Get(ctx, resourceGroup, name, "")
-			if err != nil {
-				return nil, fmt.Errorf("error fetching vm with id %q: %w", *res.ID, err)
-			}
-
-			iv, err := vmClient.InstanceView(ctx, resourceGroup, name)
-			if err != nil {
-				return nil, fmt.Errorf("error fetching instance view for vm with id %q: %w", *res.ID, err)
-			}
-			if iv.Statuses == nil {
-				return nil, fmt.Errorf("instance view statuses returned for vm with id %q is null", *res.ID)
-			}
-			var running bool
-			for _, s := range *iv.Statuses {
-				if s.Code == nil {
-					continue
-				}
-				state := strings.ToLower(*s.Code)
-				prefix := "powerstate/"
-				if !strings.HasPrefix(state, prefix) {
-					continue
-				}
-				state = strings.TrimPrefix(state, prefix)
-				if state == "running" {
-					running = true
-					break
-				}
-			}
-			if !running {
-				continue
-			}
-
-			props := vm.VirtualMachineProperties
-			if props == nil {
-				return nil, fmt.Errorf("error fetching properties for vm with id %q: %w", *res.ID, err)
-			}
-			if props.NetworkProfile == nil {
-				return nil, fmt.Errorf("error fetching network profile for vm with id %q", *res.ID)
-			}
-
-			// Within the VM, catalog the various interfaces
-			var ifaces []network.Interface
-			for _, ifaceRef := range *props.NetworkProfile.NetworkInterfaces {
-				if ifaceRef.ID == nil {
-					return nil, fmt.Errorf("nil ID for network interface for vm with id %q", *res.ID)
-				}
-				ifResGroup, ifName, err := splitId(*ifaceRef.ID, constMsNetworkService, constNetworkInterfacesResource)
-				if err != nil {
-					return nil, fmt.Errorf("error splitting network interface id %q: %w", *ifaceRef.ID, err)
-				}
-				iface, err := ifClient.Get(ctx, ifResGroup, ifName, "")
-				if err != nil {
-					return nil, fmt.Errorf("error fetching network interface with id %q: %w", *ifaceRef.ID, err)
-				}
-				if iface.InterfacePropertiesFormat == nil || iface.InterfacePropertiesFormat.IPConfigurations == nil {
-					continue
-				}
-				ifaces = append(ifaces, iface)
-			}
-			vmToIfaceMap[*res.ID] = ifaces
-		}
+	if clients.vmssClient, initErr = getVirtualMachineScaleSetsClient(commonOpts...); initErr != nil {
+		return nil, fmt.Errorf("error fetching virtual machine scale set client: %w", initErr)
 	}
 
-	// Now, fetch IP details for each of the interfaces for each VM
-	type networkInfo struct {
-		IpAddresses []string
-	}
-	vmToNetworkMap := make(map[string]networkInfo)
-	{
-		for vmId, ifaces := range vmToIfaceMap {
-			var netInfo networkInfo
-			for _, iface := range ifaces {
-				for _, ipconf := range *iface.InterfacePropertiesFormat.IPConfigurations {
-					if ipconf.PrivateIPAddress != nil {
-						netInfo.IpAddresses = append(netInfo.IpAddresses, *ipconf.PrivateIPAddress)
-					}
-					if ipconf.PublicIPAddress != nil {
-						if ipconf.PublicIPAddress.ID == nil {
-							return nil, fmt.Errorf("ip configuration %q has public IP address info but nil id", *ipconf.ID)
-						}
-						ipResGroup, ipName, err := splitId(*ipconf.PublicIPAddress.ID, constMsNetworkService, constPublicIpAddressesResource)
-						if err != nil {
-							return nil, fmt.Errorf("error splitting public ip address id %q: %w", *ipconf.PublicIPAddress.ID, err)
-						}
-						pipInfo, err := pipClient.Get(ctx, ipResGroup, ipName, "")
-						if err != nil {
-							if err != nil {
-								return nil, fmt.Errorf("error fetching public IP information with resource group %q and name %q: %w", ipResGroup, ipName, err)
-							}
-						}
-						if pipInfo.PublicIPAddressPropertiesFormat == nil {
-							return nil, fmt.Errorf("nil public ip address properties format for public ip %q", *ipconf.PublicIPAddress.ID)
-						}
-						if pipInfo.PublicIPAddressPropertiesFormat.IPAddress != nil {
-							netInfo.IpAddresses = append(netInfo.IpAddresses, *pipInfo.PublicIPAddressPropertiesFormat.IPAddress)
-						}
-					}
-				}
-			}
-			vmToNetworkMap[vmId] = netInfo
-		}
+	if clients.vmssvmClient, initErr = getVirtualMachineScaleSetVMsClient(commonOpts...); initErr != nil {
+		return nil, fmt.Errorf("error fetching virtual machine scale set vms client: %w", initErr)
 	}
 
-	ret := &pb.ListHostsResponse{}
-	for k, v := range vmToNetworkMap {
-		_, externalName, err := splitId(k, constMsComputeService, constVirtualMachinesResource)
-		if err != nil {
-			externalName = ""
-		}
-
-		host := &pb.ListHostsResponseHost{
-			ExternalId:   k,
-			ExternalName: externalName,
-			IpAddresses:  v.IpAddresses,
-			SetIds:       resourceToSetMap[k],
-		}
-		ret.Hosts = append(ret.Hosts, host)
+	if clients.ifClient, initErr = getNetworkInterfacesClient(commonOpts...); initErr != nil {
+		return nil, fmt.Errorf("error fetching network interfaces client: %w", initErr)
 	}
-	return ret, nil
+
+	if clients.pipClient, initErr = getPublicIpAddressesClient(commonOpts...); initErr != nil {
+		return nil, fmt.Errorf("error fetching public ip address client: %w", initErr)
+	}
+
+	return clients, nil
 }
 
+func (p *AzurePlugin) findMatchingResources(ctx context.Context,
+	sets []*hostsets.HostSet,
+	resClient *resources.Client) (
+	vmResources []resources.GenericResourceExpanded,
+	vmssResources []resources.GenericResourceExpanded,
+	resourceToSetMap map[string][]string,
+	_ error) {
+
+	resourceToSetMap = make(map[string][]string, len(sets)*10) // Pre-allocate map size
+
+	for _, set := range sets {
+		if err := validateHostSet(set); err != nil {
+			return nil, nil, nil, err
+		}
+
+		filter := getSetFilter(set)
+		setVMResources, setVMSSResources, setResourceToSetMap, err := listAndFilterResources(ctx, resClient, filter, set.GetId())
+		if err != nil {
+			return nil, nil, nil, err
+		}
+		vmResources = append(vmResources, setVMResources...)
+		vmssResources = append(vmssResources, setVMSSResources...)
+		for k, v := range setResourceToSetMap {
+			resourceToSetMap[k] = append(resourceToSetMap[k], v...)
+		}
+	}
+
+	return vmResources, vmssResources, resourceToSetMap, nil
+}
+
+func listAndFilterResources(ctx context.Context,
+	resClient *resources.Client,
+	filter, setID string) (
+	vmResources []resources.GenericResourceExpanded,
+	vmssResources []resources.GenericResourceExpanded,
+	resourceToSetMap map[string][]string,
+	_ error) {
+
+	resourceToSetMap = make(map[string][]string)
+	iter, err := resClient.ListComplete(ctx, filter, "", nil)
+	if err != nil {
+		return nil, nil, nil, fmt.Errorf("error listing resources: %w", err)
+	}
+
+	for iter.NotDone() {
+		val := iter.Value()
+		if err := iter.NextWithContext(ctx); err != nil {
+			return nil, nil, nil, fmt.Errorf("error iterating resources: %w", err)
+		}
+
+		if val.ID == nil || val.Type == nil {
+			continue
+		}
+		virtualMachineType := constMsComputeService + "/" + constVirtualMachinesResource
+		virtualMachineScaleSetType := constMsComputeService + "/" + constVirtualMachineScaleSetsResource
+		if strings.EqualFold(*val.Type, virtualMachineType) {
+			vmResources = append(vmResources, val)
+			resourceToSetMap[*val.ID] = append(resourceToSetMap[*val.ID], setID)
+		} else if strings.EqualFold(*val.Type, virtualMachineScaleSetType) {
+			vmssResources = append(vmssResources, val)
+			resourceToSetMap[*val.ID] = append(resourceToSetMap[*val.ID], setID)
+		}
+	}
+
+	return vmResources, vmssResources, resourceToSetMap, nil
+}
+
+// Resource Filter Management
+func getSetFilter(set *hostsets.HostSet) string {
+	setAttrFields := set.GetAttributes().GetFields()
+	if setAttrFields[constFilter] != nil {
+		if filter := setAttrFields[constFilter].GetStringValue(); filter != "" {
+			return filter
+		}
+	}
+	return ""
+}
+
+func validateHostSet(set *hostsets.HostSet) error {
+	if set == nil {
+		return errors.New("set is nil")
+	}
+	setAttrs := set.GetAttributes()
+	if setAttrs == nil {
+		return fmt.Errorf("set %s has no attributes", set.GetId())
+	}
+	if len(setAttrs.GetFields()) == 0 {
+		return fmt.Errorf("set %s attributes has no filter", set.GetId())
+	}
+	return nil
+}
+
+// ID Parsing Functions
 func splitId(in, expectedService, expectedResource string) (string, string, error) {
 	// Note: this could be stolen from the helpers directory in the TF AzureRM
 	// provider but I'm not at the moment because it does a bunch of trickery
-	// for some Azure APIs and we have a much more limited scope. I may come to
-	// regret this.
+	// for some Azure APIs and we have a much more limited scope.
 	splitId := strings.Split(strings.TrimLeft(in, "/"), "/")
+
 	// Run through some sanity checks
 	if len(splitId) != 8 ||
 		!strings.EqualFold(splitId[0], constSubscriptions) ||
@@ -423,47 +464,39 @@ func splitId(in, expectedService, expectedResource string) (string, string, erro
 	return splitId[3], splitId[7], nil
 }
 
-// Returns an invalid argument error with the problematic fields included
-// in the error message.
-func invalidArgumentError(msg string, f map[string]string) error {
-	var fieldMsgs []string
-	for field, val := range f {
-		fieldMsgs = append(fieldMsgs, fmt.Sprintf("%q: %q", field, val))
+func extractResourceType(resourceID string) (string, error) {
+	if resourceID == "" {
+		return "", fmt.Errorf("empty resource ID provided")
 	}
-	if len(fieldMsgs) > 0 {
-		msg = fmt.Sprintf("%s: [%s]", msg, strings.Join(fieldMsgs, ", "))
+
+	parts := strings.Split(strings.TrimLeft(resourceID, "/"), "/")
+
+	// Validate minimum length and basic structure
+	if len(parts) < 6 {
+		return "", fmt.Errorf("resource ID too short: %s", resourceID)
 	}
-	return status.Error(codes.InvalidArgument, msg)
-}
-
-func validateSecrets(s *structpb.Struct) error {
-	if s == nil {
-		return status.Error(codes.InvalidArgument, "Secrets not provided but are required")
-	}
-	var secrets SecretData
-	if err := mapstructure.Decode(s.AsMap(), &secrets); err != nil {
-		return status.Errorf(codes.InvalidArgument, "error decoding catalog secrets: %s", err)
-	}
-	badFields := make(map[string]string)
-	if len(secrets.SecretValue) == 0 {
-		badFields["secrets.secret_value"] = "This field is required."
-	}
-	if len(secrets.CredsLastRotatedTime) != 0 {
-		badFields["secrets.creds_last_rotated_time"] = "This field is reserved and cannot be set."
-	}
-	if len(badFields) > 0 {
-		return invalidArgumentError("Error in the secrets provided", badFields)
-	}
-	return nil
-}
-
-var allowedCatalogFields = map[string]struct{}{
-	constDisableCredentialRotation: {},
-	constSubscriptionId:            {},
-	constTenantId:                  {},
-	constClientId:                  {},
+
+	// Check for required segments in correct positions
+	if !strings.EqualFold(parts[0], constSubscriptions) ||
+		!strings.EqualFold(parts[2], constResourceGroups) ||
+		!strings.EqualFold(parts[4], constProviders) {
+		return "", fmt.Errorf("malformed resource ID, missing required segments: %s", resourceID)
+	}
+
+	// Check for Microsoft.Compute in correct position
+	if !strings.EqualFold(parts[5], constMsComputeService) {
+		return "", nil // Not a compute resource, return empty string without error
+	}
+
+	// Ensure there's a resource type after Microsoft.Compute
+	if len(parts) < 7 {
+		return "", fmt.Errorf("no resource type found after Microsoft.Compute: %s", resourceID)
+	}
+
+	return parts[6], nil
 }
 
+// Validation Functions
 func validateCatalog(c *hostcatalogs.HostCatalog) error {
 	if c == nil {
 		return status.Error(codes.InvalidArgument, "catalog is nil")
@@ -473,6 +506,7 @@ func validateCatalog(c *hostcatalogs.HostCatalog) error {
 	if err := mapstructure.Decode(attrMap, &attrs); err != nil {
 		return status.Errorf(codes.InvalidArgument, "error decoding catalog attributes: %s", err)
 	}
+
 	badFields := make(map[string]string)
 	if !attrs.DisableCredentialRotation {
 		badFields["attributes.disable_credential_rotation"] = "This field must be set to true."
@@ -499,8 +533,26 @@ func validateCatalog(c *hostcatalogs.HostCatalog) error {
 	return nil
 }
 
-var allowedSetFields = map[string]struct{}{
-	constFilter: {},
+func validateSecrets(s *structpb.Struct) error {
+	if s == nil {
+		return status.Error(codes.InvalidArgument, "Secrets not provided but are required")
+	}
+	var secrets SecretData
+	if err := mapstructure.Decode(s.AsMap(), &secrets); err != nil {
+		return status.Errorf(codes.InvalidArgument, "error decoding catalog secrets: %s", err)
+	}
+
+	badFields := make(map[string]string)
+	if len(secrets.SecretValue) == 0 {
+		badFields["secrets.secret_value"] = "This field is required."
+	}
+	if len(secrets.CredsLastRotatedTime) != 0 {
+		badFields["secrets.creds_last_rotated_time"] = "This field is reserved and cannot be set."
+	}
+	if len(badFields) > 0 {
+		return invalidArgumentError("Error in the secrets provided", badFields)
+	}
+	return nil
 }
 
 func validateSet(s *hostsets.HostSet) error {
@@ -512,6 +564,7 @@ func validateSet(s *hostsets.HostSet) error {
 	if err := mapstructure.Decode(attrMap, &attrs); err != nil {
 		return status.Errorf(codes.InvalidArgument, "error decoding set attributes: %s", err)
 	}
+
 	badFields := make(map[string]string)
 	if _, ok := attrMap[constFilter]; ok && len(attrs.Filter) == 0 {
 		badFields["attributes.filter"] = "This field must be not empty."
@@ -528,3 +581,27 @@ func validateSet(s *hostsets.HostSet) error {
 	}
 	return nil
 }
+
+// Error Handling Functions
+func invalidArgumentError(msg string, f map[string]string) error {
+	var fieldMsgs []string
+	for field, val := range f {
+		fieldMsgs = append(fieldMsgs, fmt.Sprintf("%q: %q", field, val))
+	}
+	if len(fieldMsgs) > 0 {
+		msg = fmt.Sprintf("%s: [%s]", msg, strings.Join(fieldMsgs, ", "))
+	}
+	return status.Error(codes.InvalidArgument, msg)
+}
+
+// Map Definitions for Allowed Fields
+var allowedCatalogFields = map[string]struct{}{
+	constDisableCredentialRotation: {},
+	constSubscriptionId:            {},
+	constTenantId:                  {},
+	constClientId:                  {},
+}
+
+var allowedSetFields = map[string]struct{}{
+	constFilter: {},
+}

plugin/plugin_test.go

@@ -7,6 +7,8 @@ import (
 	"context"
 	"os"
 	"path/filepath"
+	"reflect"
+	"strings"
 	"testing"
 	"time"
 
@@ -448,7 +450,7 @@ func testGetHostStructs(t *testing.T) (*hostcatalogs.HostCatalog, []*hostsets.Ho
 			Id: "set1",
 			Attrs: &hostsets.HostSet_Attributes{
 				Attributes: wrapMap(t, map[string]interface{}{
-					constFilter: constDefaultFilter,
+					constFilter: "",
 				}),
 			},
 		},
@@ -456,7 +458,7 @@ func testGetHostStructs(t *testing.T) (*hostcatalogs.HostCatalog, []*hostsets.Ho
 			Id: "set2",
 			Attrs: &hostsets.HostSet_Attributes{
 				Attributes: wrapMap(t, map[string]interface{}{
-					constFilter: constDefaultFilter,
+					constFilter: "",
 				}),
 			},
 		},
@@ -505,3 +507,287 @@ func testCreateInitialSecret(t *testing.T, hostCatalog *hostcatalogs.HostCatalog
 
 	return authzInfo, initialCred, cleanup
 }
+
+func TestSplitId(t *testing.T) {
+	tests := []struct {
+		name             string
+		input            string
+		expectedService  string
+		expectedResource string
+		wantRG           string
+		wantName         string
+		wantErr          bool
+	}{
+		{
+			name:             "valid resource id",
+			input:            "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachines/vm1",
+			expectedService:  "Microsoft.Compute",
+			expectedResource: "virtualMachines",
+			wantRG:           "rg1",
+			wantName:         "vm1",
+			wantErr:          false,
+		},
+		{
+			name:             "valid with different case",
+			input:            "/SUBSCRIPTIONS/sub1/RESOURCEGROUPS/rg1/PROVIDERS/MICROSOFT.COMPUTE/VIRTUALMACHINES/vm1",
+			expectedService:  "Microsoft.Compute",
+			expectedResource: "virtualMachines",
+			wantRG:           "rg1",
+			wantName:         "vm1",
+			wantErr:          false,
+		},
+		{
+			name:             "invalid - trailing slash",
+			input:            "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachines/vm1/",
+			expectedService:  "Microsoft.Compute",
+			expectedResource: "virtualMachines",
+			wantRG:           "",
+			wantName:         "",
+			wantErr:          true,
+		},
+		{
+			name:             "invalid - wrong service",
+			input:            "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Storage/virtualMachines/vm1",
+			expectedService:  "Microsoft.Compute",
+			expectedResource: "virtualMachines",
+			wantRG:           "",
+			wantName:         "",
+			wantErr:          true,
+		},
+		{
+			name:             "invalid - wrong resource type",
+			input:            "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/diskEncryptionSets/des1",
+			expectedService:  "Microsoft.Compute",
+			expectedResource: "virtualMachines",
+			wantRG:           "",
+			wantName:         "",
+			wantErr:          true,
+		},
+		{
+			name:             "invalid - missing segments",
+			input:            "/subscriptions/sub1/resourceGroups/rg1",
+			expectedService:  "Microsoft.Compute",
+			expectedResource: "virtualMachines",
+			wantRG:           "",
+			wantName:         "",
+			wantErr:          true,
+		},
+		{
+			name:             "invalid - empty string",
+			input:            "",
+			expectedService:  "Microsoft.Compute",
+			expectedResource: "virtualMachines",
+			wantRG:           "",
+			wantName:         "",
+			wantErr:          true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			gotRG, gotName, err := splitId(tt.input, tt.expectedService, tt.expectedResource)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("splitId() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if gotRG != tt.wantRG {
+				t.Errorf("splitId() gotRG = %v, want %v", gotRG, tt.wantRG)
+			}
+			if gotName != tt.wantName {
+				t.Errorf("splitId() gotName = %v, want %v", gotName, tt.wantName)
+			}
+		})
+	}
+}
+
+func TestExtractResourceType(t *testing.T) {
+	tests := []struct {
+		name    string
+		input   string
+		want    string
+		wantErr bool
+	}{
+		{
+			name:    "valid compute resource",
+			input:   "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachines/vm1",
+			want:    "virtualMachines",
+			wantErr: false,
+		},
+		{
+			name:    "valid compute resource different case",
+			input:   "/subscriptions/sub1/resourceGroups/rg1/providers/MICROSOFT.COMPUTE/virtualMachines/vm1",
+			want:    "virtualMachines",
+			wantErr: false,
+		},
+		{
+			name:    "valid with trailing slash",
+			input:   "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachines/vm1/",
+			want:    "virtualMachines",
+			wantErr: false,
+		},
+		{
+			name:    "different compute resource type",
+			input:   "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/diskEncryptionSets/des1",
+			want:    "diskEncryptionSets",
+			wantErr: false,
+		},
+		{
+			name:    "non-compute resource",
+			input:   "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Storage/storageAccounts/sa1",
+			want:    "",
+			wantErr: false,
+		},
+		{
+			name:    "malformed - missing provider",
+			input:   "/subscriptions/sub1/resourceGroups/rg1/Microsoft.Compute/virtualMachines/vm1",
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name:    "malformed - missing resource type",
+			input:   "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute",
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name:    "empty string",
+			input:   "",
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name:    "too short",
+			input:   "/subscriptions/sub1",
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name:    "malformed structure",
+			input:   "/something/else/entirely",
+			want:    "",
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := extractResourceType(tt.input)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("extractResourceType() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if got != tt.want {
+				t.Errorf("extractResourceType() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestCreateHostFromResource(t *testing.T) {
+	// Setup common test data
+	testNetwork := networkInfo{
+		IpAddresses: []string{"192.168.1.1", "192.168.1.2"},
+	}
+
+	// Use full Azure resource paths as keys
+	vmssPath := "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachineScaleSets/vmss1"
+	vmPath := "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachines/vm-resource1"
+
+	resourceToSetMap := map[string][]string{
+		vmssPath: {"set1", "set2"},
+		vmPath:   {"set3", "set4"},
+	}
+
+	tests := []struct {
+		name       string
+		resourceID string
+		network    networkInfo
+		setMap     map[string][]string
+		want       *pb.ListHostsResponseHost
+		wantErr    bool
+		errMessage string
+	}{
+		{
+			name:       "valid VMSS resource",
+			resourceID: vmssPath + "/virtualMachines/vm1",
+			network:    testNetwork,
+			setMap:     resourceToSetMap,
+			want: &pb.ListHostsResponseHost{
+				ExternalId:   vmssPath + "/virtualMachines/vm1",
+				ExternalName: "vmss1_vm1",
+				IpAddresses:  []string{"192.168.1.1", "192.168.1.2"},
+				SetIds:       []string{"set1", "set2"},
+			},
+			wantErr: false,
+		},
+		{
+			name:       "valid VM resource",
+			resourceID: vmPath,
+			network:    testNetwork,
+			setMap:     resourceToSetMap,
+			want: &pb.ListHostsResponseHost{
+				ExternalId:   vmPath,
+				ExternalName: "vm-resource1",
+				IpAddresses:  []string{"192.168.1.1", "192.168.1.2"},
+				SetIds:       []string{"set3", "set4"},
+			},
+			wantErr: false,
+		},
+		{
+			name:       "invalid resource ID",
+			resourceID: "invalid-resource-id",
+			network:    testNetwork,
+			setMap:     resourceToSetMap,
+			want:       nil,
+			wantErr:    true,
+			errMessage: "invalid resource ID",
+		},
+		{
+			name:       "unsupported resource type",
+			resourceID: "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/unsupportedResource/resource1",
+			network:    testNetwork,
+			setMap:     resourceToSetMap,
+			want:       nil,
+			wantErr:    true,
+			errMessage: "unsupported resource type",
+		},
+		{
+			name:       "VMSS with valid resource type but invalid structure",
+			resourceID: "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachineScaleSets", // Missing required segments after virtualMachineScaleSets
+			network:    testNetwork,
+			setMap:     resourceToSetMap,
+			want:       nil,
+			wantErr:    true,
+			errMessage: "could not determine set for resource ID: not a valid VMSS resource ID",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := createHostFromResource(tt.resourceID, tt.network, tt.setMap)
+
+			// Check error cases
+			if tt.wantErr {
+				if err == nil {
+					t.Errorf("createHostFromResource() expected error containing %q, got nil", tt.errMessage)
+					return
+				}
+				if !strings.Contains(err.Error(), tt.errMessage) {
+					t.Errorf("createHostFromResource() error = %v, want error containing %q", err, tt.errMessage)
+				}
+				return
+			}
+
+			// Check non-error cases
+			if err != nil {
+				t.Errorf("createHostFromResource() unexpected error: %v", err)
+				return
+			}
+
+			// Compare the results
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("createHostFromResource() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

plugin/vm_processor.go

@@ -0,0 +1,131 @@
+package plugin
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"sync"
+
+	"github.com/Azure/azure-sdk-for-go/profiles/latest/compute/mgmt/compute"
+	"github.com/Azure/azure-sdk-for-go/profiles/latest/resources/mgmt/resources"
+	"golang.org/x/sync/errgroup"
+)
+
+// vmJob represents a single VM processing job
+type vmJob struct {
+	resource      resources.GenericResourceExpanded
+	resourceGroup string
+	name          string
+}
+
+func (p *AzurePlugin) processStandardVMs(ctx context.Context, resourceInfos []resources.GenericResourceExpanded,
+	clients *azureClients) (map[string]networkInfo, error) {
+
+	// Initialize the map that we'll return
+	vmToNetworkMap := make(map[string]networkInfo)
+
+	// Pre-validate all resource IDs and create jobs slice
+	jobs := make([]vmJob, 0, len(resourceInfos))
+	for _, res := range resourceInfos {
+		resourceGroup, name, err := splitId(*res.ID, constMsComputeService, constVirtualMachinesResource)
+		if err != nil {
+			return nil, fmt.Errorf("invalid vm id %q: %w", *res.ID, err)
+		}
+		jobs = append(jobs, vmJob{
+			resource:      res,
+			resourceGroup: resourceGroup,
+			name:          name,
+		})
+	}
+
+	// Create error group for worker management
+	g, ctx := errgroup.WithContext(ctx)
+
+	// Create job channel
+	jobsChan := make(chan vmJob)
+
+	// Result handling
+	mu := &sync.Mutex{}
+
+	// Start worker pool
+	for i := 0; i < constDefaultWorkers; i++ {
+		g.Go(func() error {
+			for job := range jobsChan {
+				// Process single VM
+				vm, err := clients.vmClient.Get(ctx, job.resourceGroup, job.name, "")
+				if err != nil {
+					return fmt.Errorf("error fetching vm with id %q: %w", *job.resource.ID, err)
+				}
+
+				// Only process running VMs
+				running, err := isVMRunning(ctx, clients.vmClient, job.resourceGroup, job.name)
+				if err != nil {
+					return err
+				}
+				if !running {
+					continue
+				}
+
+				var netInfo networkInfo
+				if err := processVMNetworkInterfaces(ctx, vm, clients, job.resourceGroup, &netInfo); err != nil {
+					return err
+				}
+
+				mu.Lock()
+				vmToNetworkMap[*job.resource.ID] = netInfo
+				mu.Unlock()
+			}
+			return nil
+		})
+	}
+
+	// Send jobs to workers
+	go func() {
+		defer close(jobsChan)
+		for _, job := range jobs {
+			// Check context cancellation
+			if ctx.Err() != nil {
+				return
+			}
+			jobsChan <- job
+		}
+	}()
+
+	// Wait for all workers to complete
+	if err := g.Wait(); err != nil {
+		return nil, err
+	}
+
+	return vmToNetworkMap, nil
+}
+
+// VM Status Check
+func isVMRunning(
+	ctx context.Context,
+	vmClient *compute.VirtualMachinesClient,
+	resourceGroup,
+	name string,
+) (
+	bool,
+	error,
+) {
+
+	iv, err := vmClient.InstanceView(ctx, resourceGroup, name)
+	if err != nil {
+		return false, fmt.Errorf("error fetching instance view for vm %s: %w", name, err)
+	}
+	if iv.Statuses == nil {
+		return false, fmt.Errorf("instance view statuses is null for vm %s", name)
+	}
+
+	for _, s := range *iv.Statuses {
+		if s.Code == nil {
+			continue
+		}
+		state := strings.ToLower(*s.Code)
+		if strings.HasPrefix(state, "powerstate/") && strings.TrimPrefix(state, "powerstate/") == "running" {
+			return true, nil
+		}
+	}
+	return false, nil
+}

plugin/vmss_processor.go

@@ -0,0 +1,302 @@
+package plugin
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"sync"
+
+	"github.com/Azure/azure-sdk-for-go/profiles/latest/compute/mgmt/compute"
+	"github.com/Azure/azure-sdk-for-go/profiles/latest/resources/mgmt/resources"
+	"golang.org/x/sync/errgroup"
+)
+
+// vmssJob represents a single VMSS processing job
+type vmssJob struct {
+	resource      resources.GenericResourceExpanded
+	resourceGroup string
+	name          string
+}
+
+// vmssInstanceJob represents a single VMSS instance processing job
+type vmssInstanceJob struct {
+	vmssvm        compute.VirtualMachineScaleSetVM
+	resourceGroup string
+	vmssName      string
+}
+
+func (p *AzurePlugin) processVMScaleSets(
+	ctx context.Context,
+	resourceInfos []resources.GenericResourceExpanded,
+	clients *azureClients,
+) (
+	map[string]networkInfo,
+	error,
+) {
+
+	// Initialize the map that we'll return
+	vmToNetworkMap := make(map[string]networkInfo)
+
+	// Pre-validate all resource IDs and create jobs slice
+	jobs := make([]vmssJob, 0, len(resourceInfos))
+	for _, res := range resourceInfos {
+		resourceGroup, name, err := splitId(*res.ID, constMsComputeService, constVirtualMachineScaleSetsResource)
+		if err != nil {
+			return nil, fmt.Errorf("invalid vmss id %q: %w", *res.ID, err)
+		}
+		jobs = append(jobs, vmssJob{
+			resource:      res,
+			resourceGroup: resourceGroup,
+			name:          name,
+		})
+	}
+
+	// Create error group for worker management
+	g, ctx := errgroup.WithContext(ctx)
+
+	// Create job channel
+	jobsChan := make(chan vmssJob)
+
+	// Start worker pool
+	for i := 0; i < constDefaultWorkers; i++ {
+		g.Go(func() error {
+			for job := range jobsChan {
+				vmss, err := clients.vmssClient.Get(ctx, job.resourceGroup, job.name, "")
+				if err != nil {
+					return fmt.Errorf("error fetching vmss with id %q: %w", *job.resource.ID, err)
+				}
+
+				// Check if VMSS properties exist
+				if vmss.VirtualMachineScaleSetProperties == nil {
+					continue
+				}
+				vmssProperties := vmss.VirtualMachineScaleSetProperties
+
+				// Skip if OrchestrationMode is empty or Flexible
+				if vmssProperties.OrchestrationMode == "" {
+					continue
+				}
+				if vmssProperties.OrchestrationMode == compute.OrchestrationModeFlexible {
+					continue
+				}
+
+				if err := processVMSSInstancesParallel(ctx, vmss, job.resourceGroup, clients, vmToNetworkMap); err != nil {
+					return fmt.Errorf("error processing uniform vmss instances: %w", err)
+				}
+			}
+			return nil
+		})
+	}
+
+	// Send jobs to workers
+	go func() {
+		defer close(jobsChan)
+		for _, job := range jobs {
+			if ctx.Err() != nil {
+				return
+			}
+			jobsChan <- job
+		}
+	}()
+
+	// Wait for all workers to complete
+	if err := g.Wait(); err != nil {
+		return nil, err
+	}
+
+	return vmToNetworkMap, nil
+}
+
+func processVMSSInstancesParallel(
+	ctx context.Context,
+	vmss compute.VirtualMachineScaleSet,
+	resourceGroup string,
+	clients *azureClients,
+	vmToNetworkMap map[string]networkInfo,
+) error {
+
+	if vmss.Name == nil {
+		return fmt.Errorf("vmss name is nil")
+	}
+
+	vmssvms, err := clients.vmssvmClient.List(ctx, resourceGroup, *vmss.Name, "", "", "")
+	if err != nil {
+		return fmt.Errorf("error fetching vms for vmss %s: %w", *vmss.Name, err)
+	}
+
+	// Create new error group for instance processing
+	g, ctx := errgroup.WithContext(ctx)
+	instanceJobs := make(chan vmssInstanceJob)
+
+	// Start instance workers
+	instanceWorkerCount := constDefaultWorkers
+	if len(vmssvms.Values()) < instanceWorkerCount {
+		instanceWorkerCount = len(vmssvms.Values())
+	}
+
+	mu := &sync.Mutex{}
+
+	// Start instance worker pool
+	for i := 0; i < instanceWorkerCount; i++ {
+		g.Go(func() error {
+			for job := range instanceJobs {
+				if job.vmssvm.ID == nil {
+					return fmt.Errorf("nil ID for vmss vm in scale set %s", *vmss.Name)
+				}
+
+				if err := processVMSSInstanceParallel(ctx, job.vmssvm, job.resourceGroup,
+					job.vmssName, clients, vmToNetworkMap, mu); err != nil {
+					return fmt.Errorf("error processing vmss instance %s: %w", *job.vmssvm.ID, err)
+				}
+			}
+			return nil
+		})
+	}
+
+	// Send instance jobs
+	go func() {
+		defer close(instanceJobs)
+		for _, vmssvm := range vmssvms.Values() {
+			select {
+			case <-ctx.Done():
+				return
+			default:
+			}
+
+			instanceJobs <- vmssInstanceJob{
+				vmssvm:        vmssvm,
+				resourceGroup: resourceGroup,
+				vmssName:      *vmss.Name,
+			}
+		}
+	}()
+
+	return g.Wait()
+}
+
+func processVMSSInstanceParallel(
+	ctx context.Context,
+	vmssvm compute.VirtualMachineScaleSetVM,
+	resourceGroup,
+	vmssName string,
+	clients *azureClients,
+	vmToNetworkMap map[string]networkInfo,
+	mu *sync.Mutex,
+) error {
+
+	if vmssvm.InstanceID == nil {
+		return fmt.Errorf("instance ID is nil")
+	}
+
+	// Check if instance is running
+	if running, err := isVMSSInstanceRunning(ctx, clients.vmssvmClient, resourceGroup,
+		vmssName, vmssvm.InstanceID); err != nil {
+		return err
+	} else if !running {
+		return nil
+	}
+
+	// Process network interfaces for the instance
+	var netInfo networkInfo
+	if err := processVMSSNetworkInterfaces(ctx, vmssvm, resourceGroup, vmssName, clients, &netInfo); err != nil {
+		return err
+	}
+
+	func() {
+		mu.Lock()
+		defer mu.Unlock()
+		vmToNetworkMap[*vmssvm.ID] = netInfo
+	}()
+
+	return nil
+}
+
+// VMSS Instance Status Check
+func isVMSSInstanceRunning(
+	ctx context.Context,
+	vmssvmClient *compute.VirtualMachineScaleSetVMsClient,
+	resourceGroup,
+	vmssName string,
+	instanceID *string,
+) (
+	bool,
+	error,
+) {
+
+	if instanceID == nil {
+		return false, fmt.Errorf("instance ID is nil")
+	}
+
+	iv, err := vmssvmClient.GetInstanceView(ctx, resourceGroup, vmssName, *instanceID)
+	if err != nil {
+		return false, fmt.Errorf("error fetching instance view for instance %s: %w", *instanceID, err)
+	}
+	if iv.Statuses == nil {
+		return false, fmt.Errorf("instance view statuses returned for instance %s is null", *instanceID)
+	}
+
+	for _, s := range *iv.Statuses {
+		if s.Code == nil {
+			continue
+		}
+		state := strings.ToLower(*s.Code)
+		prefix := "powerstate/"
+		if !strings.HasPrefix(state, prefix) {
+			continue
+		}
+		state = strings.TrimPrefix(state, prefix)
+		if state == "running" {
+			return true, nil
+		}
+	}
+	return false, nil
+}
+
+// Helper Functions
+func getSetForVMSSInstance(in string) (string, error) {
+	if in == "" {
+		return "", fmt.Errorf("empty resource ID provided")
+	}
+
+	// Split the path into segments
+	segments := strings.Split(strings.TrimLeft(in, "/"), "/")
+
+	// Check if this is a VMSS path by looking for required segments
+	if len(segments) < 8 ||
+		!strings.EqualFold(segments[0], constSubscriptions) ||
+		!strings.EqualFold(segments[2], constResourceGroups) ||
+		!strings.EqualFold(segments[4], constProviders) ||
+		!strings.EqualFold(segments[5], constMsComputeService) ||
+		!strings.EqualFold(segments[6], constVirtualMachineScaleSetsResource) {
+		return "", fmt.Errorf("not a valid VMSS resource ID: %s", in)
+	}
+
+	// Reconstruct the VMSS path (everything up to but not including virtualMachines)
+	vmssPath := "/" + strings.Join(segments[:8], "/")
+
+	// Additional validation: if the path contains virtualMachines, it should be in the correct position
+	for i, segment := range segments[8:] {
+		if strings.EqualFold(segment, constVirtualMachinesResource) {
+			if i != 0 { // virtualMachines should be immediately after VMSS name
+				return "", fmt.Errorf("unexpected virtualMachines segment position in resource ID: %s", in)
+			}
+			break
+		}
+	}
+
+	return vmssPath, nil
+}
+
+func getExternalNameforVMSSInstance(in string) (string, error) {
+	splitId := strings.Split(strings.TrimLeft(in, "/"), "/")
+	if len(splitId) != 10 ||
+		!strings.EqualFold(splitId[0], constSubscriptions) ||
+		!strings.EqualFold(splitId[2], constResourceGroups) ||
+		!strings.EqualFold(splitId[4], constProviders) ||
+		!strings.EqualFold(splitId[5], constMsComputeService) ||
+		!strings.EqualFold(splitId[6], constVirtualMachineScaleSetsResource) ||
+		!strings.EqualFold(splitId[8], constVirtualMachinesResource) {
+		return "", fmt.Errorf("unexpected format of virtual machine stateful set ID: %v", splitId)
+	}
+	return splitId[7] + "_" + splitId[9], nil
+}

plugin/vmss_processor_test.go

@@ -0,0 +1,153 @@
+package plugin
+
+import (
+	"testing"
+)
+
+func TestGetExternalNameforVMSSInstance(t *testing.T) {
+	tests := []struct {
+		name    string
+		input   string
+		want    string
+		wantErr bool
+	}{
+		{
+			name:    "valid vmss instance id",
+			input:   "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachineScaleSets/vmss1/virtualMachines/instance1",
+			want:    "vmss1_instance1",
+			wantErr: false,
+		},
+		{
+			name:    "invalid - vmss instance id with trailing slash",
+			input:   "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachineScaleSets/vmss1/virtualMachines/instance1/",
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name:    "case insensitive segment validation",
+			input:   "/SUBSCRIPTIONS/sub1/RESOURCEGROUPS/rg1/PROVIDERS/MICROSOFT.COMPUTE/VIRTUALMACHINESCALESETS/vmss1/VIRTUALMACHINES/instance1",
+			want:    "vmss1_instance1",
+			wantErr: false,
+		},
+		{
+			name:    "invalid - missing segment",
+			input:   "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachineScaleSets/vmss1",
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name:    "invalid - wrong resource type",
+			input:   "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Storage/storageAccounts/sa1",
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name:    "invalid - empty string",
+			input:   "",
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name:    "invalid - malformed id",
+			input:   "invalid/format/string",
+			want:    "",
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := getExternalNameforVMSSInstance(tt.input)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("getExternalNameforVMSSInstance() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if got != tt.want {
+				t.Errorf("getExternalNameforVMSSInstance() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestGetSetForVMSSInstance(t *testing.T) {
+	tests := []struct {
+		name    string
+		input   string
+		want    string
+		wantErr bool
+	}{
+		{
+			name:    "valid vmss instance id",
+			input:   "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachineScaleSets/vmss1/virtualMachines/instance1",
+			want:    "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachineScaleSets/vmss1",
+			wantErr: false,
+		},
+		{
+			name:    "valid vmss path without instance",
+			input:   "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachineScaleSets/vmss1",
+			want:    "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachineScaleSets/vmss1",
+			wantErr: false,
+		},
+		{
+			name:    "case insensitive validation",
+			input:   "/SUBSCRIPTIONS/sub1/RESOURCEGROUPS/rg1/PROVIDERS/MICROSOFT.COMPUTE/VIRTUALMACHINESCALESETS/vmss1",
+			want:    "/SUBSCRIPTIONS/sub1/RESOURCEGROUPS/rg1/PROVIDERS/MICROSOFT.COMPUTE/VIRTUALMACHINESCALESETS/vmss1",
+			wantErr: false,
+		},
+		{
+			name:    "path with trailing slash",
+			input:   "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachineScaleSets/vmss1/virtualMachines/instance1/",
+			want:    "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachineScaleSets/vmss1",
+			wantErr: false,
+		},
+		{
+			name:    "empty string",
+			input:   "",
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name:    "invalid - flexible orchestration mode VM path",
+			input:   "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachines/vmss1_instance1",
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name:    "invalid - wrong provider",
+			input:   "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Storage/virtualMachineScaleSets/vmss1",
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name:    "invalid - malformed path missing segments",
+			input:   "/subscriptions/sub1/resourceGroups",
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name:    "invalid - malformed with virtualMachines in wrong position",
+			input:   "/subscriptions/sub1/resourceGroups/rg1/virtualMachines/providers/Microsoft.Compute/virtualMachineScaleSets/vmss1",
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name:    "invalid - random string",
+			input:   "not/a/valid/path",
+			want:    "",
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := getSetForVMSSInstance(tt.input)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("getSetForVMSSInstance() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if got != tt.want {
+				t.Errorf("getSetForVMSSInstance() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}