Filter while paging in GetByVMID (#50000)

This updates GetByVMID, used by azure joining, to filter by id for
each page of vms instead of filtering after accumulating _all_ vms
in memory first. While this does reduce memory consumption, the listing
process is still suboptimal in that _all_ instances in a subscription
are retrieved in response to _every_ join request.

lib/auth/join_azure.go

@@ -269,7 +269,7 @@ func verifyVMIdentity(ctx context.Context, cfg *azureRegisterConfig, accessToken
 		// If the token is from a user-assigned managed identity, the resource ID is
 		// for the identity and we need to look the VM up by VM ID.
 	} else {
-		vm, err = vmClient.GetByVMID(ctx, types.Wildcard, vmID)
+		vm, err = vmClient.GetByVMID(ctx, vmID)
 		if err != nil {
 			if trace.IsNotFound(err) {
 				return nil, trace.AccessDenied("no VM found with matching VM ID")

lib/auth/join_azure_test.go

@@ -73,13 +73,13 @@ func (m *mockAzureVMClient) Get(_ context.Context, resourceID string) (*azure.Vi
 	return vm, nil
 }
 
-func (m *mockAzureVMClient) GetByVMID(_ context.Context, resourceGroup, vmID string) (*azure.VirtualMachine, error) {
+func (m *mockAzureVMClient) GetByVMID(_ context.Context, vmID string) (*azure.VirtualMachine, error) {
 	for _, vm := range m.vms {
-		if vm.VMID == vmID && (resourceGroup == types.Wildcard || vm.ResourceGroup == resourceGroup) {
+		if vm.VMID == vmID {
 			return vm, nil
 		}
 	}
-	return nil, trace.NotFound("no vm in groups %q with id %q", resourceGroup, vmID)
+	return nil, trace.NotFound("no vm with id %q", vmID)
 }
 
 func makeVMClientGetter(clients map[string]*mockAzureVMClient) vmClientGetter {

lib/cloud/azure/vm.go

@@ -46,7 +46,7 @@ type VirtualMachinesClient interface {
 	// Get returns the virtual machine for the given resource ID.
 	Get(ctx context.Context, resourceID string) (*VirtualMachine, error)
 	// GetByVMID returns the virtual machine for a given VM ID.
-	GetByVMID(ctx context.Context, resourceGroup, vmID string) (*VirtualMachine, error)
+	GetByVMID(ctx context.Context, vmID string) (*VirtualMachine, error)
 	// ListVirtualMachines gets all of the virtual machines in the given resource group.
 	ListVirtualMachines(ctx context.Context, resourceGroup string) ([]*armcompute.VirtualMachine, error)
 }
@@ -146,15 +146,19 @@ func (c *vmClient) Get(ctx context.Context, resourceID string) (*VirtualMachine,
 }
 
 // GetByVMID returns the virtual machine for a given VM ID.
-func (c *vmClient) GetByVMID(ctx context.Context, resourceGroup, vmID string) (*VirtualMachine, error) {
-	vms, err := c.ListVirtualMachines(ctx, resourceGroup)
-	if err != nil {
-		return nil, trace.Wrap(err)
-	}
-	for _, vm := range vms {
-		if vm.Properties != nil && *vm.Properties.VMID == vmID {
-			result, err := parseVirtualMachine(vm)
-			return result, trace.Wrap(err)
+func (c *vmClient) GetByVMID(ctx context.Context, vmID string) (*VirtualMachine, error) {
+	pager := newListAllPager(c.api.NewListAllPager(&armcompute.VirtualMachinesClientListAllOptions{}))
+	for pager.more() {
+		res, err := pager.nextPage(ctx)
+		if err != nil {
+			return nil, trace.Wrap(ConvertResponseError(err))
+		}
+
+		for _, vm := range res {
+			if vm.Properties != nil && *vm.Properties.VMID == vmID {
+				result, err := parseVirtualMachine(vm)
+				return result, trace.Wrap(err)
+			}
 		}
 	}
 	return nil, trace.NotFound("no VM with ID %q", vmID)

lib/srv/discovery/discovery_test.go

@@ -2302,7 +2302,7 @@ func (m *mockAzureClient) Get(_ context.Context, _ string) (*azure.VirtualMachin
 	return nil, nil
 }
 
-func (m *mockAzureClient) GetByVMID(_ context.Context, _, _ string) (*azure.VirtualMachine, error) {
+func (m *mockAzureClient) GetByVMID(_ context.Context, _ string) (*azure.VirtualMachine, error) {
 	return nil, nil
 }