Identity mi dac fixes (Azure#45101)

billwert · web-flow · commit 894008bd36d2 · 2025-05-05T15:56:59.000-07:00
* code cleanup: remove some old MI tests which no longer are necessary and a couple constants.

* helper for getting an HttpClient that will return a sequence of responses.

* Throw CredentialUnavailableException for cases where IMDS returns invalid JSON.

* spotless

* fix test to expect new exception

* use a logger instead of println

* fix ordering

* wip

* wip

* Remove an opportunistic change I was trying to make to move away from HttpURLConnection. There's too many details of how precisely to do this to work out for the moment.

* fix error message

* remove currently unused test helper method.

* spotless
diff --git a/sdk/identity/azure-identity/src/main/java/com/azure/identity/ManagedIdentityCredential.java b/sdk/identity/azure-identity/src/main/java/com/azure/identity/ManagedIdentityCredential.java
@@ -87,12 +87,9 @@ public final class ManagedIdentityCredential implements TokenCredential {
     final ManagedIdentityServiceCredential managedIdentityServiceCredential;
     private final IdentityClientOptions identityClientOptions;
     private final String managedIdentityId;
-    static final String PROPERTY_IMDS_ENDPOINT = "IMDS_ENDPOINT";
     static final String PROPERTY_IDENTITY_SERVER_THUMBPRINT = "IDENTITY_SERVER_THUMBPRINT";
     static final String AZURE_FEDERATED_TOKEN_FILE = "AZURE_FEDERATED_TOKEN_FILE";
 
-    static final String USE_AZURE_IDENTITY_CLIENT_LIBRARY_LEGACY_MI = "USE_AZURE_IDENTITY_CLIENT_LIBRARY_LEGACY_MI";
-
     /**
      * Creates an instance of the ManagedIdentityCredential with the client ID of a
      * user-assigned identity, or app registration (when working with AKS pod-identity).
diff --git a/sdk/identity/azure-identity/src/main/java/com/azure/identity/implementation/IdentityClient.java b/sdk/identity/azure-identity/src/main/java/com/azure/identity/implementation/IdentityClient.java
@@ -33,6 +33,7 @@
 import com.microsoft.aad.msal4j.ManagedIdentityApplication;
 import com.microsoft.aad.msal4j.ManagedIdentitySourceType;
 import com.microsoft.aad.msal4j.MsalInteractionRequiredException;
+import com.microsoft.aad.msal4j.MsalJsonParsingException;
 import com.microsoft.aad.msal4j.PublicClientApplication;
 import com.microsoft.aad.msal4j.RefreshTokenParameters;
 import com.microsoft.aad.msal4j.SilentParameters;
@@ -553,7 +554,13 @@ private Mono<AccessToken> getTokenFromMsalMIClient(String resource) {
                     throw new RuntimeException(e);
                 }
             }))
-            .onErrorMap(t -> new CredentialUnavailableException("Managed Identity authentication is not available.", t))
+            .onErrorMap(t -> {
+                if (options.isChained() && t instanceof MsalJsonParsingException) {
+                    return new CredentialUnavailableException("Managed Identity authentication is not available.", t);
+                }
+                return new ClientAuthenticationException(
+                    "Managed Identity authentication failed, see inner exception" + " for more information.", null, t);
+            })
             .map(MsalToken::new);
     }
 
diff --git a/sdk/identity/azure-identity/src/test/java/LiveManagedIdentityTests.java b/sdk/identity/azure-identity/src/test/java/LiveManagedIdentityTests.java
@@ -29,6 +29,7 @@
 import static org.junit.jupiter.api.Assertions.fail;
 
 public class LiveManagedIdentityTests {
+    private static final ClientLogger LOGGER = new ClientLogger(LiveManagedIdentityTests.class);
 
     @Test
     @EnabledIfEnvironmentVariable(named = "AZURE_TEST_MODE", matches = "LIVE")
@@ -67,9 +68,9 @@ public void testManagedIdentityWebAppDeployment() {
     @Timeout(value = 15, unit = TimeUnit.MINUTES)
     @Retry(maxRetries = 3)
     public void testManagedIdentityAksDeployment() {
-        System.out.println("Environment: " + System.getenv("IDENTITY_ENVIRONMENT"));
+        LOGGER.log(LogLevel.INFORMATIONAL, () -> "Environment: " + System.getenv("IDENTITY_ENVIRONMENT"));
         String os = System.getProperty("os.name");
-        System.out.println("OS: " + os);
+        LOGGER.log(LogLevel.INFORMATIONAL, () -> "OS: " + os);
         //Setup Env
         Configuration configuration = Configuration.getGlobalConfiguration().clone();
 
@@ -82,6 +83,7 @@ public void testManagedIdentityAksDeployment() {
         assertTrue(podOutput.contains(podName), "Pod name not found in the output");
 
         String output = runCommand(kubectlPath, "exec", "-it", podName, "--", "java", "-jar", "/identity-test.jar");
+
         assertTrue(output.contains("Successfully retrieved managed identity tokens"),
             "Failed to get response from AKS");
     }
@@ -93,9 +95,9 @@ public void testManagedIdentityAksDeployment() {
     @Timeout(value = 15, unit = TimeUnit.MINUTES)
     @Retry(maxRetries = 3)
     public void testManagedIdentityVmDeployment() {
-        System.out.println("Environment: " + System.getenv("IDENTITY_ENVIRONMENT"));
+        LOGGER.log(LogLevel.INFORMATIONAL, () -> "Environment: " + System.getenv("IDENTITY_ENVIRONMENT"));
         String os = System.getProperty("os.name");
-        System.out.println("OS: " + os);
+        LOGGER.log(LogLevel.INFORMATIONAL, () -> "OS: " + os);
         //Setup Env
         Configuration configuration = Configuration.getGlobalConfiguration().clone();
 
@@ -128,7 +130,7 @@ public void testManagedIdentityVmDeployment() {
             storageAcccountName, sasToken);
         String script = String.format("curl \'%s\' -o ./testfile.jar && java -jar ./testfile.jar", vmBlob);
 
-        System.out.println("Script: " + script);
+        LOGGER.log(LogLevel.INFORMATIONAL, () -> "Script: " + script);
 
         String output = runCommand(azPath, "vm", "run-command", "invoke", "-n", vmName, "-g", resourceGroup,
             "--command-id", "RunShellScript", "--scripts", script);
@@ -163,7 +165,7 @@ private String runCommand(String... args) {
             for (String arg : args) {
                 command.append(arg).append(" ");
             }
-            System.out.println("Running command: " + command);
+            LOGGER.log(LogLevel.INFORMATIONAL, () -> "Running command: " + command);
             ProcessBuilder processBuilder = new ProcessBuilder(args);
             processBuilder.redirectErrorStream(true);
             Process process = processBuilder.start();
@@ -176,7 +178,7 @@ private String runCommand(String... args) {
                     output.append(line).append(System.lineSeparator());
                 }
             }
-            System.out.println("Output:" + System.lineSeparator() + output);
+            LOGGER.log(LogLevel.INFORMATIONAL, () -> "Output:" + System.lineSeparator() + output);
             return output.toString();
         } catch (IOException | InterruptedException e) {
             e.printStackTrace();
diff --git a/sdk/identity/azure-identity/src/test/java/com/azure/identity/ManagedIdentityCredentialTest.java b/sdk/identity/azure-identity/src/test/java/com/azure/identity/ManagedIdentityCredentialTest.java
@@ -5,21 +5,29 @@
 
 import com.azure.core.credential.TokenRequestContext;
 import com.azure.core.exception.ClientAuthenticationException;
+import com.azure.core.http.HttpClient;
+import com.azure.core.test.http.MockHttpResponse;
 import com.azure.core.test.utils.TestConfigurationSource;
 import com.azure.core.util.Configuration;
 import com.azure.identity.implementation.IdentityClient;
+import com.azure.identity.implementation.IdentityClientOptions;
 import com.azure.identity.util.TestUtils;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.MockedConstruction;
 import reactor.test.StepVerifier;
 
+import java.nio.charset.StandardCharsets;
 import java.time.OffsetDateTime;
 import java.time.ZoneOffset;
 import java.util.UUID;
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.core.IsInstanceOf.instanceOf;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.Mockito.mockConstruction;
 import static org.mockito.Mockito.when;
 
@@ -31,7 +39,39 @@ public class ManagedIdentityCredentialTest {
     @Test
     public void testVirtualMachineMSICredentialConfigurations() {
         ManagedIdentityCredential credential = new ManagedIdentityCredentialBuilder().clientId("foo").build();
-        Assertions.assertEquals("foo", credential.getClientId());
+        assertEquals("foo", credential.getClientId());
+    }
+
+    @ParameterizedTest
+    @ValueSource(booleans = { true, false })
+    public void testInvalidJsonResponse(boolean isChained) {
+        HttpClient client = TestUtils.getMockHttpClient(
+            getMockResponse(400,
+                "{\"error\":\"invalid_request\",\"error_description\":\"Required metadata header not specified\"}"),
+            getMockResponse(200, "invalid json"));
+
+        String endpoint = "http://localhost";
+        String secret = "secret";
+
+        Configuration configuration
+            = TestUtils.createTestConfiguration(new TestConfigurationSource().put("MSI_ENDPOINT", endpoint) // This must stay to signal we are in an app service context
+                .put("MSI_SECRET", secret)
+                .put("IDENTITY_ENDPOINT", endpoint)
+                .put("IDENTITY_HEADER", secret));
+
+        IdentityClientOptions options
+            = new IdentityClientOptions().setChained(isChained).setHttpClient(client).setConfiguration(configuration);
+        ManagedIdentityCredential cred = new ManagedIdentityCredential("clientId", null, null, options);
+        StepVerifier.create(cred.getToken(new TokenRequestContext().addScopes("https://management.azure.com")))
+            .expectErrorMatches(t -> {
+                if (isChained) {
+                    return t instanceof CredentialUnavailableException;
+                } else {
+                    return t instanceof ClientAuthenticationException;
+                }
+            })
+            .verify();
+
     }
 
     @Test
@@ -117,7 +157,7 @@ public void testCloudshellUserAssigned() {
         ManagedIdentityCredential credential
             = new ManagedIdentityCredentialBuilder().configuration(configuration).objectId(OBJECT_ID).build();
         StepVerifier.create(credential.getToken(request))
-            .expectErrorMatches(t -> t instanceof CredentialUnavailableException)
+            .expectErrorMatches(t -> t instanceof ClientAuthenticationException)
             .verify();
     }
 
@@ -129,82 +169,22 @@ public void testInvalidIdCombination() {
         String objectId = "2323-sd2323s-32323-32334-34343";
 
         // test
-        Assertions.assertThrows(IllegalStateException.class,
+        assertThrows(IllegalStateException.class,
             () -> new ManagedIdentityCredentialBuilder().clientId(CLIENT_ID).resourceId(resourceId).build());
 
-        Assertions.assertThrows(IllegalStateException.class,
+        assertThrows(IllegalStateException.class,
             () -> new ManagedIdentityCredentialBuilder().clientId(CLIENT_ID)
                 .resourceId(resourceId)
                 .objectId(objectId)
                 .build());
 
-        Assertions.assertThrows(IllegalStateException.class,
+        assertThrows(IllegalStateException.class,
             () -> new ManagedIdentityCredentialBuilder().clientId(CLIENT_ID).objectId(objectId).build());
 
-        Assertions.assertThrows(IllegalStateException.class,
+        assertThrows(IllegalStateException.class,
             () -> new ManagedIdentityCredentialBuilder().resourceId(resourceId).objectId(objectId).build());
     }
 
-    @Test
-    public void testArcIdentityCredentialCreated() {
-        Configuration configuration
-            = TestUtils
-                .createTestConfiguration(new TestConfigurationSource().put("IDENTITY_ENDPOINT", "http://localhost")
-                    .put("IMDS_ENDPOINT", "http://localhost"))
-                .put("USE_AZURE_IDENTITY_CLIENT_LIBRARY_LEGACY_MI", "true");
-
-        ManagedIdentityCredential cred = new ManagedIdentityCredentialBuilder().configuration(configuration).build();
-        assertThat("Received class " + cred.managedIdentityServiceCredential.getClass().toString(),
-            cred.managedIdentityServiceCredential, instanceOf(ManagedIdentityMsalCredential.class));
-    }
-
-    @Test
-    public void testServiceFabricMsiCredentialCreated() {
-        Configuration configuration = TestUtils
-            .createTestConfiguration(new TestConfigurationSource().put("IDENTITY_ENDPOINT", "http://localhost")
-                .put("IDENTITY_SERVER_THUMBPRINT", "thumbprint")
-                .put("IDENTITY_HEADER", "header"))
-            .put("USE_AZURE_IDENTITY_CLIENT_LIBRARY_LEGACY_MI", "true");
-
-        ManagedIdentityCredential cred = new ManagedIdentityCredentialBuilder().configuration(configuration).build();
-        assertThat("Received class " + cred.managedIdentityServiceCredential.getClass().toString(),
-            cred.managedIdentityServiceCredential, instanceOf(ManagedIdentityMsalCredential.class));
-    }
-
-    @Test
-    public void testAppServiceMsi2019CredentialCreated() {
-        Configuration configuration = TestUtils.createTestConfiguration(
-            new TestConfigurationSource().put("IDENTITY_ENDPOINT", "http://localhost").put("IDENTITY_HEADER", "header"))
-            .put("USE_AZURE_IDENTITY_CLIENT_LIBRARY_LEGACY_MI", "true");
-
-        ManagedIdentityCredential cred = new ManagedIdentityCredentialBuilder().configuration(configuration).build();
-        assertThat("Received class " + cred.managedIdentityServiceCredential.getClass().toString(),
-            cred.managedIdentityServiceCredential, instanceOf(ManagedIdentityMsalCredential.class));
-    }
-
-    @Test
-    public void testAppServiceMsi2017CredentialCreated() {
-        Configuration configuration = TestUtils
-            .createTestConfiguration(
-                new TestConfigurationSource().put("MSI_ENDPOINT", "http://localhost").put("MSI_SECRET", "secret"))
-            .put("USE_AZURE_IDENTITY_CLIENT_LIBRARY_LEGACY_MI", "true");
-
-        ManagedIdentityCredential cred = new ManagedIdentityCredentialBuilder().configuration(configuration).build();
-        assertThat("Received class " + cred.managedIdentityServiceCredential.getClass().toString(),
-            cred.managedIdentityServiceCredential, instanceOf(ManagedIdentityMsalCredential.class));
-    }
-
-    @Test
-    public void testCloudShellCredentialCreated() {
-        Configuration configuration
-            = TestUtils.createTestConfiguration(new TestConfigurationSource().put("MSI_ENDPOINT", "http://localhost"))
-                .put("USE_AZURE_IDENTITY_CLIENT_LIBRARY_LEGACY_MI", "true");
-
-        ManagedIdentityCredential cred = new ManagedIdentityCredentialBuilder().configuration(configuration).build();
-        assertThat("Received class " + cred.managedIdentityServiceCredential.getClass().toString(),
-            cred.managedIdentityServiceCredential, instanceOf(ManagedIdentityMsalCredential.class));
-    }
-
     @Test
     public void testAksExchangeTokenCredentialCreated() {
         Configuration configuration
@@ -217,13 +197,7 @@ public void testAksExchangeTokenCredentialCreated() {
             cred.managedIdentityServiceCredential, instanceOf(AksExchangeTokenCredential.class));
     }
 
-    @Test
-    public void testIMDSPodIdentityV1Credential() {
-        Configuration configuration = TestUtils.createTestConfiguration(new TestConfigurationSource());
-        configuration.put("USE_AZURE_IDENTITY_CLIENT_LIBRARY_LEGACY_MI", "true");
-
-        ManagedIdentityCredential cred = new ManagedIdentityCredentialBuilder().configuration(configuration).build();
-        assertThat("Received class " + cred.managedIdentityServiceCredential.getClass().toString(),
-            cred.managedIdentityServiceCredential, instanceOf(ManagedIdentityMsalCredential.class));
+    private MockHttpResponse getMockResponse(int code, String body) {
+        return new MockHttpResponse(null, code, body.getBytes(StandardCharsets.UTF_8));
     }
 }
diff --git a/sdk/identity/azure-identity/src/test/java/com/azure/identity/util/TestUtils.java b/sdk/identity/azure-identity/src/test/java/com/azure/identity/util/TestUtils.java
@@ -4,6 +4,9 @@
 package com.azure.identity.util;
 
 import com.azure.core.credential.AccessToken;
+import com.azure.core.http.HttpClient;
+import com.azure.core.http.HttpRequest;
+import com.azure.core.http.HttpResponse;
 import com.azure.core.util.Configuration;
 import com.azure.core.util.ConfigurationBuilder;
 import com.azure.core.util.ConfigurationSource;
@@ -147,6 +150,26 @@ public static Mono<AccessToken> getMockAccessToken(String accessToken, OffsetDat
         return Mono.just(new AccessToken(accessToken, expiresOn.plusMinutes(2).minus(tokenRefreshOffset)));
     }
 
+    /**
+     * Creates a mock {@link HttpClient} which simply returns the provided responses in order.
+     * @param responses The responses to return.
+     * @param delay If specified, throw a {@code RuntimeException} after the specified delay.
+     * @return A mock HttpClient that returns the provided responses.
+     */
+    public static HttpClient getMockHttpClient(HttpResponse... responses) {
+        return new HttpClient() {
+            int index = 0;
+
+            @Override
+            public Mono<HttpResponse> send(HttpRequest request) {
+                if (index >= responses.length) {
+                    throw new IllegalStateException("No more responses available");
+                }
+                return Mono.just(responses[index++]);
+            }
+        };
+    }
+
     /**
      * Creates a {@link Configuration} with the specified {@link ConfigurationSource} as the only source of
      * configurations.
