[Identity] Add support for Bridge to Kubernetes to ManagedIdentityCredential (Azure#22584)

* Add IMDS override for Bridge to K8s

* Add to change log

Author: jianghaolu

sdk/core/azure-core/src/main/java/com/azure/core/util/Configuration.java

@@ -89,6 +89,11 @@ public class Configuration implements Cloneable {
     public static final String PROPERTY_AZURE_IDENTITY_DISABLE_CP1 = "AZURE_IDENTITY_DISABLE_CP1";
 
     /**
+     * URL used by Bridge To Kubernetes to redirect IMDS calls in the development environment.
+     */
+    public static final String PROPERTY_AZURE_POD_IDENTITY_TOKEN_URL = "AZURE_POD_IDENTITY_TOKEN_URL";
+
+    /*
      * Name of Azure AAD regional authority.
      */
     public static final String PROPERTY_AZURE_REGIONAL_AUTHORITY_NAME = "AZURE_REGIONAL_AUTHORITY_NAME";

sdk/identity/azure-identity/CHANGELOG.md

@@ -3,6 +3,7 @@
 ## 1.4.0-beta.1 (Unreleased)
 ### Features Added
 
+- Added support to `ManagedIdentityCredential` for Bridge to Kubernetes local development authentication.
 - Added regional STS support to client credential types.
     - Added the `RegionalAuthority` type, that allows specifying Azure regions.
     - Added `regionalAuthority()` setter to `ClientSecretCredentialBuilder` and `ClientCertificateCredentialBuilder`.

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

@@ -15,6 +15,7 @@
 import com.azure.core.http.policy.HttpPipelinePolicy;
 import com.azure.core.http.policy.HttpPolicyProviders;
 import com.azure.core.http.policy.RetryPolicy;
+import com.azure.core.util.Configuration;
 import com.azure.core.util.CoreUtils;
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.core.util.serializer.JacksonAdapter;
@@ -25,6 +26,7 @@
 import com.azure.identity.RegionalAuthority;
 import com.azure.identity.TokenCachePersistenceOptions;
 import com.azure.identity.implementation.util.CertificateUtil;
+import com.azure.identity.implementation.util.IdentityConstants;
 import com.azure.identity.implementation.util.IdentitySslUtil;
 import com.azure.identity.implementation.util.ScopeUtil;
 import com.fasterxml.jackson.databind.JsonNode;
@@ -1114,15 +1116,18 @@ public Mono<AccessToken> authenticateToIMDSEndpoint(TokenRequestContext request)
             return Mono.error(exception);
         }
 
-        return checkIMDSAvailable().flatMap(available -> Mono.fromCallable(() -> {
+        String endpoint = Configuration.getGlobalConfiguration().get(
+            Configuration.PROPERTY_AZURE_POD_IDENTITY_TOKEN_URL,
+            IdentityConstants.DEFAULT_IMDS_ENDPOINT);
+
+        return checkIMDSAvailable(endpoint).flatMap(available -> Mono.fromCallable(() -> {
             int retry = 1;
             while (retry <= options.getMaxRetry()) {
                 URL url = null;
                 HttpURLConnection connection = null;
                 try {
                     url =
-                            new URL(String.format("http://169.254.169.254/metadata/identity/oauth2/token?%s",
-                                    payload.toString()));
+                            new URL(String.format("%s?%s", endpoint, payload.toString()));
 
                     connection = (HttpURLConnection) url.openConnection();
                     connection.setRequestMethod("GET");
@@ -1189,7 +1194,7 @@ public Mono<AccessToken> authenticateToIMDSEndpoint(TokenRequestContext request)
         }));
     }
 
-    private Mono<Boolean> checkIMDSAvailable() {
+    private Mono<Boolean> checkIMDSAvailable(String endpoint) {
         StringBuilder payload = new StringBuilder();
 
         try {
@@ -1200,8 +1205,7 @@ private Mono<Boolean> checkIMDSAvailable() {
         }
         return Mono.fromCallable(() -> {
             HttpURLConnection connection = null;
-            URL url = new URL(String.format("http://169.254.169.254/metadata/identity/oauth2/token?%s",
-                            payload.toString()));
+            URL url = new URL(String.format("%s?%s", endpoint, payload.toString()));
 
             try {
                 connection = (HttpURLConnection) url.openConnection();

sdk/identity/azure-identity/src/main/java/com/azure/identity/implementation/util/IdentityConstants.java

@@ -8,4 +8,9 @@ public class IdentityConstants {
      * The Single Sign On ID.
      */
     public static final String DEVELOPER_SINGLE_SIGN_ON_ID = "04b07795-8ddb-461a-bbee-02f9e1bf7b46";
+
+    /**
+     * The default IMDS authentication endpoint.
+     */
+    public static final String DEFAULT_IMDS_ENDPOINT = "http://169.254.169.254/metadata/identity/oauth2/token";
 }

sdk/identity/azure-identity/src/test/java/com/azure/identity/implementation/IdentityClientTests.java

@@ -9,6 +9,7 @@
 import com.azure.core.util.Configuration;
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.identity.implementation.util.CertificateUtil;
+import com.azure.identity.implementation.util.IdentityConstants;
 import com.azure.identity.util.TestUtils;
 import com.microsoft.aad.msal4j.AuthorizationCodeParameters;
 import com.microsoft.aad.msal4j.ClientCredentialFactory;
@@ -306,7 +307,31 @@ public void testValidIMDSCodeFlow() throws Exception {
 
 
         // mock
-        mockForIMDSCodeFlow(tokenJson);
+        mockForIMDSCodeFlow(IdentityConstants.DEFAULT_IMDS_ENDPOINT, tokenJson);
+
+        // mockForDeviceCodeFlow(request, accessToken, expiresOn);
+
+        // test
+        IdentityClient client = new IdentityClientBuilder().build();
+        AccessToken token = client.authenticateToIMDSEndpoint(request).block();
+        Assert.assertEquals("token1", token.getToken());
+        Assert.assertEquals(expiresOn.getSecond(), token.getExpiresAt().getSecond());
+    }
+
+    @Test
+    public void testCustomIMDSCodeFlow() throws Exception {
+        // setup
+        Configuration configuration = Configuration.getGlobalConfiguration();
+        String endpoint = "http://awesome.pod.url";
+        TokenRequestContext request = new TokenRequestContext().addScopes("https://management.azure.com");
+        OffsetDateTime expiresOn = OffsetDateTime.now(ZoneOffset.UTC).plusHours(1);
+        configuration.put(Configuration.PROPERTY_AZURE_POD_IDENTITY_TOKEN_URL, endpoint);
+        DateTimeFormatter dtf = DateTimeFormatter.ofPattern("M/d/yyyy H:mm:ss XXX");
+        String tokenJson = "{ \"access_token\" : \"token1\", \"expires_on\" : \"" + expiresOn.format(dtf) + "\" }";
+
+
+        // mock
+        mockForIMDSCodeFlow(endpoint, tokenJson);
 
         // mockForDeviceCodeFlow(request, accessToken, expiresOn);
 
@@ -570,8 +595,9 @@ private void mockForArcCodeFlow(int responseCode) throws Exception {
         when(initConnection.getResponseCode()).thenReturn(responseCode);
     }
 
-    private void mockForIMDSCodeFlow(String tokenJson) throws Exception {
+    private void mockForIMDSCodeFlow(String endpoint, String tokenJson) throws Exception {
         URL u = PowerMockito.mock(URL.class);
+        whenNew(URL.class).withArguments(ArgumentMatchers.startsWith(endpoint)).thenReturn(u);
         whenNew(URL.class).withAnyArguments().thenReturn(u);
         HttpURLConnection huc = PowerMockito.mock(HttpURLConnection.class);
         when(u.openConnection()).thenReturn(huc);