Run live Workload Identity (AKS) test in weekly pipeline (Azure#22627)

chlowell · web-flow · commit 17d2fb34df3a · 2024-03-22T12:47:16.000-07:00
diff --git a/sdk/azidentity/.gitignore b/sdk/azidentity/.gitignore
@@ -0,0 +1,4 @@
+# live test artifacts
+Dockerfile
+k8s.yaml
+sshkey*
diff --git a/sdk/azidentity/assets.json b/sdk/azidentity/assets.json
@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "go",
   "TagPrefix": "go/azidentity",
-  "Tag": "go/azidentity_db4a26f583"
+  "Tag": "go/azidentity_0a700f9ea1"
 }
diff --git a/sdk/azidentity/test-resources-post.ps1 b/sdk/azidentity/test-resources-post.ps1
@@ -4,20 +4,24 @@
 # IMPORTANT: Do not invoke this file directly. Please instead run eng/common/TestResources/New-TestResources.ps1 from the repository root.
 
 param (
-    [hashtable] $AdditionalParameters = @{},
-    [hashtable] $DeploymentOutputs
+  [hashtable] $AdditionalParameters = @{},
+  [hashtable] $DeploymentOutputs
 )
 
-if (!$AdditionalParameters['deployResources']) {
+$ErrorActionPreference = 'Stop'
+$PSNativeCommandUseErrorActionPreference = $true
+
+if ($CI) {
+  if (!$AdditionalParameters['deployResources']) {
     Write-Host "Skipping post-provisioning script because resources weren't deployed"
     return
+  }
+  az login --service-principal -u $DeploymentOutputs['AZIDENTITY_CLIENT_ID'] -p $DeploymentOutputs['AZIDENTITY_CLIENT_SECRET'] --tenant $DeploymentOutputs['AZIDENTITY_TENANT_ID']
+  az account set --subscription $DeploymentOutputs['AZIDENTITY_SUBSCRIPTION_ID']
 }
 
-$ErrorActionPreference = 'Stop'
-$PSNativeCommandUseErrorActionPreference = $true
-
 Write-Host "Building container"
-$image = "azidentity-managed-id-test"
+$image = "$($DeploymentOutputs['AZIDENTITY_ACR_LOGIN_SERVER'])/azidentity-managed-id-test"
 Set-Content -Path "$PSScriptRoot/Dockerfile" -Value @"
 FROM mcr.microsoft.com/oss/go/microsoft/golang:latest as builder
 ENV GOARCH=amd64 GOWORK=off
@@ -34,16 +38,64 @@ CMD ["./managed-id-test"]
 "@
 # build from sdk/azidentity because we need that dir in the context (because the test app uses local azidentity)
 docker build -t $image "$PSScriptRoot"
+az acr login -n $DeploymentOutputs['AZIDENTITY_ACR_NAME']
+docker push $image
 
-az login --service-principal -u $DeploymentOutputs['AZIDENTITY_CLIENT_ID'] -p $DeploymentOutputs['AZIDENTITY_CLIENT_SECRET'] --tenant $DeploymentOutputs['AZIDENTITY_TENANT_ID']
-az account set --subscription $DeploymentOutputs['AZIDENTITY_SUBSCRIPTION_ID']
+$rg = $DeploymentOutputs['AZIDENTITY_RESOURCE_GROUP']
 
 # Azure Functions deployment: copy the Windows binary from the Docker image, deploy it in a zip
 Write-Host "Deploying to Azure Functions"
 $container = docker create $image
 docker cp ${container}:managed-id-test.exe "$PSScriptRoot/testdata/managed-id-test/"
 docker rm -v $container
 Compress-Archive -Path "$PSScriptRoot/testdata/managed-id-test/*" -DestinationPath func.zip -Force
-az functionapp deploy -g $DeploymentOutputs['AZIDENTITY_RESOURCE_GROUP'] -n $DeploymentOutputs['AZIDENTITY_FUNCTION_NAME'] --src-path func.zip --type zip
+az functionapp deploy -g $rg -n $DeploymentOutputs['AZIDENTITY_FUNCTION_NAME'] --src-path func.zip --type zip
+
+Write-Host "Creating federated identity"
+$aksName = $DeploymentOutputs['AZIDENTITY_AKS_NAME']
+$idName = $DeploymentOutputs['AZIDENTITY_USER_ASSIGNED_IDENTITY_NAME']
+$issuer = az aks show -g $rg -n $aksName --query "oidcIssuerProfile.issuerUrl" -otsv
+$podName = "azidentity-test"
+$serviceAccountName = "workload-identity-sa"
+az identity federated-credential create -g $rg --identity-name $idName --issuer $issuer --name $idName --subject system:serviceaccount:default:$serviceAccountName
+Write-Host "Deploying to AKS"
+az aks get-credentials -g $rg -n $aksName
+az aks update --attach-acr $DeploymentOutputs['AZIDENTITY_ACR_NAME'] -g $rg -n $aksName
+Set-Content -Path "$PSScriptRoot/k8s.yaml" -Value @"
+apiVersion: v1
+kind: ServiceAccount
+metadata:
+  annotations:
+    azure.workload.identity/client-id: $($DeploymentOutputs['AZIDENTITY_USER_ASSIGNED_IDENTITY_CLIENT_ID'])
+  name: $serviceAccountName
+  namespace: default
+---
+apiVersion: v1
+kind: Pod
+metadata:
+  name: $podName
+  namespace: default
+  labels:
+    app: $podName
+    azure.workload.identity/use: "true"
+spec:
+  serviceAccountName: $serviceAccountName
+  containers:
+  - name: $podName
+    image: $image
+    env:
+    - name: AZIDENTITY_STORAGE_NAME
+      value: $($DeploymentOutputs['AZIDENTITY_STORAGE_NAME_USER_ASSIGNED'])
+    - name: AZIDENTITY_USE_WORKLOAD_IDENTITY
+      value: "true"
+    - name: FUNCTIONS_CUSTOMHANDLER_PORT
+      value: "80"
+  nodeSelector:
+    kubernetes.io/os: linux
+"@
+kubectl apply -f "$PSScriptRoot/k8s.yaml"
+Write-Host "##vso[task.setvariable variable=AZIDENTITY_POD_NAME;]$podName"
 
-az logout
+if ($CI) {
+  az logout
+}
diff --git a/sdk/azidentity/test-resources-pre.ps1 b/sdk/azidentity/test-resources-pre.ps1
@@ -12,6 +12,11 @@ param (
     $RemainingArguments
 )
 
+if (-not (Test-Path "$PSScriptRoot/sshkey.pub")) {
+    ssh-keygen -t rsa -b 4096 -f "$PSScriptRoot/sshkey" -N '' -C ''
+}
+$templateFileParameters['sshPubKey'] = Get-Content "$PSScriptRoot/sshkey.pub"
+
 if (!$CI) {
     # TODO: Remove this once auto-cloud config downloads are supported locally
     Write-Host "Skipping cert setup in local testing mode"
diff --git a/sdk/azidentity/test-resources.bicep b/sdk/azidentity/test-resources.bicep
@@ -1,6 +1,9 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+@description('Kubernetes cluster admin user name.')
+param adminUser string = 'azureuser'
+
 @minLength(6)
 @maxLength(23)
 @description('The base resource name.')
@@ -9,6 +12,8 @@ param baseName string = resourceGroup().name
 @description('Whether to deploy resources. When set to false, this file deploys nothing.')
 param deployResources bool = false
 
+param sshPubKey string = ''
+
 @description('The location of the resource. By default, this is the same as the resource group.')
 param location string = resourceGroup().location
 
@@ -18,7 +23,7 @@ var blobContributor = subscriptionResourceId('Microsoft.Authorization/roleDefini
 resource sa 'Microsoft.Storage/storageAccounts@2021-08-01' = if (deployResources) {
   kind: 'StorageV2'
   location: location
-  name: baseName
+  name: 'sa${uniqueString(baseName)}'
   properties: {
     accessTier: 'Hot'
   }
@@ -30,7 +35,7 @@ resource sa 'Microsoft.Storage/storageAccounts@2021-08-01' = if (deployResources
 resource saUserAssigned 'Microsoft.Storage/storageAccounts@2021-08-01' = if (deployResources) {
   kind: 'StorageV2'
   location: location
-  name: '${baseName}2'
+  name: 'sa2${uniqueString(baseName)}'
   properties: {
     accessTier: 'Hot'
   }
@@ -64,6 +69,17 @@ resource blobRoleFunc 'Microsoft.Authorization/roleAssignments@2022-04-01' = if
   scope: sa
 }
 
+resource containerRegistry 'Microsoft.ContainerRegistry/registries@2023-01-01-preview' = if (deployResources) {
+  location: location
+  name: uniqueString(resourceGroup().id)
+  properties: {
+    adminUserEnabled: true
+  }
+  sku: {
+    name: 'Basic'
+  }
+}
+
 resource farm 'Microsoft.Web/serverfarms@2021-03-01' = if (deployResources) {
   kind: 'app'
   location: location
@@ -135,7 +151,57 @@ resource azfunc 'Microsoft.Web/sites@2021-03-01' = if (deployResources) {
   }
 }
 
+resource aks 'Microsoft.ContainerService/managedClusters@2023-06-01' = if (deployResources) {
+  name: baseName
+  location: location
+  identity: {
+    type: 'SystemAssigned'
+  }
+  properties: {
+    agentPoolProfiles: [
+      {
+        count: 1
+        enableAutoScaling: false
+        kubeletDiskType: 'OS'
+        mode: 'System'
+        name: 'agentpool'
+        osDiskSizeGB: 128
+        osDiskType: 'Managed'
+        osSKU: 'Ubuntu'
+        osType: 'Linux'
+        type: 'VirtualMachineScaleSets'
+        vmSize: 'Standard_D2s_v3'
+      }
+    ]
+    dnsPrefix: 'identitytest'
+    enableRBAC: true
+    linuxProfile: {
+      adminUsername: adminUser
+      ssh: {
+        publicKeys: [
+          {
+            keyData: sshPubKey
+          }
+        ]
+      }
+    }
+    oidcIssuerProfile: {
+      enabled: true
+    }
+    securityProfile: {
+      workloadIdentity: {
+        enabled: true
+      }
+    }
+  }
+}
+
+output AZIDENTITY_ACR_LOGIN_SERVER string = deployResources ? containerRegistry.properties.loginServer : ''
+output AZIDENTITY_ACR_NAME string = deployResources ? containerRegistry.name : ''
+output AZIDENTITY_AKS_NAME string = deployResources ? aks.name : ''
 output AZIDENTITY_FUNCTION_NAME string = deployResources ? azfunc.name : ''
 output AZIDENTITY_STORAGE_NAME string = deployResources ? sa.name : ''
 output AZIDENTITY_STORAGE_NAME_USER_ASSIGNED string = deployResources ? saUserAssigned.name : ''
 output AZIDENTITY_USER_ASSIGNED_IDENTITY string = deployResources ? usermgdid.id : ''
+output AZIDENTITY_USER_ASSIGNED_IDENTITY_CLIENT_ID string = deployResources ? usermgdid.properties.clientId : ''
+output AZIDENTITY_USER_ASSIGNED_IDENTITY_NAME string = deployResources ? usermgdid.name : ''
diff --git a/sdk/azidentity/testdata/managed-id-test/main.go b/sdk/azidentity/testdata/managed-id-test/main.go
@@ -18,18 +18,38 @@ import (
 
 var (
 	config = struct {
-		id                      string
-		storageName             string
+		// resourceID of a managed identity permitted to list blobs in the account specified by storageNameUserAssigned.
+		resourceID string
+		// storageName is the name of a storage account accessible by the default or system-assigned identity
+		storageName string
+		// storageNameUserAssigned is the name of a storage account accessible by the identity specified by
+		// resourceID. The default or system-assigned identity shouldn't have any permission for this account.
 		storageNameUserAssigned string
+		// workloadID determines whether this app tests ManagedIdentityCredential or WorkloadIdentityCredential.
+		// When true, the app ignores resourceID and storageNameUserAssigned.
+		workloadID bool
 	}{
-		id:                      os.Getenv("AZIDENTITY_USER_ASSIGNED_IDENTITY"),
+		resourceID:              os.Getenv("AZIDENTITY_USER_ASSIGNED_IDENTITY"),
 		storageName:             os.Getenv("AZIDENTITY_STORAGE_NAME"),
 		storageNameUserAssigned: os.Getenv("AZIDENTITY_STORAGE_NAME_USER_ASSIGNED"),
+		workloadID:              os.Getenv("AZIDENTITY_USE_WORKLOAD_IDENTITY") != "",
 	}
 
 	missingConfig string
 )
 
+func credential(resourceID string) (azcore.TokenCredential, error) {
+	if config.workloadID {
+		// the identity is determined by service account configuration
+		return azidentity.NewWorkloadIdentityCredential(nil)
+	}
+	opts := azidentity.ManagedIdentityCredentialOptions{}
+	if resourceID != "" {
+		opts.ID = azidentity.ResourceID(resourceID)
+	}
+	return azidentity.NewManagedIdentityCredential(&opts)
+}
+
 func listContainers(account string, cred azcore.TokenCredential) error {
 	url := fmt.Sprintf("https://%s.blob.core.windows.net", account)
 	log.Printf("listing containers in %s", url)
@@ -48,20 +68,14 @@ func handler(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	cred, err := azidentity.NewManagedIdentityCredential(nil)
-	if err != nil {
-		w.WriteHeader(http.StatusInternalServerError)
-		fmt.Fprint(w, err)
-		log.Print(err)
-		return
-	}
-	err = listContainers(config.storageName, cred)
+	cred, err := credential("")
 	if err == nil {
-		cred, err = azidentity.NewManagedIdentityCredential(&azidentity.ManagedIdentityCredentialOptions{
-			ID: azidentity.ResourceID(config.id),
-		})
-		if err == nil {
-			err = listContainers(config.storageNameUserAssigned, cred)
+		err = listContainers(config.storageName, cred)
+		if !config.workloadID && err == nil {
+			cred, err = credential(config.resourceID)
+			if err == nil {
+				err = listContainers(config.storageNameUserAssigned, cred)
+			}
 		}
 	}
 
@@ -77,14 +91,19 @@ func handler(w http.ResponseWriter, r *http.Request) {
 
 func main() {
 	v := []string{}
-	if config.id == "" {
-		v = append(v, "AZIDENTITY_USER_ASSIGNED_IDENTITY")
-	}
 	if config.storageName == "" {
 		v = append(v, "AZIDENTITY_STORAGE_NAME")
 	}
-	if config.storageNameUserAssigned == "" {
-		v = append(v, "AZIDENTITY_STORAGE_NAME_USER_ASSIGNED")
+	if config.workloadID {
+		log.Print("Testing WorkloadIdentityCredential")
+	} else {
+		log.Print("Testing ManagedIdentityCredential")
+		if config.resourceID == "" {
+			v = append(v, "AZIDENTITY_USER_ASSIGNED_IDENTITY")
+		}
+		if config.storageNameUserAssigned == "" {
+			v = append(v, "AZIDENTITY_STORAGE_NAME_USER_ASSIGNED")
+		}
 	}
 	if len(v) > 0 {
 		missingConfig = strings.Join(v, ", ")
diff --git a/sdk/azidentity/workload_identity_test.go b/sdk/azidentity/workload_identity_test.go
@@ -16,6 +16,7 @@ import (
 	"fmt"
 	"net/http"
 	"os"
+	"os/exec"
 	"path/filepath"
 	"strconv"
 	"strings"
@@ -25,6 +26,7 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	"github.com/golang-jwt/jwt/v5"
 	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
 )
 
 func assertion(cert *x509.Certificate, key crypto.PrivateKey) (string, error) {
@@ -46,6 +48,22 @@ func assertion(cert *x509.Certificate, key crypto.PrivateKey) (string, error) {
 }
 
 func TestWorkloadIdentityCredential_Live(t *testing.T) {
+	// This test triggers the managed identity test app deployed to Azure Kubernetes Service.
+	// See the bicep file and test resources scripts for details.
+	// It triggers the app with kubectl because the test subscription prohibits opening ports to the internet.
+	pod := os.Getenv("AZIDENTITY_POD_NAME")
+	if pod == "" {
+		t.Skip("set AZIDENTITY_POD_NAME to run this test")
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
+	defer cancel()
+	cmd := exec.CommandContext(ctx, "kubectl", "exec", pod, "--", "wget", "-qO-", "localhost")
+	b, err := cmd.CombinedOutput()
+	require.NoError(t, err)
+	require.EqualValues(t, "test passed", b)
+}
+
+func TestWorkloadIdentityCredential_Recorded(t *testing.T) {
 	cert, err := os.ReadFile(liveSP.pemPath)
 	if err != nil {
 		t.Fatal(err)

Original file line number	Diff line number	Diff line change
`@@ -2,5 +2,5 @@`
`2`	`2`	`"AssetsRepo": "Azure/azure-sdk-assets",`
`3`	`3`	`"AssetsRepoPrefixPath": "go",`
`4`	`4`	`"TagPrefix": "go/azidentity",`
`5`		`- "Tag": "go/azidentity_db4a26f583"`
	`5`	`+ "Tag": "go/azidentity_0a700f9ea1"`
`6`	`6`	`}`