Throw CredentialUnavailableException for cases where IMDS returns invalid JSON.

Author: billwert

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

@@ -6,7 +6,9 @@
 import com.azure.core.credential.AccessToken;
 import com.azure.core.credential.TokenRequestContext;
 import com.azure.core.exception.ClientAuthenticationException;
+import com.azure.core.http.HttpMethod;
 import com.azure.core.http.HttpPipeline;
+import com.azure.core.http.HttpRequest;
 import com.azure.core.http.ProxyOptions;
 import com.azure.core.util.CoreUtils;
 import com.azure.core.util.SharedExecutorService;
@@ -33,6 +35,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;
@@ -43,7 +46,6 @@
 import reactor.core.publisher.Mono;
 
 import java.io.IOException;
-import java.net.HttpURLConnection;
 import java.net.MalformedURLException;
 import java.net.Proxy;
 import java.net.Proxy.Type;
@@ -553,7 +555,12 @@ 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 is not available.", null, t);
+            })
             .map(MsalToken::new);
     }
 
@@ -913,27 +920,18 @@ int getRetryTimeoutInMs(int retry) {
     }
 
     private Mono<Boolean> checkIMDSAvailable(String endpoint) {
-        return Mono.fromCallable(() -> {
-            HttpURLConnection connection = null;
-            URL url = getUrl(endpoint + "?api-version=2018-02-01");
-
-            try {
-                connection = (HttpURLConnection) url.openConnection();
-                connection.setRequestMethod("GET");
-                connection.setConnectTimeout(1000);
-                connection.connect();
-            } catch (Exception e) {
-                throw LoggingUtil.logCredentialUnavailableException(LOGGER, options,
-                    new CredentialUnavailableException("ManagedIdentityCredential authentication unavailable. "
-                        + "Connection to IMDS endpoint cannot be established, " + e.getMessage() + ".", e));
-            } finally {
-                if (connection != null) {
-                    connection.disconnect();
-                }
-            }
+        URL url = null;
+        try {
+            url = getUrl(endpoint + "?api-version=2018-02-01");
+        } catch (MalformedURLException e) {
+            return Mono.error(e);
+        }
+        HttpRequest request = new HttpRequest(HttpMethod.GET, url);
 
-            return true;
-        });
+        return getPipeline().send(request)
+                .onErrorMap(t -> new CredentialUnavailableException("ManagedIdentityCredential authentication unavailable. "
+                        + "Connection to IMDS endpoint cannot be established, " + t.getMessage() + ".", t))
+                .flatMap(response -> Mono.just(true));
     }
 
     private static Proxy proxyOptionsToJavaNetProxy(ProxyOptions options) {

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,40 @@ 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
@@ -129,19 +170,19 @@ 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());
     }
 
@@ -156,4 +197,8 @@ public void testAksExchangeTokenCredentialCreated() {
         assertThat("Received class " + cred.managedIdentityServiceCredential.getClass().toString(),
             cred.managedIdentityServiceCredential, instanceOf(AksExchangeTokenCredential.class));
     }
+
+    private MockHttpResponse getMockResponse(int code, String body) {
+        return new MockHttpResponse(null, code, body.getBytes(StandardCharsets.UTF_8));
+    }
 }