Azure updates

cmd/kola/options.go

@@ -123,6 +123,7 @@ func init() {
 	sv(&kola.AzureOptions.DiskController, "azure-disk-controller", "default", "Use a specific disk-controller for storage (default \"default\", also \"nvme\" and \"scsi\")")
 	sv(&kola.AzureOptions.ResourceGroup, "azure-resource-group", "", "Deploy resources in an existing resource group")
 	sv(&kola.AzureOptions.AvailabilitySet, "azure-availability-set", "", "Deploy instances with an existing availibity set")
+	sv(&kola.AzureOptions.KolaVnet, "azure-kola-vnet", "", "Pass the vnet/subnet that kola is being ran from to restrict network access to created storage accounts")
 
 	// do-specific options
 	sv(&kola.DOOptions.ConfigPath, "do-config-file", "", "DigitalOcean config file (default \"~/"+auth.DOConfigPath+"\")")

platform/api/azure/groups.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"time"
 
-	azruntime "github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources"
 )
@@ -58,11 +57,7 @@ func (a *API) TerminateResourceGroup(name string) error {
 	opts := &armresources.ResourceGroupsClientBeginDeleteOptions{
 		ForceDeletionTypes: to.Ptr("Microsoft.Compute/virtualMachines,Microsoft.Compute/virtualMachineScaleSets"),
 	}
-	poller, err := a.rgClient.BeginDelete(context.TODO(), name, opts)
-	pollOpts := &azruntime.PollUntilDoneOptions{
-		Frequency: 15 * time.Second,
-	}
-	_, err = poller.PollUntilDone(context.TODO(), pollOpts)
+	_, err := a.rgClient.BeginDelete(context.TODO(), name, opts)
 	return err
 }
 

platform/api/azure/instance.go

@@ -275,17 +275,12 @@ func (a *API) CreateInstance(name, userdata, sshkey, resourceGroup, storageAccou
 // OS disk are deleted automatically together with the VM.
 func (a *API) TerminateInstance(machine *Machine, resourceGroup string) error {
 	resourceGroup = a.getVMRG(resourceGroup)
-	poller, err := a.compClient.BeginDelete(context.TODO(), resourceGroup, machine.ID, &armcompute.VirtualMachinesClientBeginDeleteOptions{
+	_, err := a.compClient.BeginDelete(context.TODO(), resourceGroup, machine.ID, &armcompute.VirtualMachinesClientBeginDeleteOptions{
 		ForceDeletion: to.Ptr(true),
 	})
-	if err != nil {
-		return err
-	}
-	_, err = poller.PollUntilDone(context.TODO(), nil)
-	if err != nil {
-		return err
-	}
-	return nil
+	// We used to wait for the VM to be deleted here, but it's not necessary as
+	// we will also delete the resource group later.
+	return err
 }
 
 func (a *API) GetConsoleOutput(name, resourceGroup, storageAccount string) ([]byte, error) {

platform/api/azure/network.go

@@ -31,45 +31,49 @@ var (
 	kolaVnet                   = "kola-vn"
 )
 
-func (a *API) PrepareNetworkResources(resourceGroup string) (Network, error) {
-	if a.Opts.VnetSubnetName != "" {
-		parts := strings.SplitN(a.Opts.VnetSubnetName, "/", 2)
-		vnetName := parts[0]
-		subnetName := "default"
-		if len(parts) > 1 {
-			subnetName = parts[1]
+func (a *API) findVnetSubnet(vnetSubnetStr string) (Network, error) {
+	parts := strings.SplitN(vnetSubnetStr, "/", 2)
+	vnetName := parts[0]
+	subnetName := "default"
+	if len(parts) > 1 {
+		subnetName = parts[1]
+	}
+	var net *armnetwork.VirtualNetwork
+	pager := a.netClient.NewListAllPager(nil)
+	for pager.More() {
+		page, err := pager.NextPage(context.TODO())
+		if err != nil {
+			return Network{}, fmt.Errorf("failed to iterate vnets: %w", err)
 		}
-		var net *armnetwork.VirtualNetwork
-		pager := a.netClient.NewListAllPager(nil)
-		for pager.More() {
-			page, err := pager.NextPage(context.TODO())
-			if err != nil {
-				return Network{}, fmt.Errorf("failed to iterate vnets: %w", err)
-			}
-			for _, vnet := range page.Value {
-				if vnet.Name != nil && *vnet.Name == vnetName {
-					net = vnet
-					break
-				}
-			}
-			if net != nil {
+		for _, vnet := range page.Value {
+			if vnet.Name != nil && *vnet.Name == vnetName {
+				net = vnet
 				break
 			}
 		}
-		if net == nil {
-			return Network{}, fmt.Errorf("failed to find vnet %s", vnetName)
-		}
-		subnets := net.Properties.Subnets
-		if subnets == nil {
-			return Network{}, fmt.Errorf("failed to find subnet %s in vnet %s", subnetName, vnetName)
-		}
-		for _, subnet := range subnets {
-			if subnet != nil && subnet.Name != nil && *subnet.Name == subnetName {
-				return Network{*subnet}, nil
-			}
+		if net != nil {
+			break
 		}
+	}
+	if net == nil {
+		return Network{}, fmt.Errorf("failed to find vnet %s", vnetName)
+	}
+	subnets := net.Properties.Subnets
+	if subnets == nil {
 		return Network{}, fmt.Errorf("failed to find subnet %s in vnet %s", subnetName, vnetName)
 	}
+	for _, subnet := range subnets {
+		if subnet != nil && subnet.Name != nil && *subnet.Name == subnetName {
+			return Network{*subnet}, nil
+		}
+	}
+	return Network{}, fmt.Errorf("failed to find subnet %s in vnet %s", subnetName, vnetName)
+}
+
+func (a *API) PrepareNetworkResources(resourceGroup string) (Network, error) {
+	if a.Opts.VnetSubnetName != "" {
+		return a.findVnetSubnet(a.Opts.VnetSubnetName)
+	}
 
 	if err := a.createVirtualNetwork(resourceGroup); err != nil {
 		return Network{}, err

platform/api/azure/options.go

@@ -49,6 +49,7 @@ type Options struct {
 	Location         string
 	HyperVGeneration string
 	VnetSubnetName   string
+	KolaVnet         string
 	UseGallery       bool
 	UsePrivateIPs    bool
 

platform/api/azure/storage.go

@@ -260,6 +260,21 @@ func (a *API) CreateStorageAccount(resourceGroup string) (string, error) {
 			AllowSharedKeyAccess: to.Ptr(false),
 		},
 	}
+	if a.Opts.KolaVnet != "" {
+		net, err := a.findVnetSubnet(a.Opts.KolaVnet)
+		if err != nil {
+			return "", fmt.Errorf("CreateStorageAccount: %v", err)
+		}
+		parameters.Properties.NetworkRuleSet = &armstorage.NetworkRuleSet{
+			DefaultAction: to.Ptr(armstorage.DefaultActionDeny),
+			VirtualNetworkRules: []*armstorage.VirtualNetworkRule{
+				{
+					VirtualNetworkResourceID: net.subnet.ID,
+				},
+			},
+		}
+	}
+
 	plog.Infof("Creating StorageAccount %s", name)
 	poller, err := a.accClient.BeginCreate(ctx, resourceGroup, name, parameters, nil)
 	if err != nil {