feat: new e2e uses bastion hosts for ssh

Author: awesomenix

.pipelines/scripts/e2e_run.sh

@@ -64,11 +64,32 @@ if [ -n "${SIG_GALLERY_NAME}" ]; then
   export GALLERY_NAME=$SIG_GALLERY_NAME
 fi
 
+az extension add --name bastion
+
 # this software is used to take the output of "go test" and produce a junit report that we can upload to the pipeline
 # and see fancy test results.
 cd e2e
 mkdir -p bin
-GOBIN=`pwd`/bin/ go install gotest.tools/gotestsum@latest
+architecture=$(uname -m)
+
+case "$architecture" in
+  x86_64 | amd64) architecture="amd64" ;;
+  aarch64 | arm64) architecture="arm64" ;;
+  *)
+    echo "Unsupported architecture: $architecture"
+    exit 1
+    ;;
+esac
+
+gotestsum_version="1.13.0"
+gotestsum_archive="gotestsum_${gotestsum_version}_linux_${architecture}.tar.gz"
+gotestsum_url="https://github.com/gotestyourself/gotestsum/releases/download/v${gotestsum_version}/${gotestsum_archive}"
+
+temp_file="$(mktemp)"
+curl -fsSL "$gotestsum_url" -o "$temp_file"
+tar -xzf "$temp_file" -C bin
+chmod +x bin/gotestsum
+rm -f "$temp_file"
 
 # gotestsum configure to only show logs for failed tests, json file for detailed logs
 # Run the tests! Yey!

.pipelines/templates/e2e-template.yaml

@@ -34,6 +34,8 @@ jobs:
         displayName: Run AgentBaker E2E
         env:
           E2E_SUBSCRIPTION_ID: $(E2E_SUBSCRIPTION_ID)
+          SYS_SSH_PUBLIC_KEY: $(SYS_SSH_PUBLIC_KEY)
+          SYS_SSH_PRIVATE_KEY_B64: $(SYS_SSH_PRIVATE_KEY_B64)
           BUILD_SRC_DIR: $(System.DefaultWorkingDirectory)
           DefaultWorkingDirectory: $(Build.SourcesDirectory)
           VHD_BUILD_ID: $(VHD_BUILD_ID)

e2e/aks_model.go

@@ -175,6 +175,16 @@ func getBaseClusterModel(clusterName, location, k8sSystemPoolSKU string) *armcon
 					Enabled: to.Ptr(false),
 				},
 			},
+			LinuxProfile: &armcontainerservice.LinuxProfile{
+				AdminUsername: to.Ptr("azureuser"),
+				SSH: &armcontainerservice.SSHConfiguration{
+					PublicKeys: []*armcontainerservice.SSHPublicKey{
+						{
+							KeyData: to.Ptr(string(config.SysSSHPublicKey)),
+						},
+					},
+				},
+			},
 		},
 		Identity: &armcontainerservice.ManagedClusterIdentity{
 			Type: to.Ptr(armcontainerservice.ResourceIdentityTypeSystemAssigned),

e2e/cache.go

@@ -157,7 +157,7 @@ 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-v3", request.Location, request.K8sSystemPoolSKU), false, false)
+	return prepareCluster(ctx, getKubenetClusterModel("abe2e-kubenet-v4", request.Location, request.K8sSystemPoolSKU), false, false)
 }
 
 var ClusterKubenetAirgap = cachedFunc(clusterKubenetAirgap)

e2e/cluster.go

@@ -8,6 +8,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"net/netip"
 	"strings"
 	"time"
 
@@ -17,6 +18,7 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/authorization/armauthorization/v2"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v6"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v6"
+	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/network/armnetwork/v6"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources"
 	"github.com/google/uuid"
 	corev1 "k8s.io/api/core/v1"
@@ -39,6 +41,7 @@ type Cluster struct {
 	ClusterParams   *ClusterParams
 	Maintenance     *armcontainerservice.MaintenanceConfiguration
 	DebugPod        *corev1.Pod
+	Bastion         *armnetwork.BastionHost
 }
 
 // Returns true if the cluster is configured with Azure CNI
@@ -66,6 +69,11 @@ func prepareCluster(ctx context.Context, cluster *armcontainerservice.ManagedClu
 		return nil, fmt.Errorf("get or create cluster: %w", err)
 	}
 
+	bastion, err := getOrCreateBastion(ctx, cluster)
+	if err != nil {
+		return nil, fmt.Errorf("get or create bastion: %w", err)
+	}
+
 	maintenance, err := getOrCreateMaintenanceConfiguration(ctx, cluster)
 	if err != nil {
 		return nil, fmt.Errorf("get or create maintenance configuration: %w", err)
@@ -143,6 +151,7 @@ func prepareCluster(ctx context.Context, cluster *armcontainerservice.ManagedClu
 		Maintenance:     maintenance,
 		ClusterParams:   clusterParams,
 		DebugPod:        hostPod,
+		Bastion:         bastion,
 	}, nil
 }
 
@@ -470,6 +479,205 @@ func createNewMaintenanceConfiguration(ctx context.Context, cluster *armcontaine
 	return &maintenance, nil
 }
 
+func getOrCreateBastion(ctx context.Context, cluster *armcontainerservice.ManagedCluster) (*armnetwork.BastionHost, error) {
+	nodeRG := *cluster.Properties.NodeResourceGroup
+	bastionName := fmt.Sprintf("%s-bastion", *cluster.Name)
+
+	existing, err := config.Azure.BastionHosts.Get(ctx, nodeRG, bastionName, nil)
+	var azErr *azcore.ResponseError
+	if errors.As(err, &azErr) && azErr.StatusCode == http.StatusNotFound {
+		return createNewBastion(ctx, cluster)
+	}
+	if err != nil {
+		return nil, fmt.Errorf("failed to get bastion %q in rg %q: %w", bastionName, nodeRG, err)
+	}
+	return &existing.BastionHost, nil
+}
+
+func createNewBastion(ctx context.Context, cluster *armcontainerservice.ManagedCluster) (*armnetwork.BastionHost, error) {
+	nodeRG := *cluster.Properties.NodeResourceGroup
+	location := *cluster.Location
+	bastionName := fmt.Sprintf("%s-bastion", *cluster.Name)
+	publicIPName := fmt.Sprintf("%s-bastion-pip", *cluster.Name)
+	publicIPName = sanitizeAzureResourceName(publicIPName)
+
+	vnet, err := getClusterVNet(ctx, nodeRG)
+	if err != nil {
+		return nil, fmt.Errorf("get cluster vnet in rg %q: %w", nodeRG, err)
+	}
+
+	// Azure Bastion requires a dedicated subnet named AzureBastionSubnet. Standard SKU (required for
+	// native client support/tunneling) requires at least a /26.
+	bastionSubnetName := "AzureBastionSubnet"
+	bastionSubnetPrefix := "10.226.0.0/26"
+	if _, err := netip.ParsePrefix(bastionSubnetPrefix); err != nil {
+		return nil, fmt.Errorf("invalid bastion subnet prefix %q: %w", bastionSubnetPrefix, err)
+	}
+
+	var bastionSubnetID string
+	bastionSubnet, subnetGetErr := config.Azure.Subnet.Get(ctx, nodeRG, vnet.name, bastionSubnetName, nil)
+	if subnetGetErr != nil {
+		var subnetAzErr *azcore.ResponseError
+		if !errors.As(subnetGetErr, &subnetAzErr) || subnetAzErr.StatusCode != http.StatusNotFound {
+			return nil, fmt.Errorf("get subnet %q in vnet %q rg %q: %w", bastionSubnetName, vnet.name, nodeRG, subnetGetErr)
+		}
+
+		logf(ctx, "creating subnet %s in VNet %s (rg %s)", bastionSubnetName, vnet.name, nodeRG)
+		subnetParams := armnetwork.Subnet{
+			Properties: &armnetwork.SubnetPropertiesFormat{
+				AddressPrefix: to.Ptr(bastionSubnetPrefix),
+			},
+		}
+		subnetPoller, err := config.Azure.Subnet.BeginCreateOrUpdate(ctx, nodeRG, vnet.name, bastionSubnetName, subnetParams, nil)
+		if err != nil {
+			return nil, fmt.Errorf("failed to start creating bastion subnet: %w", err)
+		}
+		bastionSubnet, err := subnetPoller.PollUntilDone(ctx, config.DefaultPollUntilDoneOptions)
+		if err != nil {
+			return nil, fmt.Errorf("failed to create bastion subnet: %w", err)
+		}
+		bastionSubnetID = *bastionSubnet.ID
+	} else {
+		bastionSubnetID = *bastionSubnet.ID
+	}
+
+	// Public IP for Bastion
+	pipParams := 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 bastion public IP %s (rg %s)", publicIPName, nodeRG)
+	pipPoller, err := config.Azure.PublicIPAddresses.BeginCreateOrUpdate(ctx, nodeRG, publicIPName, pipParams, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to start creating bastion public IP: %w", err)
+	}
+	pipResp, err := pipPoller.PollUntilDone(ctx, config.DefaultPollUntilDoneOptions)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create bastion public IP: %w", err)
+	}
+	if pipResp.ID == nil {
+		return nil, fmt.Errorf("bastion public IP response missing ID")
+	}
+
+	bastion := armnetwork.BastionHost{
+		Location: to.Ptr(location),
+		SKU: &armnetwork.SKU{
+			Name: to.Ptr(armnetwork.BastionHostSKUNameStandard),
+		},
+		Properties: &armnetwork.BastionHostPropertiesFormat{
+			// Native client support is enabled via tunneling.
+			EnableTunneling: to.Ptr(true),
+			IPConfigurations: []*armnetwork.BastionHostIPConfiguration{
+				{
+					Name: to.Ptr("bastion-ipcfg"),
+					Properties: &armnetwork.BastionHostIPConfigurationPropertiesFormat{
+						Subnet: &armnetwork.SubResource{
+							ID: to.Ptr(bastionSubnetID),
+						},
+						PublicIPAddress: &armnetwork.SubResource{
+							ID: pipResp.ID,
+						},
+					},
+				},
+			},
+		},
+	}
+
+	logf(ctx, "creating bastion %s (native client/tunneling enabled) in rg %s", bastionName, nodeRG)
+	bastionPoller, err := config.Azure.BastionHosts.BeginCreateOrUpdate(ctx, nodeRG, bastionName, bastion, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to start creating bastion: %w", err)
+	}
+	resp, err := bastionPoller.PollUntilDone(ctx, config.DefaultPollUntilDoneOptions)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create bastion: %w", err)
+	}
+
+	if err := verifyBastion(ctx, cluster, &resp.BastionHost); err != nil {
+		return nil, fmt.Errorf("failed to verify bastion: %w", err)
+	}
+	return &resp.BastionHost, nil
+}
+
+func verifyBastion(ctx context.Context, cluster *armcontainerservice.ManagedCluster, bastion *armnetwork.BastionHost) error {
+	nodeRG := *cluster.Properties.NodeResourceGroup
+	vmssName, err := getSystemPoolVMSSName(ctx, cluster)
+	if err != nil {
+		return err
+	}
+
+	var vmssVM *armcompute.VirtualMachineScaleSetVM
+	pager := config.Azure.VMSSVM.NewListPager(nodeRG, vmssName, nil)
+	if pager.More() {
+		page, err := pager.NextPage(ctx)
+		if err != nil {
+			return fmt.Errorf("list vmss vms for %q in rg %q: %w", vmssName, nodeRG, err)
+		}
+		if len(page.Value) > 0 {
+			vmssVM = page.Value[0]
+		}
+	}
+
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	localPort, pid, err := startBastionTunnel(ctx, *bastion.Name, nodeRG, *vmssVM.ID)
+	if err != nil {
+		return err
+	}
+
+	defer cleanupBastionTunnel(localPort, pid)
+
+	result, err := runSSHCommandWithPrivateKeyFile(ctx, localPort, "uname -a", config.SysSSHPrivateKey)
+	if err != nil {
+		return err
+	}
+	if strings.Contains(result.stdout, *vmssVM.Name) {
+		return nil
+	}
+	return fmt.Errorf("Executed ssh on wrong VM: %s", result.stdout)
+}
+
+func getSystemPoolVMSSName(ctx context.Context, cluster *armcontainerservice.ManagedCluster) (string, error) {
+	nodeRG := *cluster.Properties.NodeResourceGroup
+	var systemPoolName string
+	for _, pool := range cluster.Properties.AgentPoolProfiles {
+		if strings.EqualFold(string(*pool.Mode), "System") {
+			systemPoolName = *pool.Name
+		}
+	}
+	pager := config.Azure.VMSS.NewListPager(nodeRG, nil)
+	if pager.More() {
+		page, err := pager.NextPage(ctx)
+		if err != nil {
+			return "", fmt.Errorf("list vmss in rg %q: %w", nodeRG, err)
+		}
+		for _, vmss := range page.Value {
+			if strings.Contains(strings.ToLower(*vmss.Name), strings.ToLower(systemPoolName)) {
+				return *vmss.Name, nil
+			}
+		}
+	}
+	return "", fmt.Errorf("no matching VMSS found for system pool %q in rg %q", systemPoolName, nodeRG)
+}
+
+func sanitizeAzureResourceName(name string) string {
+	// Azure resource name restrictions vary by type. For our usage here (Public IP name) we just
+	// keep it simple and strip problematic characters.
+	replacer := strings.NewReplacer("/", "-", "\\", "-", ":", "-", "_", "-", " ", "-")
+	name = replacer.Replace(name)
+	name = strings.Trim(name, "-")
+	if len(name) > 80 {
+		name = name[:80]
+	}
+	return name
+}
+
 type VNet struct {
 	name     string
 	subnetId string

e2e/config/azure.go

@@ -40,6 +40,7 @@ import (
 type AzureClient struct {
 	AKS                       *armcontainerservice.ManagedClustersClient
 	AzureFirewall             *armnetwork.AzureFirewallsClient
+	BastionHosts              *armnetwork.BastionHostsClient
 	Blob                      *azblob.Client
 	StorageContainers         *armstorage.BlobContainersClient
 	CacheRulesClient          *armcontainerregistry.CacheRulesClient
@@ -156,6 +157,16 @@ func NewAzureClient() (*AzureClient, error) {
 		return nil, fmt.Errorf("create public ip addresses client: %w", err)
 	}
 
+	cloud.BastionHosts, err = armnetwork.NewBastionHostsClient(Config.SubscriptionID, credential, opts)
+	if err != nil {
+		return nil, fmt.Errorf("create bastion hosts client: %w", err)
+	}
+
+	cloud.BastionHosts, err = armnetwork.NewBastionHostsClient(Config.SubscriptionID, credential, opts)
+	if err != nil {
+		return nil, fmt.Errorf("create bastion hosts 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)