Add token refresh options to identity library (#13029)

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

@@ -6,6 +6,7 @@
 import com.azure.core.annotation.Immutable;
 import com.azure.core.credential.AccessToken;
 import com.azure.core.credential.TokenCredential;
+import com.azure.core.credential.TokenRefreshOptions;
 import com.azure.core.credential.TokenRequestContext;
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.identity.implementation.IdentityClient;
@@ -27,6 +28,7 @@ public class AuthorizationCodeCredential implements TokenCredential {
     private final String authCode;
     private final URI redirectUri;
     private final IdentityClient identityClient;
+    private final IdentityClientOptions identityClientOptions;
     private final AtomicReference<MsalAuthenticationAccount> cachedToken;
     private final ClientLogger logger = new ClientLogger(AuthorizationCodeCredential.class);
 
@@ -46,6 +48,7 @@ public class AuthorizationCodeCredential implements TokenCredential {
             .clientId(clientId)
             .identityClientOptions(identityClientOptions)
             .build();
+        this.identityClientOptions = identityClientOptions;
         this.cachedToken = new AtomicReference<>();
         this.authCode = authCode;
         this.redirectUri = redirectUri;
@@ -71,4 +74,9 @@ public Mono<AccessToken> getToken(TokenRequestContext request) {
             .doOnNext(token -> LoggingUtil.logTokenSuccess(logger, request))
             .doOnError(error -> LoggingUtil.logTokenError(logger, request, error));
     }
+
+    @Override
+    public TokenRefreshOptions getTokenRefreshOptions() {
+        return identityClientOptions.getTokenRefreshOptions();
+    }
 }

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

@@ -6,6 +6,7 @@
 import com.azure.core.annotation.Immutable;
 import com.azure.core.credential.AccessToken;
 import com.azure.core.credential.TokenCredential;
+import com.azure.core.credential.TokenRefreshOptions;
 import com.azure.core.credential.TokenRequestContext;
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.identity.implementation.IdentityClient;
@@ -21,6 +22,7 @@
 @Immutable
 public class AzureCliCredential implements TokenCredential {
     private final IdentityClient identityClient;
+    private final IdentityClientOptions identityClientOptions;
     private final ClientLogger logger = new ClientLogger(AzureCliCredential.class);
 
     /**
@@ -29,6 +31,7 @@ public class AzureCliCredential implements TokenCredential {
      */
     AzureCliCredential(IdentityClientOptions identityClientOptions) {
         identityClient = new IdentityClientBuilder().identityClientOptions(identityClientOptions).build();
+        this.identityClientOptions = identityClientOptions;
     }
 
     @Override
@@ -37,4 +40,9 @@ public Mono<AccessToken> getToken(TokenRequestContext request) {
             .doOnNext(token -> LoggingUtil.logTokenSuccess(logger, request))
             .doOnError(error -> LoggingUtil.logTokenError(logger, request, error));
     }
+
+    @Override
+    public TokenRefreshOptions getTokenRefreshOptions() {
+        return identityClientOptions.getTokenRefreshOptions();
+    }
 }

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

@@ -6,6 +6,7 @@
 import com.azure.core.annotation.Immutable;
 import com.azure.core.credential.AccessToken;
 import com.azure.core.credential.TokenCredential;
+import com.azure.core.credential.TokenRefreshOptions;
 import com.azure.core.credential.TokenRequestContext;
 import com.azure.core.exception.ClientAuthenticationException;
 import com.azure.core.util.logging.ClientLogger;
@@ -25,6 +26,7 @@
  */
 @Immutable
 public class ChainedTokenCredential implements TokenCredential {
+    private volatile TokenRefreshOptions tokenRefreshOptions;
     private final ClientLogger logger = new ClientLogger(getClass());
     private final List<TokenCredential> credentials;
     private final String unavailableError = this.getClass().getSimpleName() + " authentication failed. ---> ";
@@ -55,7 +57,7 @@ public Mono<AccessToken> getToken(TokenRequestContext request) {
                     logger.info("Azure Identity => Attempted credential {} is unavailable.",
                         p.getClass().getSimpleName());
                     return Mono.empty();
-                }), 1)
+                }).doOnNext(t -> tokenRefreshOptions = p.getTokenRefreshOptions()), 1)
             .next()
             .switchIfEmpty(Mono.defer(() -> {
                 // Chain Exceptions.
@@ -69,6 +71,11 @@ public Mono<AccessToken> getToken(TokenRequestContext request) {
             }));
     }
 
+    @Override
+    public TokenRefreshOptions getTokenRefreshOptions() {
+        return tokenRefreshOptions;
+    }
+
 
     /**
      * Get the read-only list of credentials sequentially used to attempt authentication.

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

@@ -6,6 +6,7 @@
 import com.azure.core.annotation.Immutable;
 import com.azure.core.credential.AccessToken;
 import com.azure.core.credential.TokenCredential;
+import com.azure.core.credential.TokenRefreshOptions;
 import com.azure.core.credential.TokenRequestContext;
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.identity.implementation.IdentityClient;
@@ -28,6 +29,7 @@
 @Immutable
 public class ClientCertificateCredential implements TokenCredential {
     private final IdentityClient identityClient;
+    private final IdentityClientOptions identityClientOptions;
     private final ClientLogger logger = new ClientLogger(ClientCertificateCredential.class);
 
     /**
@@ -48,6 +50,7 @@ public class ClientCertificateCredential implements TokenCredential {
             .certificatePassword(certificatePassword)
             .identityClientOptions(identityClientOptions)
             .build();
+        this.identityClientOptions = identityClientOptions;
     }
 
     @Override
@@ -58,4 +61,9 @@ public Mono<AccessToken> getToken(TokenRequestContext request) {
             .doOnNext(token -> LoggingUtil.logTokenSuccess(logger, request))
             .doOnError(error -> LoggingUtil.logTokenError(logger, request, error));
     }
+
+    @Override
+    public TokenRefreshOptions getTokenRefreshOptions() {
+        return identityClientOptions.getTokenRefreshOptions();
+    }
 }

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

@@ -6,6 +6,7 @@
 import com.azure.core.annotation.Immutable;
 import com.azure.core.credential.AccessToken;
 import com.azure.core.credential.TokenCredential;
+import com.azure.core.credential.TokenRefreshOptions;
 import com.azure.core.credential.TokenRequestContext;
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.identity.implementation.IdentityClient;
@@ -28,6 +29,7 @@
 @Immutable
 public class ClientSecretCredential implements TokenCredential {
     private final IdentityClient identityClient;
+    private final IdentityClientOptions identityClientOptions;
     private final ClientLogger logger = new ClientLogger(ClientSecretCredential.class);
 
     /**
@@ -48,6 +50,7 @@ public class ClientSecretCredential implements TokenCredential {
             .clientSecret(clientSecret)
             .identityClientOptions(identityClientOptions)
             .build();
+        this.identityClientOptions = identityClientOptions;
     }
 
     @Override
@@ -58,4 +61,9 @@ public Mono<AccessToken> getToken(TokenRequestContext request) {
             .doOnNext(token -> LoggingUtil.logTokenSuccess(logger, request))
             .doOnError(error -> LoggingUtil.logTokenError(logger, request, error));
     }
+
+    @Override
+    public TokenRefreshOptions getTokenRefreshOptions() {
+        return identityClientOptions.getTokenRefreshOptions();
+    }
 }

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

@@ -3,6 +3,7 @@
 
 package com.azure.identity;
 
+import com.azure.core.credential.SimpleTokenCache;
 import com.azure.core.http.HttpClient;
 import com.azure.core.http.HttpPipeline;
 import com.azure.core.http.ProxyOptions;
@@ -95,8 +96,16 @@ public T httpClient(HttpClient client) {
      * token will be considered expired at and after the time of (actual
      * expiry - token refresh offset). The default offset is 2 minutes.
      *
-     * This is useful when network is congested and a request containing the
-     * token takes longer than normal to get to the server.
+     * This is used in {@link SimpleTokenCache} and {@link com.azure.core.http.policy.BearerTokenAuthenticationPolicy}
+     * to proactively retrieve a more up-to-date token before the cached token gets too close to its expiry.
+     *
+     * Extending this offset is recommended if it takes > 2 minutes to reach the service (application is running on
+     * high load), or you would like to simply keep a more up-to-date token in the cache (more robust against token
+     * refresh API down times). The user is responsible for specifying a valid offset.
+     *
+     * When a proactive token refresh fails but the previously cached token is still valid,
+     * {@link com.azure.core.http.policy.BearerTokenAuthenticationPolicy} will NOT fail but return the previous valid
+     * token. Another proactive refresh will be attempted in 30 seconds.
      *
      * @param tokenRefreshOffset the duration before the actual expiry of a token to refresh it
      * @return An updated instance of this builder with the token refresh offset set as specified.

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

@@ -5,6 +5,8 @@
 
 import com.azure.core.annotation.Immutable;
 import com.azure.core.credential.TokenCredential;
+import com.azure.core.credential.TokenRefreshOptions;
+import com.azure.identity.implementation.IdentityClientOptions;
 
 import java.util.List;
 
@@ -22,6 +24,7 @@
  */
 @Immutable
 public final class DefaultAzureCredential extends ChainedTokenCredential {
+    private final IdentityClientOptions identityClientOptions;
 
     /**
      * Creates default DefaultAzureCredential instance to use. This will use AZURE_CLIENT_ID,
@@ -32,9 +35,16 @@ public final class DefaultAzureCredential extends ChainedTokenCredential {
      * token cache.
      *
      * @param tokenCredentials the list of credentials to execute for authentication.
+     * @param identityClientOptions the options for configuring the identity client.
      */
-    DefaultAzureCredential(List<TokenCredential> tokenCredentials) {
+    DefaultAzureCredential(List<TokenCredential> tokenCredentials, IdentityClientOptions identityClientOptions) {
         super(tokenCredentials);
+        this.identityClientOptions = identityClientOptions;
+    }
+
+    @Override
+    public TokenRefreshOptions getTokenRefreshOptions() {
+        return identityClientOptions.getTokenRefreshOptions();
     }
 
 

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

@@ -95,7 +95,7 @@ public DefaultAzureCredentialBuilder executorService(ExecutorService executorSer
      * @return a {@link DefaultAzureCredential} with the current configurations.
      */
     public DefaultAzureCredential build() {
-        return new DefaultAzureCredential(getCredentialsChain());
+        return new DefaultAzureCredential(getCredentialsChain(), identityClientOptions);
     }
 
     private ArrayList<TokenCredential> getCredentialsChain() {

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

@@ -6,6 +6,7 @@
 import com.azure.core.annotation.Immutable;
 import com.azure.core.credential.AccessToken;
 import com.azure.core.credential.TokenCredential;
+import com.azure.core.credential.TokenRefreshOptions;
 import com.azure.core.credential.TokenRequestContext;
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.identity.implementation.IdentityClient;
@@ -26,6 +27,7 @@
 public class DeviceCodeCredential implements TokenCredential {
     private final Consumer<DeviceCodeInfo> challengeConsumer;
     private final IdentityClient identityClient;
+    private final IdentityClientOptions identityClientOptions;
     private final AtomicReference<MsalAuthenticationAccount> cachedToken;
     private final String authorityHost;
     private final boolean automaticAuthentication;
@@ -50,6 +52,7 @@ public class DeviceCodeCredential implements TokenCredential {
             .identityClientOptions(identityClientOptions)
             .build();
         this.cachedToken = new AtomicReference<>();
+        this.identityClientOptions = identityClientOptions;
         this.authorityHost = identityClientOptions.getAuthorityHost();
         this.automaticAuthentication = automaticAuthentication;
         if (identityClientOptions.getAuthenticationRecord() != null) {
@@ -126,4 +129,9 @@ private AccessToken updateCache(MsalToken msalToken) {
                                 identityClient.getTenantId())));
         return msalToken;
     }
+
+    @Override
+    public TokenRefreshOptions getTokenRefreshOptions() {
+        return identityClientOptions.getTokenRefreshOptions();
+    }
 }

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

@@ -6,6 +6,7 @@
 import com.azure.core.annotation.Immutable;
 import com.azure.core.credential.AccessToken;
 import com.azure.core.credential.TokenCredential;
+import com.azure.core.credential.TokenRefreshOptions;
 import com.azure.core.credential.TokenRequestContext;
 import com.azure.core.util.Configuration;
 import com.azure.core.util.logging.ClientLogger;
@@ -128,6 +129,11 @@ public Mono<AccessToken> getToken(TokenRequestContext request) {
         }
     }
 
+    @Override
+    public TokenRefreshOptions getTokenRefreshOptions() {
+        return identityClientOptions.getTokenRefreshOptions();
+    }
+
     private boolean verifyNotNull(String... configs) {
         for (String config: configs) {
             if (config == null) {

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

@@ -6,6 +6,7 @@
 import com.azure.core.annotation.Immutable;
 import com.azure.core.credential.AccessToken;
 import com.azure.core.credential.TokenCredential;
+import com.azure.core.credential.TokenRefreshOptions;
 import com.azure.core.credential.TokenRequestContext;
 import com.azure.core.util.CoreUtils;
 import com.azure.core.util.logging.ClientLogger;
@@ -30,6 +31,7 @@
 public class IntelliJCredential implements TokenCredential {
     private static final String AZURE_TOOLS_FOR_INTELLIJ_CLIENT_ID = "61d65f5a-6e3b-468b-af73-a033f5098c5c";
     private final IdentityClient identityClient;
+    private final IdentityClientOptions identityClientOptions;
     private final AtomicReference<MsalToken> cachedToken;
     private final ClientLogger logger = new ClientLogger(IntelliJCredential.class);
 
@@ -72,6 +74,7 @@ public class IntelliJCredential implements TokenCredential {
                              .build();
 
         this.cachedToken = new AtomicReference<>();
+        this.identityClientOptions = identityClientOptions;
     }
 
     @Override
@@ -92,4 +95,9 @@ public Mono<AccessToken> getToken(TokenRequestContext request) {
             .doOnNext(token -> LoggingUtil.logTokenSuccess(logger, request))
             .doOnError(error -> LoggingUtil.logTokenError(logger, request, error));
     }
+
+    @Override
+    public TokenRefreshOptions getTokenRefreshOptions() {
+        return identityClientOptions.getTokenRefreshOptions();
+    }
 }

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

@@ -6,6 +6,7 @@
 import com.azure.core.annotation.Immutable;
 import com.azure.core.credential.AccessToken;
 import com.azure.core.credential.TokenCredential;
+import com.azure.core.credential.TokenRefreshOptions;
 import com.azure.core.credential.TokenRequestContext;
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.identity.implementation.IdentityClient;
@@ -31,6 +32,7 @@
 public class InteractiveBrowserCredential implements TokenCredential {
     private final int port;
     private final IdentityClient identityClient;
+    private final IdentityClientOptions identityClientOptions;
     private final AtomicReference<MsalAuthenticationAccount> cachedToken;
     private final boolean automaticAuthentication;
     private final String authorityHost;
@@ -55,6 +57,7 @@ public class InteractiveBrowserCredential implements TokenCredential {
             .clientId(clientId)
             .identityClientOptions(identityClientOptions)
             .build();
+        this.identityClientOptions = identityClientOptions;
         cachedToken = new AtomicReference<>();
         this.authorityHost = identityClientOptions.getAuthorityHost();
         this.automaticAuthentication = automaticAuthentication;
@@ -115,12 +118,16 @@ public Mono<AuthenticationRecord> authenticate() {
         return authenticate(new TokenRequestContext().addScopes(defaultScope));
     }
 
+    @Override
+    public TokenRefreshOptions getTokenRefreshOptions() {
+        return identityClientOptions.getTokenRefreshOptions();
+    }
+
     private AccessToken updateCache(MsalToken msalToken) {
         cachedToken.set(
                 new MsalAuthenticationAccount(
                         new AuthenticationRecord(msalToken.getAuthenticationResult(),
                                 identityClient.getTenantId())));
         return msalToken;
     }
-
 }