Enable Custom Proxy in WorkloadIdentityCredential

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,15 @@ public class WorkloadIdentityCredential implements TokenCredential {
         ClientAssertionCredential tempClientAssertionCredential = null;
         String tempClientId = null;
 
+        if (identityClientOptions.isKubernetesTokenProxyEnabled()) {
+            if (!CustomTokenProxyConfiguration.isConfigured(configuration)) {
+                throw LOGGER.logExceptionAsError(
+                    new IllegalArgumentException("Kubernetes token proxy is enabled but not configured."));
+            }
+            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,11 @@ public WorkloadIdentityCredentialBuilder tokenFilePath(String tokenFilePath) {
         return this;
     }
 
+    public WorkloadIdentityCredentialBuilder enableKubernetesTokenProxy(boolean enable) {
+        this.enableTokenProxy = enable;
+        return this;
+    }
+
     /**
      * Creates new {@link WorkloadIdentityCredential} with the configured options set.
      *
@@ -88,6 +94,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,118 @@
+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 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 IllegalStateException(
+                    "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 IllegalStateException(
+                "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));
+            }
+        }
+
+        return new ProxyConfig(proxyUrl, sniName, caFile, caCertBytes);
+    }
+
+    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);
+        }
+    }
+
+    // public static String getTokenProxyUrl() {
+    //     String tokenProxyUrl = System.getenv(AZURE_KUBERNETES_TOKEN_PROXY);
+    //     if (tokenProxyUrl == null || tokenProxyUrl.isEmpty()) {
+    //         throw LOGGER.logExceptionAsError(new IllegalStateException(
+    //             String.format("Environment variable '%s' is not set or is empty. It must be set to the URL of the"
+    //                 + " token proxy.", AZURE_KUBERNETES_TOKEN_PROXY)));
+    //     }
+    //     return tokenProxyUrl;
+    // }
+
+}