chore(e2e): add firewall to e2es (#7350)

lilypan26 · web-flow · commit 9d81ad7ee5d5 · 2025-11-21T13:31:11.000-08:00
diff --git a/e2e/aks_model.go b/e2e/aks_model.go
@@ -182,6 +182,258 @@ func getBaseClusterModel(clusterName, location, k8sSystemPoolSKU string) *armcon
 	}
 }
 
+func getFirewall(ctx context.Context, location, firewallSubnetID, publicIPID string) *armnetwork.AzureFirewall {
+	var (
+		natRuleCollections []*armnetwork.AzureFirewallNatRuleCollection
+		netRuleCollections []*armnetwork.AzureFirewallNetworkRuleCollection
+	)
+
+	// Application rule for AKS FQDN tags
+	aksAppRule := armnetwork.AzureFirewallApplicationRule{
+		Name:            to.Ptr("aks-fqdn"),
+		SourceAddresses: []*string{to.Ptr("*")},
+		Protocols: []*armnetwork.AzureFirewallApplicationRuleProtocol{
+			{
+				ProtocolType: to.Ptr(armnetwork.AzureFirewallApplicationRuleProtocolTypeHTTP),
+				Port:         to.Ptr[int32](80),
+			},
+			{
+				ProtocolType: to.Ptr(armnetwork.AzureFirewallApplicationRuleProtocolTypeHTTPS),
+				Port:         to.Ptr[int32](443),
+			},
+		},
+		FqdnTags: []*string{to.Ptr("AzureKubernetesService")},
+	}
+
+	// needed for scriptless e2e hack
+	blobStorageFqdn := config.Config.BlobStorageAccount() + ".blob.core.windows.net"
+	blobStorageAppRule := armnetwork.AzureFirewallApplicationRule{
+		Name:            to.Ptr("blob-storage-fqdn"),
+		SourceAddresses: []*string{to.Ptr("*")},
+		Protocols: []*armnetwork.AzureFirewallApplicationRuleProtocol{
+			{
+				ProtocolType: to.Ptr(armnetwork.AzureFirewallApplicationRuleProtocolTypeHTTPS),
+				Port:         to.Ptr[int32](443),
+			},
+		},
+		TargetFqdns: []*string{to.Ptr(blobStorageFqdn)},
+	}
+
+	appRuleCollection := armnetwork.AzureFirewallApplicationRuleCollection{
+		Name: to.Ptr("aksfwar"),
+		Properties: &armnetwork.AzureFirewallApplicationRuleCollectionPropertiesFormat{
+			Priority: to.Ptr[int32](100),
+			Action: &armnetwork.AzureFirewallRCAction{
+				Type: to.Ptr(armnetwork.AzureFirewallRCActionTypeAllow),
+			},
+			Rules: []*armnetwork.AzureFirewallApplicationRule{&aksAppRule, &blobStorageAppRule},
+		},
+	}
+
+	ipConfigurations := []*armnetwork.AzureFirewallIPConfiguration{
+		{
+			Name: to.Ptr("firewall-ip-config"),
+			Properties: &armnetwork.AzureFirewallIPConfigurationPropertiesFormat{
+				Subnet: &armnetwork.SubResource{
+					ID: to.Ptr(firewallSubnetID),
+				},
+				PublicIPAddress: &armnetwork.SubResource{
+					ID: to.Ptr(publicIPID),
+				},
+			},
+		},
+	}
+
+	logf(ctx, "Firewall rules configured successfully")
+	return &armnetwork.AzureFirewall{
+		Location: to.Ptr(location),
+		Properties: &armnetwork.AzureFirewallPropertiesFormat{
+			ApplicationRuleCollections: []*armnetwork.AzureFirewallApplicationRuleCollection{&appRuleCollection},
+			NetworkRuleCollections:     netRuleCollections,
+			NatRuleCollections:         natRuleCollections,
+			IPConfigurations:           ipConfigurations,
+		},
+	}
+}
+
+func addFirewallRules(
+	ctx context.Context, clusterModel *armcontainerservice.ManagedCluster,
+	location string,
+) error {
+
+	vnet, err := getClusterVNet(ctx, *clusterModel.Properties.NodeResourceGroup)
+	if err != nil {
+		return err
+	}
+
+	// Create AzureFirewallSubnet - this subnet name is required by Azure Firewall
+	firewallSubnetName := "AzureFirewallSubnet"
+	firewallSubnetParams := armnetwork.Subnet{
+		Properties: &armnetwork.SubnetPropertiesFormat{
+			AddressPrefix: to.Ptr("10.225.0.0/24"), // Use a different CIDR that doesn't overlap with 10.224.0.0/16
+		},
+	}
+
+	logf(ctx, "Creating subnet %s in VNet %s", firewallSubnetName, vnet.name)
+	subnetPoller, err := config.Azure.Subnet.BeginCreateOrUpdate(
+		ctx,
+		*clusterModel.Properties.NodeResourceGroup,
+		vnet.name,
+		firewallSubnetName,
+		firewallSubnetParams,
+		nil,
+	)
+	if err != nil {
+		return fmt.Errorf("failed to start creating firewall subnet: %w", err)
+	}
+
+	subnetResp, err := subnetPoller.PollUntilDone(ctx, config.DefaultPollUntilDoneOptions)
+	if err != nil {
+		return fmt.Errorf("failed to create firewall subnet: %w", err)
+	}
+
+	firewallSubnetID := *subnetResp.ID
+	logf(ctx, "Created firewall subnet with ID: %s", firewallSubnetID)
+
+	// Create public IP for the firewall
+	publicIPName := "abe2e-fw-pip"
+	publicIPParams := armnetwork.PublicIPAddress{
+		Location: to.Ptr(location),
+		SKU: &armnetwork.PublicIPAddressSKU{
+			Name: to.Ptr(armnetwork.PublicIPAddressSKUNameStandard),
+		},
+		Properties: &armnetwork.PublicIPAddressPropertiesFormat{
+			PublicIPAllocationMethod: to.Ptr(armnetwork.IPAllocationMethodStatic),
+		},
+	}
+
+	logf(ctx, "Creating public IP %s", publicIPName)
+	pipPoller, err := config.Azure.PublicIPAddresses.BeginCreateOrUpdate(
+		ctx,
+		*clusterModel.Properties.NodeResourceGroup,
+		publicIPName,
+		publicIPParams,
+		nil,
+	)
+	if err != nil {
+		return fmt.Errorf("failed to start creating public IP: %w", err)
+	}
+
+	pipResp, err := pipPoller.PollUntilDone(ctx, config.DefaultPollUntilDoneOptions)
+	if err != nil {
+		return fmt.Errorf("failed to create public IP: %w", err)
+	}
+
+	publicIPID := *pipResp.ID
+	logf(ctx, "Created public IP with ID: %s", publicIPID)
+
+	firewallName := "abe2e-fw"
+	firewall := getFirewall(ctx, location, firewallSubnetID, publicIPID)
+	fwPoller, err := config.Azure.AzureFirewall.BeginCreateOrUpdate(ctx, *clusterModel.Properties.NodeResourceGroup, firewallName, *firewall, nil)
+	if err != nil {
+		return fmt.Errorf("failed to start Firewall creation: %w", err)
+	}
+	fwResp, err := fwPoller.PollUntilDone(ctx, nil)
+	if err != nil {
+		return fmt.Errorf("failed to create Firewall: %w", err)
+	}
+
+	// Get the firewall's private IP address
+	var firewallPrivateIP string
+	if fwResp.Properties != nil && fwResp.Properties.IPConfigurations != nil && len(fwResp.Properties.IPConfigurations) > 0 {
+		if fwResp.Properties.IPConfigurations[0].Properties != nil && fwResp.Properties.IPConfigurations[0].Properties.PrivateIPAddress != nil {
+			firewallPrivateIP = *fwResp.Properties.IPConfigurations[0].Properties.PrivateIPAddress
+			logf(ctx, "Firewall private IP: %s", firewallPrivateIP)
+		}
+	}
+
+	if firewallPrivateIP == "" {
+		return fmt.Errorf("failed to get firewall private IP address")
+	}
+
+	routeTableName := "abe2e-fw-rt"
+	routeTableParams := armnetwork.RouteTable{
+		Location: to.Ptr(location),
+		Properties: &armnetwork.RouteTablePropertiesFormat{
+			Routes: []*armnetwork.Route{
+				// Allow internal VNet traffic to bypass the firewall
+				{
+					Name: to.Ptr("vnet-local"),
+					Properties: &armnetwork.RoutePropertiesFormat{
+						AddressPrefix: to.Ptr("10.224.0.0/16"), // AKS subnet CIDR
+						NextHopType:   to.Ptr(armnetwork.RouteNextHopTypeVnetLocal),
+					},
+				},
+				// Route all other traffic (internet-bound) through the firewall
+				{
+					Name: to.Ptr("default-route-to-firewall"),
+					Properties: &armnetwork.RoutePropertiesFormat{
+						AddressPrefix:    to.Ptr("0.0.0.0/0"),
+						NextHopType:      to.Ptr(armnetwork.RouteNextHopTypeVirtualAppliance),
+						NextHopIPAddress: to.Ptr(firewallPrivateIP),
+					},
+				},
+			},
+			DisableBgpRoutePropagation: to.Ptr(true),
+		},
+	}
+
+	logf(ctx, "Creating route table %s", routeTableName)
+	rtPoller, err := config.Azure.RouteTables.BeginCreateOrUpdate(
+		ctx,
+		*clusterModel.Properties.NodeResourceGroup,
+		routeTableName,
+		routeTableParams,
+		nil,
+	)
+	if err != nil {
+		return fmt.Errorf("failed to start creating route table: %w", err)
+	}
+
+	rtResp, err := rtPoller.PollUntilDone(ctx, config.DefaultPollUntilDoneOptions)
+	if err != nil {
+		return fmt.Errorf("failed to create route table: %w", err)
+	}
+
+	logf(ctx, "Created route table with ID: %s", *rtResp.ID)
+
+	// Get the AKS subnet and associate it with the route table
+	aksSubnetResp, err := config.Azure.Subnet.Get(ctx, *clusterModel.Properties.NodeResourceGroup, vnet.name, "aks-subnet", nil)
+	if err != nil {
+		return fmt.Errorf("failed to get AKS subnet: %w", err)
+	}
+
+	// Update subnet to associate with route table
+	aksSubnet := aksSubnetResp.Subnet
+	if aksSubnet.Properties == nil {
+		aksSubnet.Properties = &armnetwork.SubnetPropertiesFormat{}
+	}
+	aksSubnet.Properties.RouteTable = &armnetwork.RouteTable{
+		ID: rtResp.ID,
+	}
+
+	logf(ctx, "Associating route table with AKS subnet")
+	subnetUpdatePoller, err := config.Azure.Subnet.BeginCreateOrUpdate(
+		ctx,
+		*clusterModel.Properties.NodeResourceGroup,
+		vnet.name,
+		"aks-subnet",
+		aksSubnet,
+		nil,
+	)
+	if err != nil {
+		return fmt.Errorf("failed to start updating subnet with route table: %w", err)
+	}
+
+	_, err = subnetUpdatePoller.PollUntilDone(ctx, config.DefaultPollUntilDoneOptions)
+	if err != nil {
+		return fmt.Errorf("failed to associate route table with subnet: %w", err)
+	}
+
+	logf(ctx, "Successfully configured firewall and routing for AKS cluster")
+	return nil
+}
+
 func addAirgapNetworkSettings(ctx context.Context, clusterModel *armcontainerservice.ManagedCluster, privateACRName, location string) error {
 	logf(ctx, "Adding network settings for airgap cluster %s in rg %s", *clusterModel.Name, *clusterModel.Properties.NodeResourceGroup)
 
diff --git a/e2e/cache.go b/e2e/cache.go
@@ -157,49 +157,49 @@ var ClusterKubenet = cachedFunc(clusterKubenet)
 
 // clusterKubenet creates a basic cluster using kubenet networking
 func clusterKubenet(ctx context.Context, request ClusterRequest) (*Cluster, error) {
-	return prepareCluster(ctx, getKubenetClusterModel("abe2e-kubenet-v2", request.Location, request.K8sSystemPoolSKU), false, false)
+	return prepareCluster(ctx, getKubenetClusterModel("abe2e-kubenet-v3", request.Location, request.K8sSystemPoolSKU), false, false)
 }
 
 var ClusterKubenetAirgap = cachedFunc(clusterKubenetAirgap)
 
 // clusterKubenetAirgap creates an airgapped kubenet cluster (no internet access)
 func clusterKubenetAirgap(ctx context.Context, request ClusterRequest) (*Cluster, error) {
-	return prepareCluster(ctx, getKubenetClusterModel("abe2e-kubenet-airgap", request.Location, request.K8sSystemPoolSKU), true, false)
+	return prepareCluster(ctx, getKubenetClusterModel("abe2e-kubenet-airgap-v2", request.Location, request.K8sSystemPoolSKU), true, false)
 }
 
 var ClusterKubenetAirgapNonAnon = cachedFunc(clusterKubenetAirgapNonAnon)
 
 // clusterKubenetAirgapNonAnon creates an airgapped kubenet cluster with non-anonymous image pulls
 func clusterKubenetAirgapNonAnon(ctx context.Context, request ClusterRequest) (*Cluster, error) {
-	return prepareCluster(ctx, getKubenetClusterModel("abe2e-kubenet-nonanonpull-airgap", request.Location, request.K8sSystemPoolSKU), true, true)
+	return prepareCluster(ctx, getKubenetClusterModel("abe2e-kubenet-nonanonpull-airgap-v2", request.Location, request.K8sSystemPoolSKU), true, true)
 }
 
 var ClusterAzureNetwork = cachedFunc(clusterAzureNetwork)
 
 // clusterAzureNetwork creates a cluster with Azure CNI networking
 func clusterAzureNetwork(ctx context.Context, request ClusterRequest) (*Cluster, error) {
-	return prepareCluster(ctx, getAzureNetworkClusterModel("abe2e-azure-network", request.Location, request.K8sSystemPoolSKU), false, false)
+	return prepareCluster(ctx, getAzureNetworkClusterModel("abe2e-azure-network-v2", request.Location, request.K8sSystemPoolSKU), false, false)
 }
 
 var ClusterAzureOverlayNetwork = cachedFunc(clusterAzureOverlayNetwork)
 
 // clusterAzureOverlayNetwork creates a cluster with Azure CNI Overlay networking
 func clusterAzureOverlayNetwork(ctx context.Context, request ClusterRequest) (*Cluster, error) {
-	return prepareCluster(ctx, getAzureOverlayNetworkClusterModel("abe2e-azure-overlay-network", request.Location, request.K8sSystemPoolSKU), false, false)
+	return prepareCluster(ctx, getAzureOverlayNetworkClusterModel("abe2e-azure-overlay-network-v2", request.Location, request.K8sSystemPoolSKU), false, false)
 }
 
 var ClusterAzureOverlayNetworkDualStack = cachedFunc(clusterAzureOverlayNetworkDualStack)
 
 // clusterAzureOverlayNetworkDualStack creates a dual-stack (IPv4+IPv6) Azure CNI Overlay cluster
 func clusterAzureOverlayNetworkDualStack(ctx context.Context, request ClusterRequest) (*Cluster, error) {
-	return prepareCluster(ctx, getAzureOverlayNetworkDualStackClusterModel("abe2e-azure-overlay-dualstack", request.Location, request.K8sSystemPoolSKU), false, false)
+	return prepareCluster(ctx, getAzureOverlayNetworkDualStackClusterModel("abe2e-azure-overlay-dualstack-v2", request.Location, request.K8sSystemPoolSKU), false, false)
 }
 
 var ClusterCiliumNetwork = cachedFunc(clusterCiliumNetwork)
 
 // clusterCiliumNetwork creates a cluster with Cilium CNI networking
 func clusterCiliumNetwork(ctx context.Context, request ClusterRequest) (*Cluster, error) {
-	return prepareCluster(ctx, getCiliumNetworkClusterModel("abe2e-cilium-network", request.Location, request.K8sSystemPoolSKU), false, false)
+	return prepareCluster(ctx, getCiliumNetworkClusterModel("abe2e-cilium-network-v2", request.Location, request.K8sSystemPoolSKU), false, false)
 }
 
 // isNotFoundErr checks if an error represents a "not found" response from Azure API
diff --git a/e2e/cluster.go b/e2e/cluster.go
@@ -100,6 +100,10 @@ func prepareCluster(ctx context.Context, cluster *armcontainerservice.ManagedClu
 		}
 	}
 
+	if err := addFirewallRules(ctx, cluster, *cluster.Location); err != nil {
+		return nil, fmt.Errorf("add firewall rules: %w", err)
+	}
+
 	kubeletIdentity, err := getClusterKubeletIdentity(cluster)
 	if err != nil {
 		return nil, fmt.Errorf("getting cluster kubelet identity: %w", err)
diff --git a/e2e/config/azure.go b/e2e/config/azure.go
@@ -39,6 +39,7 @@ import (
 
 type AzureClient struct {
 	AKS                       *armcontainerservice.ManagedClustersClient
+	AzureFirewall             *armnetwork.AzureFirewallsClient
 	Blob                      *azblob.Client
 	StorageContainers         *armstorage.BlobContainersClient
 	CacheRulesClient          *armcontainerregistry.CacheRulesClient
@@ -57,6 +58,8 @@ type AzureClient struct {
 	SecurityGroup             *armnetwork.SecurityGroupsClient
 	StorageAccounts           *armstorage.AccountsClient
 	Subnet                    *armnetwork.SubnetsClient
+	PublicIPAddresses         *armnetwork.PublicIPAddressesClient
+	RouteTables               *armnetwork.RouteTablesClient
 	UserAssignedIdentities    *armmsi.UserAssignedIdentitiesClient
 	VMSS                      *armcompute.VirtualMachineScaleSetsClient
 	VMSSVM                    *armcompute.VirtualMachineScaleSetVMsClient
@@ -148,6 +151,11 @@ func NewAzureClient() (*AzureClient, error) {
 		return nil, fmt.Errorf("create core client: %w", err)
 	}
 
+	cloud.PublicIPAddresses, err = armnetwork.NewPublicIPAddressesClient(Config.SubscriptionID, credential, opts)
+	if err != nil {
+		return nil, fmt.Errorf("create public ip addresses client: %w", err)
+	}
+
 	cloud.RegistriesClient, err = armcontainerregistry.NewRegistriesClient(Config.SubscriptionID, credential, opts)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create registry client: %w", err)
@@ -193,6 +201,11 @@ func NewAzureClient() (*AzureClient, error) {
 		return nil, fmt.Errorf("create subnet client: %w", err)
 	}
 
+	cloud.RouteTables, err = armnetwork.NewRouteTablesClient(Config.SubscriptionID, credential, opts)
+	if err != nil {
+		return nil, fmt.Errorf("create route tables client: %w", err)
+	}
+
 	cloud.AKS, err = armcontainerservice.NewManagedClustersClient(Config.SubscriptionID, credential, opts)
 	if err != nil {
 		return nil, fmt.Errorf("create aks client: %w", err)
@@ -258,6 +271,16 @@ func NewAzureClient() (*AzureClient, error) {
 		return nil, fmt.Errorf("create vnet client: %w", err)
 	}
 
+	cloud.AzureFirewall, err = armnetwork.NewAzureFirewallsClient(Config.SubscriptionID, credential, opts)
+	if err != nil {
+		return nil, fmt.Errorf("create firewall client: %w", err)
+	}
+
+	cloud.PublicIPAddresses, err = armnetwork.NewPublicIPAddressesClient(Config.SubscriptionID, credential, opts)
+	if err != nil {
+		return nil, fmt.Errorf("create public ip addresses client: %w", err)
+	}
+
 	cloud.Blob, err = azblob.NewClient(Config.BlobStorageAccountURL(), credential, nil)
 	if err != nil {
 		return nil, fmt.Errorf("create blob container client: %w", err)

Original file line number	Diff line number	Diff line change
`@@ -100,6 +100,10 @@ func prepareCluster(ctx context.Context, cluster *armcontainerservice.ManagedClu`
`100`	`100`	`}`
`101`	`101`	`}`
`102`	`102`
	`103`	`+ if err := addFirewallRules(ctx, cluster, *cluster.Location); err != nil {`
	`104`	`+ return nil, fmt.Errorf("add firewall rules: %w", err)`
	`105`	`+ }`
	`106`	`+`
`103`	`107`	`kubeletIdentity, err := getClusterKubeletIdentity(cluster)`
`104`	`108`	`if err != nil {`
`105`	`109`	`return nil, fmt.Errorf("getting cluster kubelet identity: %w", err)`