Enable Custom Proxy in WorkloadIdentityCredential

sdk/identity/azure-identity/pom.xml

@@ -115,6 +115,13 @@
       <version>1.17.7</version> <!-- {x-version-update;testdep_net.bytebuddy:byte-buddy-agent;external_dependency} -->
       <scope>test</scope>
     </dependency>
+    <!-- BouncyCastle is required for generating self-signed certificates with SANs in tests -->
+    <dependency>
+      <groupId>org.bouncycastle</groupId>
+      <artifactId>bcpkix-lts8on</artifactId>
+      <version>2.73.8</version> <!-- {x-version-update;org.bouncycastle:bcpkix-lts8on;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>

sdk/identity/azure-identity/src/main/java/com/azure/identity/WorkloadIdentityCredential.java

@@ -10,6 +10,9 @@
 import com.azure.core.util.CoreUtils;
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.identity.implementation.IdentityClientOptions;
+import com.azure.identity.implementation.customtokenproxy.CustomTokenProxyConfiguration;
+import com.azure.identity.implementation.customtokenproxy.CustomTokenProxyHttpClient;
+import com.azure.identity.implementation.customtokenproxy.ProxyConfig;
 import com.azure.identity.implementation.util.LoggingUtil;
 import com.azure.identity.implementation.util.ValidationUtil;
 import reactor.core.publisher.Mono;
@@ -89,6 +92,13 @@ public class WorkloadIdentityCredential implements TokenCredential {
         ClientAssertionCredential tempClientAssertionCredential = null;
         String tempClientId = null;
 
+        if (identityClientOptions.isKubernetesTokenProxyEnabled()) {
+            if (CustomTokenProxyConfiguration.isConfigured(configuration)) {
+                ProxyConfig proxyConfig = CustomTokenProxyConfiguration.parseAndValidate(configuration);
+                identityClientOptions.setHttpClient(new CustomTokenProxyHttpClient(proxyConfig));
+            }
+        }
+
         if (!(CoreUtils.isNullOrEmpty(tenantIdInput)
             || CoreUtils.isNullOrEmpty(federatedTokenFilePathInput)
             || CoreUtils.isNullOrEmpty(clientIdInput)

sdk/identity/azure-identity/src/main/java/com/azure/identity/WorkloadIdentityCredentialBuilder.java

@@ -47,6 +47,7 @@
 public class WorkloadIdentityCredentialBuilder extends AadCredentialBuilderBase<WorkloadIdentityCredentialBuilder> {
     private static final ClientLogger LOGGER = new ClientLogger(WorkloadIdentityCredentialBuilder.class);
     private String tokenFilePath;
+    private boolean enableTokenProxy = false;
 
     /**
      * Creates an instance of a WorkloadIdentityCredentialBuilder.
@@ -66,6 +67,19 @@ public WorkloadIdentityCredentialBuilder tokenFilePath(String tokenFilePath) {
         return this;
     }
 
+    /**
+     * Enables the Kubernetes token proxy feature for AKS workload identity scenarios.
+     * When enabled, the credential will attempt to use a custom token proxy configured through
+     * environment variables (AZURE_KUBERNETES_TOKEN_PROXY, AZURE_KUBERNETES_CA_FILE,
+     * AZURE_KUBERNETES_CA_DATA, AZURE_KUBERNETES_SNI_NAME).
+     *
+     * @return An updated instance of this builder with Kubernetes token proxy enabled.
+     */
+    public WorkloadIdentityCredentialBuilder enableKubernetesTokenProxy() {
+        this.enableTokenProxy = true;
+        return this;
+    }
+
     /**
      * Creates new {@link WorkloadIdentityCredential} with the configured options set.
      *
@@ -88,6 +102,8 @@ public WorkloadIdentityCredential build() {
         ValidationUtil.validate(this.getClass().getSimpleName(), LOGGER, "Client ID", clientIdInput, "Tenant ID",
             tenantIdInput, "Service Token File Path", federatedTokenFilePathInput);
 
+        identityClientOptions.setEnableKubernetesTokenProxy(this.enableTokenProxy);
+
         return new WorkloadIdentityCredential(tenantIdInput, clientIdInput, federatedTokenFilePathInput,
             identityClientOptions.clone());
     }

sdk/identity/azure-identity/src/main/java/com/azure/identity/implementation/IdentityClientOptions.java

@@ -71,6 +71,7 @@ public final class IdentityClientOptions implements Cloneable {
     private List<HttpPipelinePolicy> perRetryPolicies;
     private boolean instanceDiscovery;
     private String dacEnvConfiguredCredential;
+    private boolean enableKubernetesTokenProxy;
 
     private Duration credentialProcessTimeout = Duration.ofSeconds(10);
 
@@ -833,6 +834,15 @@ public String getDACEnvConfiguredCredential() {
         return dacEnvConfiguredCredential;
     }
 
+    public boolean isKubernetesTokenProxyEnabled() {
+        return enableKubernetesTokenProxy;
+    }
+
+    public IdentityClientOptions setEnableKubernetesTokenProxy(boolean enableTokenProxy) {
+        this.enableKubernetesTokenProxy = enableTokenProxy;
+        return this;
+    }
+
     public IdentityClientOptions clone() {
         IdentityClientOptions clone
             = new IdentityClientOptions().setAdditionallyAllowedTenants(this.additionallyAllowedTenants)
@@ -863,7 +873,8 @@ public IdentityClientOptions clone() {
                 .setPerRetryPolicies(this.perRetryPolicies)
                 .setBrowserCustomizationOptions(this.browserCustomizationOptions)
                 .setChained(this.isChained)
-                .subscription(this.subscription);
+                .subscription(this.subscription)
+                .setEnableKubernetesTokenProxy(this.enableKubernetesTokenProxy);
 
         if (isBrokerEnabled()) {
             clone.setBrokerWindowHandle(this.brokerWindowHandle);

sdk/identity/azure-identity/src/main/java/com/azure/identity/implementation/customtokenproxy/CustomTokenProxyConfiguration.java

@@ -0,0 +1,112 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.identity.implementation.customtokenproxy;
+
+import com.azure.core.util.logging.ClientLogger;
+
+import java.net.URI;
+import java.net.URL;
+import java.nio.charset.StandardCharsets;
+import java.net.URISyntaxException;
+
+import com.azure.core.util.Configuration;
+import com.azure.core.util.CoreUtils;
+
+public final class CustomTokenProxyConfiguration {
+
+    private static final ClientLogger LOGGER = new ClientLogger(CustomTokenProxyConfiguration.class);
+
+    public static final String AZURE_KUBERNETES_TOKEN_PROXY = "AZURE_KUBERNETES_TOKEN_PROXY";
+    public static final String AZURE_KUBERNETES_CA_FILE = "AZURE_KUBERNETES_CA_FILE";
+    public static final String AZURE_KUBERNETES_CA_DATA = "AZURE_KUBERNETES_CA_DATA";
+    public static final String AZURE_KUBERNETES_SNI_NAME = "AZURE_KUBERNETES_SNI_NAME";
+
+    private CustomTokenProxyConfiguration() {
+    }
+
+    public static boolean isConfigured(Configuration configuration) {
+        String tokenProxyUrl = configuration.get(AZURE_KUBERNETES_TOKEN_PROXY);
+        return !CoreUtils.isNullOrEmpty(tokenProxyUrl);
+    }
+
+    public static ProxyConfig parseAndValidate(Configuration configuration) {
+        String tokenProxyUrl = configuration.get(AZURE_KUBERNETES_TOKEN_PROXY);
+        String caFile = configuration.get(AZURE_KUBERNETES_CA_FILE);
+        String caData = configuration.get(AZURE_KUBERNETES_CA_DATA);
+        String sniName = configuration.get(AZURE_KUBERNETES_SNI_NAME);
+
+        if (CoreUtils.isNullOrEmpty(tokenProxyUrl)) {
+            if (!CoreUtils.isNullOrEmpty(sniName)
+                || !CoreUtils.isNullOrEmpty(caFile)
+                || !CoreUtils.isNullOrEmpty(caData)) {
+                throw LOGGER.logExceptionAsError(new IllegalArgumentException(
+                    "AZURE_KUBERNETES_TOKEN_PROXY is not set but other custom endpoint-related environment variables are present"));
+            }
+            return null;
+        }
+
+        if (!CoreUtils.isNullOrEmpty(caFile) && !CoreUtils.isNullOrEmpty(caData)) {
+            throw LOGGER.logExceptionAsError(new IllegalArgumentException(
+                "Only one of AZURE_KUBERNETES_CA_FILE or AZURE_KUBERNETES_CA_DATA can be set."));
+        }
+
+        URL proxyUrl = validateProxyUrl(tokenProxyUrl);
+
+        byte[] caCertBytes = null;
+        if (!CoreUtils.isNullOrEmpty(caData)) {
+            try {
+                caCertBytes = caData.getBytes(StandardCharsets.UTF_8);
+            } catch (Exception e) {
+                throw LOGGER.logExceptionAsError(new IllegalArgumentException(
+                    "Failed to decode CA certificate data from AZURE_KUBERNETES_CA_DATA", e));
+            }
+        }
+
+        ProxyConfig config = new ProxyConfig(proxyUrl, sniName, caFile, caCertBytes);
+        return config;
+    }
+
+    private static URL validateProxyUrl(String endpoint) {
+        if (CoreUtils.isNullOrEmpty(endpoint)) {
+            throw LOGGER.logExceptionAsError(new IllegalArgumentException("Proxy endpoint cannot be null or empty"));
+        }
+
+        try {
+            URI tokenProxy = new URI(endpoint);
+
+            if (!"https".equals(tokenProxy.getScheme())) {
+                throw LOGGER.logExceptionAsError(new IllegalArgumentException(
+                    "Custom token endpoint must use https scheme, got: " + tokenProxy.getScheme()));
+            }
+
+            if (tokenProxy.getRawUserInfo() != null) {
+                throw LOGGER.logExceptionAsError(
+                    new IllegalArgumentException("Custom token endpoint URL must not contain user info: " + endpoint));
+            }
+
+            if (tokenProxy.getRawQuery() != null) {
+                throw LOGGER.logExceptionAsError(
+                    new IllegalArgumentException("Custom token endpoint URL must not contain a query: " + endpoint));
+            }
+
+            if (tokenProxy.getRawFragment() != null) {
+                throw LOGGER.logExceptionAsError(
+                    new IllegalArgumentException("Custom token endpoint URL must not contain a fragment: " + endpoint));
+            }
+
+            if (tokenProxy.getRawPath() == null || tokenProxy.getRawPath().isEmpty()) {
+                tokenProxy = new URI(tokenProxy.getScheme(), null, tokenProxy.getHost(), tokenProxy.getPort(), "/",
+                    null, null);
+            }
+
+            return tokenProxy.toURL();
+
+        } catch (URISyntaxException | IllegalArgumentException e) {
+            throw LOGGER.logExceptionAsError(new IllegalArgumentException("Failed to normalize proxy URL path", e));
+        } catch (Exception e) {
+            throw new RuntimeException("Unexpected error while validating proxy URL: " + endpoint, e);
+        }
+    }
+
+}