feat: Add airgap nsg to abe2e (#4136)

.pipelines/e2e.yaml

@@ -19,12 +19,15 @@ pr:
 
 variables:
 - group: ab-e2e
+- name: defaultTimeout
+  value: 90
 
 pool:
   name: nodesigtest-pool
 
 jobs:
 - job: Run_AgentBaker_E2E
+  timeoutInMinutes: ${{ variables.defaultTimeout }}
   steps:
   - bash: |
       az login --identity
@@ -39,7 +42,7 @@ jobs:
       export BUILD_ID=$(Build.BuildId)
 
       cd e2e
-      go test -timeout 90m -v -run Test_All ./
+      go test -timeout ${{ variables.defaultTimeout }}m -v -run Test_All ./
     displayName: Run AgentBaker E2E
     env:
       VHD_BUILD_ID: $(VHD_BUILD_ID)

e2e/airgap.go

@@ -0,0 +1,151 @@
+package e2e_test
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"net"
+
+	"github.com/Azure/agentbakere2e/suite"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
+	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/network/armnetwork"
+)
+
+var (
+	nsgName           string = "abe2e-airgap-securityGroup"
+	defaultSubnetName string = "aks-subnet"
+)
+
+func airGapSecurityGroup(location, kubernetesEndpont string) armnetwork.SecurityGroup {
+	allowVnet := armnetwork.SecurityRule{
+		Name: to.Ptr("AllowVnetOutBound"),
+		Properties: &armnetwork.SecurityRulePropertiesFormat{
+			Protocol:                 to.Ptr(armnetwork.SecurityRuleProtocolAsterisk),
+			Access:                   to.Ptr(armnetwork.SecurityRuleAccessAllow),
+			Direction:                to.Ptr(armnetwork.SecurityRuleDirectionOutbound),
+			SourceAddressPrefix:      to.Ptr("VirtualNetwork"),
+			SourcePortRange:          to.Ptr("*"),
+			DestinationAddressPrefix: to.Ptr("VirtualNetwork"),
+			DestinationPortRange:     to.Ptr("*"),
+			Priority:                 to.Ptr[int32](2000),
+		},
+	}
+
+	blockOutbound := armnetwork.SecurityRule{
+		Name: to.Ptr("block-all-outbound"),
+		Properties: &armnetwork.SecurityRulePropertiesFormat{
+			Protocol:                 to.Ptr(armnetwork.SecurityRuleProtocolAsterisk),
+			Access:                   to.Ptr(armnetwork.SecurityRuleAccessDeny),
+			Direction:                to.Ptr(armnetwork.SecurityRuleDirectionOutbound),
+			SourceAddressPrefix:      to.Ptr("*"),
+			SourcePortRange:          to.Ptr("*"),
+			DestinationAddressPrefix: to.Ptr("*"),
+			DestinationPortRange:     to.Ptr("*"),
+			Priority:                 to.Ptr[int32](2001),
+		},
+	}
+
+	securityRules := []*armnetwork.SecurityRule{
+		getSecurityRule("allow-mcr-microsoft-com", "204.79.197.219", 100),
+		getSecurityRule("allow-acs-mirror.azureedge.net", "72.21.81.200", 101),
+		getSecurityRule("allow-management.azure.com", "4.150.240.10", 102),
+		getSecurityRule("allow-kubernetes-endpoint", kubernetesEndpont, 103),
+		&allowVnet,
+		&blockOutbound,
+	}
+
+	return armnetwork.SecurityGroup{
+		Location:   &location,
+		Name:       &nsgName,
+		Properties: &armnetwork.SecurityGroupPropertiesFormat{SecurityRules: securityRules},
+	}
+}
+
+func getSecurityRule(name, destinationAddressPrefix string, priority int32) *armnetwork.SecurityRule {
+	return &armnetwork.SecurityRule{
+		Name: to.Ptr(name),
+		Properties: &armnetwork.SecurityRulePropertiesFormat{
+			Protocol:                 to.Ptr(armnetwork.SecurityRuleProtocolAsterisk),
+			Access:                   to.Ptr(armnetwork.SecurityRuleAccessAllow),
+			Direction:                to.Ptr(armnetwork.SecurityRuleDirectionOutbound),
+			SourceAddressPrefix:      to.Ptr("*"),
+			SourcePortRange:          to.Ptr("*"),
+			DestinationAddressPrefix: to.Ptr(destinationAddressPrefix),
+			DestinationPortRange:     to.Ptr("*"),
+			Priority:                 to.Ptr[int32](priority),
+		},
+	}
+}
+
+func addAirgapNetworkSettings(ctx context.Context, cloud *azureClient, suiteConfig *suite.Config, clusterConfig clusterConfig) error {
+	log.Printf("Adding network settings for airgap cluster %s in rg %s\n", *clusterConfig.cluster.Name, *clusterConfig.cluster.Properties.NodeResourceGroup)
+
+	vnet, err := getClusterVNet(ctx, cloud, *clusterConfig.cluster.Properties.NodeResourceGroup)
+	if err != nil {
+		return err
+	}
+	clusterConfig.subnetId = vnet.subnetId
+
+	ipAddresses, err := net.LookupIP(*clusterConfig.cluster.Properties.Fqdn)
+	if err != nil {
+		return err
+	}
+	nsgParams := airGapSecurityGroup(suiteConfig.Location, ipAddresses[0].String())
+
+	nsg, err := createAirgapSecurityGroup(ctx, cloud, clusterConfig, nsgParams, nil)
+	if err != nil {
+		return err
+	}
+
+	subnetParameters := armnetwork.Subnet{
+		ID: to.Ptr(clusterConfig.subnetId),
+		Properties: &armnetwork.SubnetPropertiesFormat{
+			AddressPrefix: to.Ptr("10.224.0.0/16"),
+			NetworkSecurityGroup: &armnetwork.SecurityGroup{
+				ID: nsg.ID,
+			},
+		},
+	}
+	err = updateSubnet(ctx, cloud, clusterConfig, subnetParameters, vnet.name)
+	if err != nil {
+		return err
+	}
+
+	fmt.Printf("Updated the subnet to airgap %s\n", *clusterConfig.cluster.Name)
+	return nil
+}
+
+func createAirgapSecurityGroup(ctx context.Context, cloud *azureClient, clusterConfig clusterConfig, nsgParams armnetwork.SecurityGroup, options *armnetwork.SecurityGroupsClientBeginCreateOrUpdateOptions) (*armnetwork.SecurityGroupsClientCreateOrUpdateResponse, error) {
+	poller, err := cloud.securityGroupClient.BeginCreateOrUpdate(ctx, *clusterConfig.cluster.Properties.NodeResourceGroup, nsgName, nsgParams, options)
+	if err != nil {
+		return nil, err
+	}
+	nsg, err := poller.PollUntilDone(ctx, nil)
+	if err != nil {
+		return nil, err
+	}
+	return &nsg, nil
+}
+
+func updateSubnet(ctx context.Context, cloud *azureClient, clusterConfig clusterConfig, subnetParameters armnetwork.Subnet, vnetName string) error {
+	poller, err := cloud.subnetClient.BeginCreateOrUpdate(ctx, *clusterConfig.cluster.Properties.NodeResourceGroup, vnetName, defaultSubnetName, subnetParameters, nil)
+	if err != nil {
+		return err
+	}
+	_, err = poller.PollUntilDone(ctx, nil)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func isNetworkSecurityGroupAirgap(cloud *azureClient, resourceGroupName string) (bool, error) {
+	_, err := cloud.securityGroupClient.Get(context.Background(), resourceGroupName, nsgName, nil)
+	if err != nil {
+		if isNotFoundError(err) {
+			return false, nil
+		}
+		return false, fmt.Errorf("failed to get network security group: %w", err)
+	}
+	return true, nil
+}

e2e/cloud.go

@@ -27,6 +27,8 @@ type azureClient struct {
 	resourceClient      *armresources.Client
 	resourceGroupClient *armresources.ResourceGroupsClient
 	aksClient           *armcontainerservice.ManagedClustersClient
+	securityGroupClient *armnetwork.SecurityGroupsClient
+	subnetClient        *armnetwork.SubnetsClient
 }
 
 func newAzureClient(subscription string) (*azureClient, error) {
@@ -89,6 +91,16 @@ func newAzureClient(subscription string) (*azureClient, error) {
 		return nil, fmt.Errorf("failed to create core client: %w", err)
 	}
 
+	securityGroupClient, err := armnetwork.NewSecurityGroupsClient(subscription, credential, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create security group client: %w", err)
+	}
+
+	subnetClient, err := armnetwork.NewSubnetsClient(subscription, credential, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create subnet client: %w", err)
+	}
+
 	aksClient, err := armcontainerservice.NewManagedClustersClient(subscription, credential, opts)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create aks client: %w", err)
@@ -127,6 +139,8 @@ func newAzureClient(subscription string) (*azureClient, error) {
 		vmssClient:          vmssClient,
 		vmssVMClient:        vmssVMClient,
 		vnetClient:          vnetClient,
+		securityGroupClient: securityGroupClient,
+		subnetClient:        subnetClient,
 	}
 
 	return cloud, nil

e2e/cluster.go

@@ -24,11 +24,17 @@ const (
 type clusterParameters map[string]string
 
 type clusterConfig struct {
-	cluster      *armcontainerservice.ManagedCluster
-	kube         *kubeclient
-	parameters   clusterParameters
-	subnetId     string
-	isNewCluster bool
+	cluster         *armcontainerservice.ManagedCluster
+	kube            *kubeclient
+	parameters      clusterParameters
+	subnetId        string
+	isNewCluster    bool
+	isAirgapCluster bool
+}
+
+type VNet struct {
+	name     string
+	subnetId string
 }
 
 // Returns true if the cluster is configured with Azure CNI
@@ -145,7 +151,8 @@ func createNewCluster(
 	ctx context.Context,
 	cloud *azureClient,
 	resourceGroupName string,
-	clusterModel *armcontainerservice.ManagedCluster) (*armcontainerservice.ManagedCluster, error) {
+	clusterConfig clusterConfig) (*armcontainerservice.ManagedCluster, error) {
+	clusterModel := clusterConfig.cluster
 	pollerResp, err := cloud.aksClient.BeginCreateOrUpdate(
 		ctx,
 		resourceGroupName,
@@ -179,23 +186,21 @@ func deleteExistingCluster(ctx context.Context, cloud *azureClient, resourceGrou
 	return nil
 }
 
-func getClusterSubnetID(ctx context.Context, cloud *azureClient, location, mcResourceGroupName, clusterName string) (string, error) {
+func getClusterVNet(ctx context.Context, cloud *azureClient, mcResourceGroupName string) (VNet, error) {
 	pager := cloud.vnetClient.NewListPager(mcResourceGroupName, nil)
-
 	for pager.More() {
 		nextResult, err := pager.NextPage(ctx)
 		if err != nil {
-			return "", fmt.Errorf("failed to advance page: %w", err)
+			return VNet{}, fmt.Errorf("failed to advance page: %w", err)
 		}
 		for _, v := range nextResult.Value {
 			if v == nil {
-				return "", fmt.Errorf("aks vnet id was empty")
+				return VNet{}, fmt.Errorf("aks vnet was empty")
 			}
-			return fmt.Sprintf("%s/subnets/%s", *v.ID, "aks-subnet"), nil
+			return VNet{name: *v.Name, subnetId: fmt.Sprintf("%s/subnets/%s", *v.ID, "aks-subnet")}, nil
 		}
 	}
-
-	return "", fmt.Errorf("failed to find aks vnet")
+	return VNet{}, fmt.Errorf("failed to find aks vnet")
 }
 
 func getInitialClusterConfigs(ctx context.Context, cloud *azureClient, resourceGroupName string) ([]clusterConfig, error) {
@@ -226,36 +231,40 @@ func getInitialClusterConfigs(ctx context.Context, cloud *azureClient, resourceG
 					continue
 				}
 
-				log.Printf("found agentbaker e2e cluster %q in provisioning state %q", *resource.Name, *cluster.Properties.ProvisioningState)
-				configs = append(configs, clusterConfig{cluster: &cluster.ManagedCluster})
+				clusterConfig := clusterConfig{cluster: &cluster.ManagedCluster}
+				isAirgap, err := isNetworkSecurityGroupAirgap(cloud, *clusterConfig.cluster.Properties.NodeResourceGroup)
+				if err != nil {
+					return nil, fmt.Errorf("failed to verify if aks subnet is for an airgap cluster: %w", err)
+				}
+
+				clusterConfig.isAirgapCluster = isAirgap
+				log.Printf("found agentbaker e2e cluster %q in provisioning state %q is Airgap %v", *resource.Name, *cluster.Properties.ProvisioningState, clusterConfig.isAirgapCluster)
+				configs = append(configs, clusterConfig)
 			}
 		}
 	}
-
 	return configs, nil
 }
 
 func hasViableConfig(scenario *scenario.Scenario, clusterConfigs []clusterConfig) bool {
 	for _, config := range clusterConfigs {
+		if scenario.Airgap && !config.isAirgapCluster {
+			continue
+		}
 		if scenario.Config.ClusterSelector(config.cluster) {
 			return true
 		}
 	}
 	return false
 }
 
-func createMissingClusters(
-	ctx context.Context,
-	r *mrand.Rand,
-	cloud *azureClient,
-	suiteConfig *suite.Config,
-	scenarios scenario.Table,
-	clusterConfigs *[]clusterConfig) error {
+func createMissingClusters(ctx context.Context, r *mrand.Rand, cloud *azureClient, suiteConfig *suite.Config,
+	scenarios scenario.Table, clusterConfigs *[]clusterConfig) error {
 	var newConfigs []clusterConfig
 	for _, scenario := range scenarios {
 		if !hasViableConfig(scenario, *clusterConfigs) && !hasViableConfig(scenario, newConfigs) {
 			newClusterModel := getNewClusterModelForScenario(generateClusterName(r), suiteConfig.Location, scenario)
-			newConfigs = append(newConfigs, clusterConfig{cluster: &newClusterModel, isNewCluster: true})
+			newConfigs = append(newConfigs, clusterConfig{cluster: &newClusterModel, isNewCluster: true, isAirgapCluster: scenario.Airgap})
 		}
 	}
 
@@ -267,7 +276,7 @@ func createMissingClusters(
 			clusterName := *config.cluster.Name
 
 			log.Printf("creating cluster %q...", clusterName)
-			liveCluster, err := createNewCluster(ctx, cloud, suiteConfig.ResourceGroupName, config.cluster)
+			liveCluster, err := createNewCluster(ctx, cloud, suiteConfig.ResourceGroupName, config)
 			if err != nil {
 				return fmt.Errorf("unable to create new cluster: %w", err)
 			}
@@ -307,9 +316,14 @@ func chooseCluster(
 	suiteConfig *suite.Config,
 	scenario *scenario.Scenario,
 	clusterConfigs []clusterConfig) (clusterConfig, error) {
+
 	var chosenConfig clusterConfig
 	for i := range clusterConfigs {
 		config := &clusterConfigs[i]
+		if (scenario.Airgap && !config.isAirgapCluster) || (config.isAirgapCluster && !scenario.Airgap) {
+			continue
+		} 
+
 		if scenario.Config.ClusterSelector(config.cluster) {
 			// only validate + prep the cluster for testing if we didn't just create it and it hasn't already been prepared
 			if !config.isNewCluster && config.needsPreparation() {
@@ -331,6 +345,21 @@ func chooseCluster(
 		return clusterConfig{}, fmt.Errorf("tried to chose a cluster without a node resource group: %+v", *chosenConfig.cluster)
 	}
 
+	if chosenConfig.isAirgapCluster {
+		hasAirgapNSG, err := isNetworkSecurityGroupAirgap(cloud, *chosenConfig.cluster.Properties.NodeResourceGroup)
+		if err != nil {
+			return clusterConfig{}, fmt.Errorf("failed to check if airgap settings are present: %w", err)
+		}
+
+		if !hasAirgapNSG {
+			log.Printf("adding airgap network settings to cluster %q...", *chosenConfig.cluster.Name)
+			err = addAirgapNetworkSettings(ctx, cloud, suiteConfig, chosenConfig)
+			if err != nil {
+				return clusterConfig{}, fmt.Errorf("failed to add airgap network settings: %w", err)
+			}
+		}
+	}
+
 	return chosenConfig, nil
 }
 
@@ -346,7 +375,7 @@ func validateAndPrepareCluster(ctx context.Context, r *mrand.Rand, cloud *azureC
 		if err != nil {
 			return err
 		}
-		newCluster, err := createNewCluster(ctx, cloud, suiteConfig.ResourceGroupName, newModel)
+		newCluster, err := createNewCluster(ctx, cloud, suiteConfig.ResourceGroupName, *config)
 		if err != nil {
 			return err
 		}
@@ -372,7 +401,7 @@ func prepareClusterForTests(
 	cluster *armcontainerservice.ManagedCluster) (*kubeclient, string, clusterParameters, error) {
 	clusterName := *cluster.Name
 
-	subnetId, err := getClusterSubnetID(ctx, cloud, suiteConfig.Location, *cluster.Properties.NodeResourceGroup, clusterName)
+	vnet, err := getClusterVNet(ctx, cloud, *cluster.Properties.NodeResourceGroup)
 	if err != nil {
 		return nil, "", nil, fmt.Errorf("unable get subnet ID of cluster %q: %w", clusterName, err)
 	}
@@ -390,8 +419,7 @@ func prepareClusterForTests(
 	if err != nil {
 		return nil, "", nil, fmt.Errorf("unable to extract cluster parameters from %q: %w", clusterName, err)
 	}
-
-	return kube, subnetId, clusterParams, nil
+	return kube, vnet.subnetId, clusterParams, nil
 }
 
 // TODO(cameissner): figure out a better way to reconcile server-side and client-side properties,