feat: azure/aks workload identity support

.gitignore

@@ -47,6 +47,11 @@ inference-extension-conformance-test-report.yaml
 .claude/
 .env
 .mcp.json
+.serena/
+CLAUDE.md
 
 .goose
 /aigw
+
+# Go test binaries
+*.test

api/v1alpha1/backendsecurity_policy.go

@@ -232,22 +232,29 @@ type BackendSecurityPolicyGCPCredentials struct {
 }
 
 // BackendSecurityPolicyAzureCredentials contains the supported authentication mechanisms to access Azure.
-// Only one of ClientSecretRef or OIDCExchangeToken must be specified. Credentials will not be generated if
-// neither are set.
+// One of ClientSecretRef, OIDCExchangeToken, or UseManagedIdentity must be specified.
+// When UseManagedIdentity is true, neither ClientSecretRef nor OIDCExchangeToken should be set.
+// Otherwise, exactly one of ClientSecretRef or OIDCExchangeToken must be specified.
 //
-// +kubebuilder:validation:XValidation:rule="(has(self.clientSecretRef) && !has(self.oidcExchangeToken)) || (!has(self.clientSecretRef) && has(self.oidcExchangeToken))",message="Exactly one of clientSecretRef or oidcExchangeToken must be specified"
+// +kubebuilder:validation:XValidation:rule="has(self.useManagedIdentity) && self.useManagedIdentity ? (!has(self.clientSecretRef) && !has(self.oidcExchangeToken)) : ((has(self.clientSecretRef) && !has(self.oidcExchangeToken)) || (!has(self.clientSecretRef) && has(self.oidcExchangeToken)))",message="When useManagedIdentity is true, clientSecretRef and oidcExchangeToken must not be specified. Otherwise, exactly one of clientSecretRef or oidcExchangeToken must be specified"
 type BackendSecurityPolicyAzureCredentials struct {
 	// ClientID is a unique identifier for an application in Azure.
+	// This field is optional when using managed identity or workload identity,
+	// as the value will be provided via environment variables (AZURE_CLIENT_ID).
+	// Required for other authentication methods.
 	//
-	// +kubebuilder:validation:Required
+	// +optional
 	// +kubebuilder:validation:MinLength=1
-	ClientID string `json:"clientID"`
+	ClientID string `json:"clientID,omitempty"`
 
 	// TenantId is a unique identifier for an Azure Active Directory instance.
+	// This field is optional when using workload identity with service account annotations,
+	// as the value will be provided via environment variables (AZURE_TENANT_ID).
+	// Required for other authentication methods.
 	//
-	// +kubebuilder:validation:Required
+	// +optional
 	// +kubebuilder:validation:MinLength=1
-	TenantID string `json:"tenantID"`
+	TenantID string `json:"tenantID,omitempty"`
 
 	// ClientSecretRef is the reference to the secret containing the Azure client secret.
 	// ai-gateway must be given the permission to read this secret.
@@ -261,6 +268,16 @@ type BackendSecurityPolicyAzureCredentials struct {
 	//
 	// +optional
 	OIDCExchangeToken *AzureOIDCExchangeToken `json:"oidcExchangeToken,omitempty"`
+
+	// UseManagedIdentity enables Azure Managed Identity authentication.
+	// When set to true, the gateway will use DefaultAzureCredential which supports:
+	// - Environment variables (AZURE_CLIENT_ID, AZURE_TENANT_ID, AZURE_CLIENT_SECRET, AZURE_FEDERATED_TOKEN_FILE)
+	// - AKS Workload Identity (via service account annotations)
+	// - System-assigned managed identity (when clientID is not specified)
+	// - User-assigned managed identity (when clientID is specified)
+	//
+	// +optional
+	UseManagedIdentity *bool `json:"useManagedIdentity,omitempty"`
 }
 
 // AzureOIDCExchangeToken specifies credentials to obtain oidc token from a sso server.

examples/basic/azure_openai.yaml

@@ -3,6 +3,29 @@
 # The full text of the Apache license is available in the LICENSE file at
 # the root of the repo.
 
+# Azure OpenAI Authentication Examples
+#
+# This file demonstrates two authentication methods for Azure OpenAI:
+# 1. Client Secret (default, uncommented)
+# 2. Managed Identity / Workload Identity (commented alternatives)
+#
+# For AKS Workload Identity setup, configure the service account during Helm install:
+#
+#   helm upgrade --install ai-gateway-controller envoyproxy/ai-gateway-helm \
+#     --set controller.serviceAccount.annotations."azure\.workload\.identity/client-id"="<your-managed-identity-client-id>" \
+#     --set controller.serviceAccount.annotations."azure\.workload\.identity/tenant-id"="<your-tenant-id>" \
+#     --create-namespace \
+#     -n envoy-ai-gateway-system
+#
+#   # Add the required label:
+#   kubectl label serviceaccount ai-gateway-controller \
+#     azure.workload.identity/use=true \
+#     -n envoy-ai-gateway-system
+#
+# For complete Azure infrastructure setup, see:
+# https://learn.microsoft.com/en-us/azure/aks/workload-identity-deploy-cluster
+
+---
 apiVersion: aigateway.envoyproxy.io/v1alpha1
 kind: AIGatewayRoute
 metadata:
@@ -39,6 +62,7 @@ spec:
     kind: Backend
     group: gateway.envoyproxy.io
 ---
+# Option 1: Client Secret Authentication (Default)
 apiVersion: aigateway.envoyproxy.io/v1alpha1
 kind: BackendSecurityPolicy
 metadata:
@@ -57,6 +81,26 @@ spec:
       name: envoy-ai-gateway-basic-azure-client-secret
       namespace: default
 ---
+# Option 2: Managed Identity / Workload Identity Authentication
+# Uncomment this section and comment out Option 1 to use Managed Identity.
+#
+# For AKS Workload Identity: Configure service account annotations during Helm install
+# (see header comment above for helm upgrade command)
+#
+# apiVersion: aigateway.envoyproxy.io/v1alpha1
+# kind: BackendSecurityPolicy
+# metadata:
+#   name: envoy-ai-gateway-basic-azure-credentials
+#   namespace: default
+# spec:
+#   targetRefs:
+#     - group: aigateway.envoyproxy.io
+#       kind: AIServiceBackend
+#       name: envoy-ai-gateway-basic-azure
+#   type: AzureCredentials
+#   azureCredentials:
+#     useManagedIdentity: true
+---
 apiVersion: gateway.envoyproxy.io/v1alpha1
 kind: Backend
 metadata:

internal/backendauth/azure.go

@@ -8,24 +8,168 @@ package backendauth
 import (
 	"context"
 	"fmt"
+	"net/http"
+	"net/url"
+	"os"
 	"strings"
+	"sync"
+	"time"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
 
 	"github.com/envoyproxy/ai-gateway/internal/filterapi"
 	"github.com/envoyproxy/ai-gateway/internal/internalapi"
 )
 
+const azureScopeURL = "https://cognitiveservices.azure.com/.default"
+
 type azureHandler struct {
+	// For controller-managed token rotation (client secret, OIDC)
 	azureAccessToken string
+	// For extproc-managed workload identity
+	useManagedIdentity bool
+	credential         azcore.TokenCredential
+	tokenOptions       policy.TokenRequestOptions
+	cachedToken        string
+	tokenExpiry        time.Time
+	mu                 sync.RWMutex
 }
 
 func newAzureHandler(auth *filterapi.AzureAuth) (filterapi.BackendAuthHandler, error) {
-	return &azureHandler{azureAccessToken: strings.TrimSpace(auth.AccessToken)}, nil
+	if auth.UseManagedIdentity {
+		// Extproc-managed workload identity: obtain tokens dynamically
+		credential, err := createAzureCredential(auth.ClientID, auth.TenantID)
+		if err != nil {
+			return nil, fmt.Errorf("failed to create Azure credential: %w", err)
+		}
+		return &azureHandler{
+			useManagedIdentity: true,
+			credential:         credential,
+			tokenOptions:       policy.TokenRequestOptions{Scopes: []string{azureScopeURL}},
+		}, nil
+	}
+	// Controller-managed token rotation: use pre-obtained token
+	return &azureHandler{
+		useManagedIdentity: false,
+		azureAccessToken:   strings.TrimSpace(auth.AccessToken),
+	}, nil
+}
+
+// createAzureCredential creates an Azure credential based on the environment and configuration.
+// Supports AKS Workload Identity, user-assigned managed identity, and system-assigned managed identity.
+func createAzureCredential(clientID, tenantID string) (azcore.TokenCredential, error) {
+	clientOptions := getDefaultAzureCredentialOptions()
+
+	// Check if running in AKS Workload Identity environment
+	federatedTokenFile := os.Getenv("AZURE_FEDERATED_TOKEN_FILE")
+	envTenantID := os.Getenv("AZURE_TENANT_ID")
+	envClientID := os.Getenv("AZURE_CLIENT_ID")
+
+	if federatedTokenFile != "" && (tenantID != "" || envTenantID != "") {
+		// Use Workload Identity - this is the AKS Workload Identity pattern
+		if tenantID == "" {
+			tenantID = envTenantID
+		}
+		if clientID == "" {
+			clientID = envClientID
+		}
+		workloadIDOptions := &azidentity.WorkloadIdentityCredentialOptions{
+			ClientID:      clientID,
+			TenantID:      tenantID,
+			TokenFilePath: federatedTokenFile,
+		}
+		if clientOptions != nil {
+			workloadIDOptions.ClientOptions = clientOptions.ClientOptions
+		}
+		return azidentity.NewWorkloadIdentityCredential(workloadIDOptions)
+	} else if clientID != "" {
+		// User-assigned managed identity - specify the client ID.
+		// This uses Azure VM/VMSS Managed Identity via IMDS
+		managedIDOptions := &azidentity.ManagedIdentityCredentialOptions{
+			ID: azidentity.ClientID(clientID),
+		}
+		if clientOptions != nil {
+			managedIDOptions.ClientOptions = clientOptions.ClientOptions
+		}
+		return azidentity.NewManagedIdentityCredential(managedIDOptions)
+	}
+	// Use DefaultAzureCredential which will try multiple credential types,
+	// including system-assigned managed identity
+	return azidentity.NewDefaultAzureCredential(clientOptions)
+}
+
+// getDefaultAzureCredentialOptions returns the client options for Azure credentials,
+// including proxy configuration if set via environment variable.
+func getDefaultAzureCredentialOptions() *azidentity.DefaultAzureCredentialOptions {
+	if azureProxyURL := os.Getenv("AI_GATEWAY_AZURE_PROXY_URL"); azureProxyURL != "" {
+		proxyURL, err := url.Parse(azureProxyURL)
+		if err == nil {
+			customTransport := &http.Transport{Proxy: http.ProxyURL(proxyURL)}
+			customHTTPClient := &http.Client{Transport: customTransport}
+			return &azidentity.DefaultAzureCredentialOptions{
+				ClientOptions: azcore.ClientOptions{
+					Transport: customHTTPClient,
+				},
+			}
+		}
+	}
+	return nil
 }
 
 // Do implements [Handler.Do].
 //
-// Extracts the azure access token from the local file and set it as an authorization header.
-func (a *azureHandler) Do(_ context.Context, requestHeaders map[string]string, _ []byte) ([]internalapi.Header, error) {
-	requestHeaders["Authorization"] = fmt.Sprintf("Bearer %s", a.azureAccessToken)
-	return []internalapi.Header{{"Authorization", fmt.Sprintf("Bearer %s", a.azureAccessToken)}}, nil
+// For controller-managed tokens: Uses the pre-obtained access token.
+// For extproc-managed workload identity: Obtains tokens dynamically using Azure SDK.
+func (a *azureHandler) Do(ctx context.Context, requestHeaders map[string]string, _ []byte) ([]internalapi.Header, error) {
+	var token string
+	var err error
+
+	if a.useManagedIdentity {
+		// Get token dynamically using workload identity
+		token, err = a.getToken(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get Azure token: %w", err)
+		}
+	} else {
+		// Use pre-obtained token from controller
+		token = a.azureAccessToken
+	}
+
+	authorizationValue := fmt.Sprintf("Bearer %s", token)
+	requestHeaders["Authorization"] = authorizationValue
+	return []internalapi.Header{{"Authorization", authorizationValue}}, nil
+}
+
+// getToken retrieves an Azure access token, using cached token if still valid.
+func (a *azureHandler) getToken(ctx context.Context) (string, error) {
+	// Check if cached token is still valid (with 5-minute buffer)
+	a.mu.RLock()
+	if a.cachedToken != "" && time.Now().Add(5*time.Minute).Before(a.tokenExpiry) {
+		token := a.cachedToken
+		a.mu.RUnlock()
+		return token, nil
+	}
+	a.mu.RUnlock()
+
+	// Need to get a new token
+	a.mu.Lock()
+	defer a.mu.Unlock()
+
+	// Double-check after acquiring write lock
+	if a.cachedToken != "" && time.Now().Add(5*time.Minute).Before(a.tokenExpiry) {
+		return a.cachedToken, nil
+	}
+
+	// Get new token from Azure
+	azureToken, err := a.credential.GetToken(ctx, a.tokenOptions)
+	if err != nil {
+		return "", err
+	}
+
+	// Cache the new token
+	a.cachedToken = azureToken.Token
+	a.tokenExpiry = azureToken.ExpiresOn
+	return azureToken.Token, nil
 }

internal/backendauth/azure_test.go

@@ -40,3 +40,5 @@ func TestNewAzureHandler_Do(t *testing.T) {
 	require.Equal(t, "Authorization", headers[0][0])
 	require.Equal(t, "Bearer some-access-token", headers[0][1])
 }
+
+// investigate how to ruun test foor workload identity cases.

internal/controller/backend_security_policy.go

@@ -136,8 +136,11 @@
 		var provider tokenprovider.TokenProvider
 		options := policy.TokenRequestOptions{Scopes: []string{azureScopeURL}}
 
-		oidc := getBackendSecurityPolicyAuthOIDC(bsp.Spec)
-		if oidc != nil {
+		if bsp.Spec.AzureCredentials.UseManagedIdentity != nil && *bsp.Spec.AzureCredentials.UseManagedIdentity {
+			// Managed identity: extproc will obtain tokens dynamically, controller doesn't need to rotate
+			c.enqueueAIServiceBackendsForBackendSecurityPolicy(ctx, bsp)
+			return ctrl.Result{}, nil
+		} else if oidc := getBackendSecurityPolicyAuthOIDC(bsp.Spec); oidc != nil {
 			var oidcProvider tokenprovider.TokenProvider
 			oidcProvider, err = tokenprovider.NewOidcTokenProvider(ctx, c.client, oidc)
 			if err != nil {
@@ -169,7 +172,7 @@
 				return ctrl.Result{}, err
 			}
 		} else {
-			return ctrl.Result{}, fmt.Errorf("one of secret ref or oidc must be defined, namespace %s name %s", bsp.Namespace, bsp.Name)
+			return ctrl.Result{}, fmt.Errorf("one of secret ref, oidc, or managed identity must be defined, namespace %s name %s", bsp.Namespace, bsp.Name)
 		}
 
 		rotator, err = rotators.NewAzureTokenRotator(c.client, c.kube, c.logger, bsp.Namespace, bsp.Name, preRotationWindow, provider)