Persistent service principal token cache for MSAL confidential client

sdk/core/azure-core/src/main/java/com/azure/core/http/policy/BearerTokenAuthenticationPolicy.java

@@ -3,7 +3,6 @@
 
 package com.azure.core.http.policy;
 
-import com.azure.core.credential.SimpleTokenCache;
 import com.azure.core.credential.TokenCredential;
 import com.azure.core.credential.TokenRequestContext;
 import com.azure.core.http.HttpPipelineCallContext;
@@ -22,8 +21,7 @@ public class BearerTokenAuthenticationPolicy implements HttpPipelinePolicy {
     private static final String BEARER = "Bearer";
 
     private final TokenCredential credential;
-    private final String[] scopes;
-    private final SimpleTokenCache cache;
+    private final TokenRequestContext tokenRequestContext;
 
     /**
      * Creates BearerTokenAuthenticationPolicy.
@@ -36,16 +34,15 @@ public BearerTokenAuthenticationPolicy(TokenCredential credential, String... sco
         Objects.requireNonNull(scopes);
         assert scopes.length > 0;
         this.credential = credential;
-        this.scopes = scopes;
-        this.cache = new SimpleTokenCache(() -> credential.getToken(new TokenRequestContext().addScopes(scopes)));
+        this.tokenRequestContext = new TokenRequestContext().addScopes(scopes);
     }
 
     @Override
     public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
         if ("http".equals(context.getHttpRequest().getUrl().getProtocol())) {
             return Mono.error(new RuntimeException("token credentials require a URL using the HTTPS protocol scheme"));
         }
-        return cache.getToken()
+        return credential.getToken(tokenRequestContext)
             .flatMap(token -> {
                 context.getHttpRequest().getHeaders().put(AUTHORIZATION_HEADER, BEARER + " " + token.getToken());
                 return next.process();

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

@@ -52,7 +52,7 @@ public class AuthorizationCodeCredential implements TokenCredential {
     public Mono<AccessToken> getToken(TokenRequestContext request) {
         return Mono.defer(() -> {
             if (cachedToken.get() != null) {
-                return identityClient.authenticateWithMsalAccount(request, cachedToken.get().getAccount())
+                return identityClient.authenticateWithPublicClientCache(request, cachedToken.get().getAccount())
                     .onErrorResume(t -> Mono.empty());
             } else {
                 return Mono.empty();

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

@@ -25,8 +25,6 @@
  */
 @Immutable
 public class ClientCertificateCredential implements TokenCredential {
-    private final String clientCertificate;
-    private final String clientCertificatePassword;
     private final IdentityClient identityClient;
 
     /**
@@ -40,22 +38,19 @@ public class ClientCertificateCredential implements TokenCredential {
     ClientCertificateCredential(String tenantId, String clientId, String certificatePath, String certificatePassword,
                                 IdentityClientOptions identityClientOptions) {
         Objects.requireNonNull(certificatePath, "'certificatePath' cannot be null.");
-        this.clientCertificate = certificatePath;
-        this.clientCertificatePassword = certificatePassword;
-        identityClient =
-            new IdentityClientBuilder()
-                .tenantId(tenantId)
-                .clientId(clientId)
-                .identityClientOptions(identityClientOptions)
-                .build();
+        identityClient = new IdentityClientBuilder()
+            .tenantId(tenantId)
+            .clientId(clientId)
+            .certificatePath(certificatePath)
+            .certificatePassword(certificatePassword)
+            .identityClientOptions(identityClientOptions)
+            .build();
     }
 
     @Override
     public Mono<AccessToken> getToken(TokenRequestContext request) {
-        if (clientCertificatePassword != null) {
-            return identityClient.authenticateWithPfxCertificate(clientCertificate, clientCertificatePassword, request);
-        } else {
-            return identityClient.authenticateWithPemCertificate(clientCertificate, request);
-        }
+        return identityClient.authenticateWithConfidentialClientCache(request)
+            .onErrorResume(t -> Mono.empty())
+            .switchIfEmpty(Mono.defer(() -> identityClient.authenticateWithConfidentialClient(request)));
     }
 }

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

@@ -40,6 +40,18 @@ public ClientCertificateCredentialBuilder pfxCertificate(String certificatePath,
         return this;
     }
 
+    /**
+     * Sets whether to enable using the shared token cache. This is disabled by default.
+     *
+     * @param enabled indicates whether to enable using the shared token cache.
+     *
+     * @return An updated instance of this builder with if the shared token cache enabled specified.
+     */
+    public ClientCertificateCredentialBuilder enablePersistentCache(boolean enabled) {
+        this.identityClientOptions.enablePersistentCache(enabled);
+        return this;
+    }
+
     /**
      * Creates a new {@link ClientCertificateCredential} with the current configurations.
      *

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

@@ -25,8 +25,6 @@
  */
 @Immutable
 public class ClientSecretCredential implements TokenCredential {
-    /* The client secret value. */
-    private final String clientSecret;
     private final IdentityClient identityClient;
 
     /**
@@ -44,13 +42,15 @@ public class ClientSecretCredential implements TokenCredential {
         identityClient = new IdentityClientBuilder()
             .tenantId(tenantId)
             .clientId(clientId)
+            .clientSecret(clientSecret)
             .identityClientOptions(identityClientOptions)
             .build();
-        this.clientSecret = clientSecret;
     }
 
     @Override
     public Mono<AccessToken> getToken(TokenRequestContext request) {
-        return identityClient.authenticateWithClientSecret(clientSecret, request);
+        return identityClient.authenticateWithConfidentialClientCache(request)
+            .onErrorResume(t -> Mono.empty())
+            .switchIfEmpty(Mono.defer(() -> identityClient.authenticateWithConfidentialClient(request)));
     }
 }

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

@@ -25,6 +25,18 @@ public ClientSecretCredentialBuilder clientSecret(String clientSecret) {
         return this;
     }
 
+    /**
+     * Sets whether to enable using the shared token cache. This is disabled by default.
+     *
+     * @param enabled indicates whether to enable using the shared token cache.
+     *
+     * @return An updated instance of this builder with if the shared token cache enabled specified.
+     */
+    public ClientSecretCredentialBuilder enablePersistentCache(boolean enabled) {
+        this.identityClientOptions.enablePersistentCache(enabled);
+        return this;
+    }
+
     /**
      * Creates a new {@link ClientCertificateCredential} with the current configurations.
      *

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

@@ -48,7 +48,7 @@ public class DeviceCodeCredential implements TokenCredential {
     public Mono<AccessToken> getToken(TokenRequestContext request) {
         return Mono.defer(() -> {
             if (cachedToken.get() != null) {
-                return identityClient.authenticateWithMsalAccount(request, cachedToken.get().getAccount())
+                return identityClient.authenticateWithPublicClientCache(request, cachedToken.get().getAccount())
                     .onErrorResume(t -> Mono.empty());
             } else {
                 return Mono.empty();

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

@@ -75,7 +75,7 @@ class IntelliJCredential implements TokenCredential {
     public Mono<AccessToken> getToken(TokenRequestContext request) {
         return Mono.defer(() -> {
             if (cachedToken.get() != null) {
-                return identityClient.authenticateWithMsalAccount(request, cachedToken.get().getAccount())
+                return identityClient.authenticateWithPublicClientCache(request, cachedToken.get().getAccount())
                            .onErrorResume(t -> Mono.empty());
             } else {
                 return Mono.empty();

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

@@ -54,7 +54,7 @@ public class InteractiveBrowserCredential implements TokenCredential {
     public Mono<AccessToken> getToken(TokenRequestContext request) {
         return Mono.defer(() -> {
             if (cachedToken.get() != null) {
-                return identityClient.authenticateWithMsalAccount(request, cachedToken.get().getAccount())
+                return identityClient.authenticateWithPublicClientCache(request, cachedToken.get().getAccount())
                     .onErrorResume(t -> Mono.empty());
             } else {
                 return Mono.empty();

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

@@ -70,7 +70,7 @@ public class SharedTokenCacheCredential implements TokenCredential {
     public Mono<AccessToken> getToken(TokenRequestContext request) {
         return Mono.defer(() -> {
             if (cachedToken.get() != null) {
-                return identityClient.authenticateWithMsalAccount(request, cachedToken.get().getAccount())
+                return identityClient.authenticateWithPublicClientCache(request, cachedToken.get().getAccount())
                     .onErrorResume(t -> Mono.empty());
             } else {
                 return Mono.empty();

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

@@ -56,7 +56,7 @@ public class UsernamePasswordCredential implements TokenCredential {
     public Mono<AccessToken> getToken(TokenRequestContext request) {
         return Mono.defer(() -> {
             if (cachedToken.get() != null) {
-                return identityClient.authenticateWithMsalAccount(request, cachedToken.get().getAccount())
+                return identityClient.authenticateWithPublicClientCache(request, cachedToken.get().getAccount())
                     .onErrorResume(t -> Mono.empty());
             } else {
                 return Mono.empty();

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

@@ -67,7 +67,7 @@ class VisualStudioCodeCredential implements TokenCredential {
     public Mono<AccessToken> getToken(TokenRequestContext request) {
         return Mono.defer(() -> {
             if (cachedToken.get() != null) {
-                return identityClient.authenticateWithMsalAccount(request, cachedToken.get().getAccount())
+                return identityClient.authenticateWithPublicClientCache(request, cachedToken.get().getAccount())
                            .onErrorResume(t -> Mono.empty());
             } else {
                 return Mono.empty();