Add Token Caching Support For Managed Identity

sdk/identity/azure-identity/pom.xml

@@ -43,7 +43,7 @@
     <dependency>
       <groupId>com.microsoft.azure</groupId>
       <artifactId>msal4j</artifactId>
-      <version>1.12.0</version> <!-- {x-version-update;com.microsoft.azure:msal4j;external_dependency} -->
+      <version>1.13.0</version> <!-- {x-version-update;com.microsoft.azure:msal4j;external_dependency} -->
     </dependency>
     <dependency>
       <groupId>com.microsoft.azure</groupId>
@@ -104,7 +104,7 @@
           <rules>
             <bannedDependencies>
               <includes>
-                <include>com.microsoft.azure:msal4j:[1.12.0]</include> <!-- {x-include-update;com.microsoft.azure:msal4j;external_dependency} -->
+                <include>com.microsoft.azure:msal4j:[1.13.0]</include> <!-- {x-include-update;com.microsoft.azure:msal4j;external_dependency} -->
                 <include>com.microsoft.azure:msal4j-persistence-extension:[1.1.0]</include> <!-- {x-include-update;com.microsoft.azure:msal4j-persistence-extension;external_dependency} -->
                 <include>net.java.dev.jna:jna-platform:[5.6.0]</include> <!-- {x-include-update;net.java.dev.jna:jna-platform;external_dependency} -->
                 <include>org.linguafranca.pwdb:KeePassJava2:[2.1.4]</include> <!-- {x-include-update;org.linguafranca.pwdb:KeePassJava2;external_dependency} -->

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

@@ -33,6 +33,6 @@ public Mono<AccessToken> authenticate(TokenRequestContext request) {
                 + " 'AZURE_CLIENT_ID' environment variable or through the credential builder."
                 + " Please ensure client id is provided to authenticate via token exchange in AKS environment.")));
         }
-        return identityClient.authenticateWithExchangeToken(request);
+        return identityClient.authenticateWithManagedIdentityConfidentialClient(request);
     }
 }

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

@@ -51,8 +51,6 @@ class AppServiceMsiCredential extends ManagedIdentityServiceCredential {
      * @return A publisher that emits an {@link AccessToken}.
      */
     public Mono<AccessToken> authenticate(TokenRequestContext request) {
-        return identityClient.authenticateToManagedIdentityEndpoint(identityEndpoint, identityHeader,
-            msiEndpoint, msiSecret,
-            request);
+        return identityClient.authenticateWithManagedIdentityConfidentialClient(request);
     }
 }

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

@@ -49,6 +49,6 @@ public Mono<AccessToken> authenticate(TokenRequestContext request) {
                     + "with the system assigned identity omit the client id when constructing the"
                     + " ManagedIdentityCredential.", null)));
         }
-        return identityClient.authenticateToArcManagedIdentityEndpoint(identityEndpoint, request);
+        return identityClient.authenticateWithManagedIdentityConfidentialClient(request);
     }
 }

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

@@ -11,6 +11,8 @@
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.identity.implementation.IdentityClientBuilder;
 import com.azure.identity.implementation.IdentityClientOptions;
+import com.azure.identity.implementation.ManagedIdentityParameters;
+import com.azure.identity.implementation.ManagedIdentityType;
 import com.azure.identity.implementation.util.LoggingUtil;
 import reactor.core.publisher.Mono;
 
@@ -48,7 +50,6 @@ public final class ManagedIdentityCredential implements TokenCredential {
         Configuration configuration = identityClientOptions.getConfiguration() == null
             ? Configuration.getGlobalConfiguration().clone() : identityClientOptions.getConfiguration();
 
-
         /*
          * Choose credential based on available environment variables in this order:
          *
@@ -62,18 +63,33 @@ public final class ManagedIdentityCredential implements TokenCredential {
          */
 
         if (configuration.contains(Configuration.PROPERTY_MSI_ENDPOINT)) {
-            managedIdentityServiceCredential = new AppServiceMsiCredential(clientId, clientBuilder.build());
+            managedIdentityServiceCredential = new AppServiceMsiCredential(clientId, clientBuilder
+                .identityClientOptions(updateIdentityClientOptions(ManagedIdentityType.APP_SERVICE,
+                    identityClientOptions, configuration))
+                .build());
         } else if (configuration.contains(Configuration.PROPERTY_IDENTITY_ENDPOINT)) {
             if (configuration.contains(Configuration.PROPERTY_IDENTITY_HEADER)) {
                 if (configuration.get(PROPERTY_IDENTITY_SERVER_THUMBPRINT) != null) {
-                    managedIdentityServiceCredential = new ServiceFabricMsiCredential(clientId, clientBuilder.build());
+                    managedIdentityServiceCredential = new ServiceFabricMsiCredential(clientId, clientBuilder
+                        .identityClientOptions(updateIdentityClientOptions(ManagedIdentityType.SERVICE_FABRIC,
+                            identityClientOptions, configuration))
+                        .build());
                 } else {
-                    managedIdentityServiceCredential = new AppServiceMsiCredential(clientId, clientBuilder.build());
+                    managedIdentityServiceCredential = new AppServiceMsiCredential(clientId, clientBuilder
+                        .identityClientOptions(updateIdentityClientOptions(ManagedIdentityType.APP_SERVICE,
+                            identityClientOptions, configuration))
+                        .build());
                 }
             } else if (configuration.get(PROPERTY_IMDS_ENDPOINT) != null) {
-                managedIdentityServiceCredential = new ArcIdentityCredential(clientId, clientBuilder.build());
+                managedIdentityServiceCredential = new ArcIdentityCredential(clientId, clientBuilder
+                    .identityClientOptions(updateIdentityClientOptions(ManagedIdentityType.ARC,
+                        identityClientOptions, configuration))
+                    .build());
             } else {
-                managedIdentityServiceCredential = new VirtualMachineMsiCredential(clientId, clientBuilder.build());
+                managedIdentityServiceCredential = new VirtualMachineMsiCredential(clientId, clientBuilder
+                    .identityClientOptions(updateIdentityClientOptions(ManagedIdentityType.VM,
+                        identityClientOptions, configuration))
+                    .build());
             }
         } else if (configuration.contains(Configuration.PROPERTY_AZURE_TENANT_ID)
                 && configuration.get(AZURE_FEDERATED_TOKEN_FILE) != null) {
@@ -83,13 +99,51 @@ public final class ManagedIdentityCredential implements TokenCredential {
             clientBuilder.tenantId(configuration.get(Configuration.PROPERTY_AZURE_TENANT_ID));
             clientBuilder.clientAssertionPath(configuration.get(AZURE_FEDERATED_TOKEN_FILE));
             clientBuilder.clientAssertionTimeout(Duration.ofMinutes(5));
-            managedIdentityServiceCredential = new AksExchangeTokenCredential(clientIdentifier, clientBuilder.build());
+            managedIdentityServiceCredential = new AksExchangeTokenCredential(clientIdentifier, clientBuilder
+                .identityClientOptions(updateIdentityClientOptions(ManagedIdentityType.AKS,
+                    identityClientOptions, configuration))
+                .build());
         } else {
-            managedIdentityServiceCredential = new VirtualMachineMsiCredential(clientId, clientBuilder.build());
+            managedIdentityServiceCredential = new VirtualMachineMsiCredential(clientId, clientBuilder
+                .identityClientOptions(updateIdentityClientOptions(ManagedIdentityType.VM,
+                    identityClientOptions, configuration))
+                .build());
         }
         LoggingUtil.logAvailableEnvironmentVariables(LOGGER, configuration);
     }
 
+    private IdentityClientOptions updateIdentityClientOptions(ManagedIdentityType managedIdentityType,
+                                             IdentityClientOptions clientOptions, Configuration configuration) {
+        switch (managedIdentityType) {
+            case APP_SERVICE:
+                return clientOptions
+                    .setManagedIdentityType(ManagedIdentityType.APP_SERVICE)
+                    .setManagedIdentityParameters(new ManagedIdentityParameters()
+                        .setMsiEndpoint(configuration.get(Configuration.PROPERTY_MSI_ENDPOINT))
+                        .setMsiSecret(configuration.get(Configuration.PROPERTY_MSI_SECRET))
+                        .setIdentityEndpoint(configuration.get(Configuration.PROPERTY_IDENTITY_ENDPOINT))
+                        .setIdentityHeader(configuration.get(Configuration.PROPERTY_IDENTITY_HEADER)));
+            case SERVICE_FABRIC:
+                return clientOptions
+                    .setManagedIdentityType(ManagedIdentityType.SERVICE_FABRIC)
+                    .setManagedIdentityParameters(new ManagedIdentityParameters()
+                        .setIdentityServerThumbprint(configuration.get(PROPERTY_IDENTITY_SERVER_THUMBPRINT))
+                        .setIdentityEndpoint(configuration.get(Configuration.PROPERTY_IDENTITY_ENDPOINT))
+                        .setIdentityHeader(configuration.get(Configuration.PROPERTY_IDENTITY_HEADER)));
+            case ARC:
+                return clientOptions
+                    .setManagedIdentityType(ManagedIdentityType.ARC)
+                    .setManagedIdentityParameters(new ManagedIdentityParameters()
+                        .setIdentityEndpoint(configuration.get(Configuration.PROPERTY_IDENTITY_ENDPOINT)));
+            case VM:
+                return clientOptions.setManagedIdentityType(ManagedIdentityType.VM);
+            case AKS:
+                return clientOptions.setManagedIdentityType(ManagedIdentityType.AKS);
+            default:
+                return clientOptions.setManagedIdentityType(ManagedIdentityType.NONE);
+        }
+    }
+
     /**
      * Gets the client ID of user assigned or system assigned identity.
      * @return the client ID of user assigned or system assigned identity.

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

@@ -49,7 +49,6 @@ class ServiceFabricMsiCredential extends ManagedIdentityServiceCredential {
      * @return A publisher that emits an {@link AccessToken}.
      */
     public Mono<AccessToken> authenticate(TokenRequestContext request) {
-        return identityClient.authenticateToServiceFabricManagedIdentityEndpoint(identityEndpoint, identityHeader,
-            identityServerThumbprint, request);
+        return identityClient.authenticateWithManagedIdentityConfidentialClient(request);
     }
 }

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

@@ -31,6 +31,6 @@ class VirtualMachineMsiCredential extends ManagedIdentityServiceCredential {
      * @return A publisher that emits an {@link AccessToken}.
      */
     public Mono<AccessToken> authenticate(TokenRequestContext request) {
-        return identityClient.authenticateToIMDSEndpoint(request);
+        return identityClient.authenticateWithManagedIdentityConfidentialClient(request);
     }
 }

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

@@ -30,23 +30,7 @@
 import com.azure.identity.implementation.util.ScopeUtil;
 import com.azure.identity.implementation.util.LoggingUtil;
 import com.fasterxml.jackson.databind.JsonNode;
-import com.microsoft.aad.msal4j.AuthorizationCodeParameters;
-import com.microsoft.aad.msal4j.ClaimsRequest;
-import com.microsoft.aad.msal4j.ClientCredentialFactory;
-import com.microsoft.aad.msal4j.ClientCredentialParameters;
-import com.microsoft.aad.msal4j.ConfidentialClientApplication;
-import com.microsoft.aad.msal4j.DeviceCodeFlowParameters;
-import com.microsoft.aad.msal4j.IAccount;
-import com.microsoft.aad.msal4j.IAuthenticationResult;
-import com.microsoft.aad.msal4j.IClientCredential;
-import com.microsoft.aad.msal4j.InteractiveRequestParameters;
-import com.microsoft.aad.msal4j.MsalInteractionRequiredException;
-import com.microsoft.aad.msal4j.OnBehalfOfParameters;
-import com.microsoft.aad.msal4j.Prompt;
-import com.microsoft.aad.msal4j.PublicClientApplication;
-import com.microsoft.aad.msal4j.RefreshTokenParameters;
-import com.microsoft.aad.msal4j.SilentParameters;
-import com.microsoft.aad.msal4j.UserNamePasswordParameters;
+import com.microsoft.aad.msal4j.*;
 import com.sun.jna.Platform;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -137,6 +121,7 @@ public class IdentityClient {
     private HttpPipelineAdapter httpPipelineAdapter;
     private final SynchronizedAccessor<PublicClientApplication> publicClientApplicationAccessor;
     private final SynchronizedAccessor<ConfidentialClientApplication> confidentialClientApplicationAccessor;
+    private final SynchronizedAccessor<ConfidentialClientApplication> managedIdentityConfidentialClientApplicationAccessor;
     private final SynchronizedAccessor<String> clientAssertionAccessor;
 
 
@@ -182,6 +167,9 @@ public class IdentityClient {
         this.confidentialClientApplicationAccessor = new SynchronizedAccessor<>(() ->
             getConfidentialClientApplication());
 
+        this.managedIdentityConfidentialClientApplicationAccessor = new SynchronizedAccessor<>(() ->
+            getManagedIdentityConfidentialClient());
+
         this.clientAssertionAccessor = clientAssertionTimeout == null
             ? new SynchronizedAccessor<>(() -> parseClientAssertion(), Duration.ofMinutes(5))
                 : new SynchronizedAccessor<>(() -> parseClientAssertion(), clientAssertionTimeout);
@@ -241,6 +229,26 @@ private Mono<ConfidentialClientApplication> getConfidentialClientApplication() {
 
             applicationBuilder.sendX5c(options.isIncludeX5c());
 
+
+            if (!options.getManagedIdentityType().equals(ManagedIdentityType.NONE)) {
+                applicationBuilder.appTokenProvider(appTokenProviderParameters -> {
+                    TokenRequestContext trc = new TokenRequestContext()
+                        .setScopes(new ArrayList<>(appTokenProviderParameters.scopes))
+                        .setClaims(appTokenProviderParameters.claims)
+                        .setTenantId(appTokenProviderParameters.tenantId);
+
+                    Mono<AccessToken> accessTokenMono = getTokenFromTargetManagedIdentity(options.getManagedIdentityType(), trc);
+
+                    return accessTokenMono.toFuture().thenApply(accessToken -> {
+                        TokenProviderResult result =  new TokenProviderResult();
+                        result.setAccessToken(accessToken.getToken());
+                        result.setTenantId(trc.getTenantId());
+                        result.setExpiresInSeconds(accessToken.getExpiresAt().toEpochSecond());
+                        return result;
+                    });
+                });
+            }
+
             initializeHttpPipelineAdapter();
             if (httpPipelineAdapter != null) {
                 applicationBuilder.httpClient(httpPipelineAdapter);
@@ -277,6 +285,74 @@ private Mono<ConfidentialClientApplication> getConfidentialClientApplication() {
         });
     }
 
+    private Mono<ConfidentialClientApplication> getManagedIdentityConfidentialClient() {
+        return Mono.defer(() -> {
+            String authorityUrl = TRAILING_FORWARD_SLASHES.matcher(options.getAuthorityHost()).replaceAll("")
+                + "/" + tenantId;
+
+            // Temporarily pass in Dummy Client secret and Client ID. until MSal removes its requirements.
+            IClientCredential credential = ClientCredentialFactory.createFromSecret("dummy-secret");
+            ConfidentialClientApplication.Builder applicationBuilder =
+                ConfidentialClientApplication.builder(clientId == null ? IdentityConstants.DEVELOPER_SINGLE_SIGN_ON_ID
+                    : clientId, credential);
+            try {
+                applicationBuilder = applicationBuilder.authority(authorityUrl);
+            } catch (MalformedURLException e) {
+                return Mono.error(LOGGER.logExceptionAsWarning(new IllegalStateException(e)));
+            }
+
+            applicationBuilder.appTokenProvider(appTokenProviderParameters -> {
+                TokenRequestContext trc = new TokenRequestContext()
+                    .setScopes(new ArrayList<>(appTokenProviderParameters.scopes))
+                    .setClaims(appTokenProviderParameters.claims)
+                    .setTenantId(appTokenProviderParameters.tenantId);
+
+                Mono<AccessToken> accessTokenMono = getTokenFromTargetManagedIdentity(options.getManagedIdentityType(), trc);
+
+                return accessTokenMono.toFuture().thenApply(accessToken -> {
+                    TokenProviderResult result =  new TokenProviderResult();
+                    result.setAccessToken(accessToken.getToken());
+                    result.setTenantId(trc.getTenantId());
+                    result.setExpiresInSeconds(accessToken.getExpiresAt().toEpochSecond());
+                    return result;
+                });
+            });
+
+
+            initializeHttpPipelineAdapter();
+            if (httpPipelineAdapter != null) {
+                applicationBuilder.httpClient(httpPipelineAdapter);
+            } else {
+                applicationBuilder.proxy(proxyOptionsToJavaNetProxy(options.getProxyOptions()));
+            }
+
+            if (options.getExecutorService() != null) {
+                applicationBuilder.executorService(options.getExecutorService());
+            }
+
+            ConfidentialClientApplication confidentialClientApplication = applicationBuilder.build();
+            return Mono.just(confidentialClientApplication);
+        });
+    }
+
+    private Mono<AccessToken> getTokenFromTargetManagedIdentity(ManagedIdentityType managedIdentityType, TokenRequestContext trc) {
+        ManagedIdentityParameters parameters = options.getManagedIdentityParameters();
+        switch (managedIdentityType) {
+            case APP_SERVICE:
+                return authenticateToManagedIdentityEndpoint(parameters.getIdentityEndpoint(),
+                    parameters.getIdentityHeader(), parameters.getMsiEndpoint(), parameters.getMsiSecret(), trc);
+            case SERVICE_FABRIC:
+                return authenticateToServiceFabricManagedIdentityEndpoint(parameters.getIdentityEndpoint(),
+                    parameters.getIdentityHeader(), parameters.getIdentityServerThumbprint(), trc);
+            case ARC:
+                return authenticateToArcManagedIdentityEndpoint(parameters.getIdentityEndpoint(), trc);
+            case AKS:
+                return authenticateWithExchangeToken(trc);
+            default:
+                return authenticateToIMDSEndpoint(trc);
+        }
+    }
+
     private Mono<String> parseClientAssertion() {
         return Mono.fromCallable(() -> {
             if (clientAssertionFilePath != null) {
@@ -632,7 +708,6 @@ public Mono<AccessToken> authenticateWithOBO(TokenRequestContext request) {
                 .map(MsalToken::new));
     }
 
-
     private Mono<AccessToken> getAccessTokenFromPowerShell(TokenRequestContext request,
                                                            PowershellManager powershellManager) {
         return powershellManager.initSession()
@@ -705,6 +780,18 @@ public Mono<AccessToken> authenticateWithConfidentialClient(TokenRequestContext
         )).map(MsalToken::new);
     }
 
+    public Mono<AccessToken> authenticateWithManagedIdentityConfidentialClient(TokenRequestContext request) {
+        return managedIdentityConfidentialClientApplicationAccessor.getValue()
+            .flatMap(confidentialClient -> Mono.fromFuture(() -> {
+                    ClientCredentialParameters.ClientCredentialParametersBuilder builder =
+                        ClientCredentialParameters.builder(new HashSet<>(request.getScopes()))
+                            .tenant(IdentityUtil
+                                .resolveTenantId(tenantId, request, options));
+                    return confidentialClient.acquireToken(builder.build());
+                }
+            )).map(MsalToken::new);
+    }
+
     private HttpPipeline setupPipeline(HttpClient httpClient) {
         List<HttpPipelinePolicy> policies = new ArrayList<>();
         HttpLogOptions httpLogOptions = new HttpLogOptions();