Add configuration for refresh token before expiry

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

@@ -84,4 +84,21 @@ public T httpClient(HttpClient client) {
         this.identityClientOptions.setHttpClient(client);
         return (T) this;
     }
+
+    /**
+     * Sets how long before the actual token expiry to refresh the token. The
+     * 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.
+     *
+     * @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.
+     */
+    @SuppressWarnings("unchecked")
+    public T tokenRefreshOffset(Duration tokenRefreshOffset) {
+        this.identityClientOptions.setTokenRefreshOffset(tokenRefreshOffset);
+        return (T) this;
+    }
 }

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

@@ -46,7 +46,6 @@
 import java.nio.file.Paths;
 import java.time.Duration;
 import java.time.OffsetDateTime;
-import java.time.ZoneOffset;
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
@@ -147,10 +146,8 @@ public Mono<AccessToken> authenticateWithClientSecret(String clientSecret, Token
             ConfidentialClientApplication application = applicationBuilder.build();
             return Mono.fromFuture(application.acquireToken(
                 ClientCredentialParameters.builder(new HashSet<>(request.getScopes()))
-                   .build()))
-                   .map(ar -> new AccessToken(ar.accessToken(), OffsetDateTime.ofInstant(ar.expiresOnDate().toInstant(),
-                               ZoneOffset.UTC)));
-
+                    .build()))
+                .map(ar -> new MsalToken(ar, options));
         } catch (MalformedURLException e) {
             return Mono.error(e);
         }
@@ -194,8 +191,7 @@ public Mono<AccessToken> authenticateWithPfxCertificate(String pfxCertificatePat
             return applicationBuilder.build();
         }).flatMap(application -> Mono.fromFuture(application.acquireToken(
                 ClientCredentialParameters.builder(new HashSet<>(request.getScopes())).build())))
-        .map(ar -> new AccessToken(ar.accessToken(), OffsetDateTime.ofInstant(ar.expiresOnDate().toInstant(),
-                ZoneOffset.UTC)));
+        .map(ar -> new MsalToken(ar, options));
     }
 
     /**
@@ -226,8 +222,7 @@ public Mono<AccessToken> authenticateWithPemCertificate(String pemCertificatePat
             return Mono.fromFuture(application.acquireToken(
                 ClientCredentialParameters.builder(new HashSet<>(request.getScopes()))
                     .build()))
-                .map(ar -> new AccessToken(ar.accessToken(), OffsetDateTime.ofInstant(ar.expiresOnDate().toInstant(),
-                    ZoneOffset.UTC)));
+                .map(ar -> new MsalToken(ar, options));
         } catch (IOException e) {
             return Mono.error(e);
         }
@@ -246,7 +241,7 @@ public Mono<MsalToken> authenticateWithUsernamePassword(TokenRequestContext requ
         return Mono.fromFuture(publicClientApplication.acquireToken(
             UserNamePasswordParameters.builder(new HashSet<>(request.getScopes()), username, password.toCharArray())
                 .build()))
-            .map(MsalToken::new);
+            .map(ar -> new MsalToken(ar, options));
     }
 
     /**
@@ -264,7 +259,8 @@ public Mono<MsalToken> authenticateWithUserRefreshToken(TokenRequestContext requ
         }
         return Mono.defer(() -> {
             try {
-                return Mono.fromFuture(publicClientApplication.acquireTokenSilently(parameters)).map(MsalToken::new);
+                return Mono.fromFuture(publicClientApplication.acquireTokenSilently(parameters))
+                        .map(ar -> new MsalToken(ar, options));
             } catch (MalformedURLException e) {
                 return Mono.error(e);
             }
@@ -288,7 +284,7 @@ public Mono<MsalToken> authenticateWithDeviceCode(TokenRequestContext request,
                 dc -> deviceCodeConsumer.accept(new DeviceCodeInfo(dc.userCode(), dc.deviceCode(),
                     dc.verificationUri(), OffsetDateTime.now().plusSeconds(dc.expiresIn()), dc.message()))).build();
             return publicClientApplication.acquireToken(parameters);
-        }).map(MsalToken::new);
+        }).map(ar -> new MsalToken(ar, options));
     }
 
     /**
@@ -305,7 +301,7 @@ public Mono<MsalToken> authenticateWithAuthorizationCode(TokenRequestContext req
             AuthorizationCodeParameters.builder(authorizationCode, redirectUrl)
                 .scopes(new HashSet<>(request.getScopes()))
                 .build()))
-            .map(MsalToken::new);
+            .map(ar -> new MsalToken(ar, options));
     }
 
     /**
@@ -362,11 +358,11 @@ public Mono<MsalToken> authenticateWithBrowserInteraction(TokenRequestContext re
      */
     public Mono<AccessToken> authenticateToManagedIdentityEndpoint(String msiEndpoint, String msiSecret,
                                                                    TokenRequestContext request) {
-        String resource = ScopeUtil.scopesToResource(request.getScopes());
-        HttpURLConnection connection = null;
-        StringBuilder payload = new StringBuilder();
+        return Mono.fromCallable(() -> {
+            String resource = ScopeUtil.scopesToResource(request.getScopes());
+            HttpURLConnection connection = null;
+            StringBuilder payload = new StringBuilder();
 
-        try {
             payload.append("resource=");
             payload.append(URLEncoder.encode(resource, "UTF-8"));
             payload.append("&api-version=");
@@ -375,32 +371,30 @@ public Mono<AccessToken> authenticateToManagedIdentityEndpoint(String msiEndpoin
                 payload.append("&clientid=");
                 payload.append(URLEncoder.encode(clientId, "UTF-8"));
             }
-        } catch (IOException exception) {
-            return Mono.error(exception);
-        }
-        try {
-            URL url = new URL(String.format("%s?%s", msiEndpoint, payload));
-            connection = (HttpURLConnection) url.openConnection();
+            try {
+                URL url = new URL(String.format("%s?%s", msiEndpoint, payload));
+                connection = (HttpURLConnection) url.openConnection();
 
-            connection.setRequestMethod("GET");
-            if (msiSecret != null) {
-                connection.setRequestProperty("Secret", msiSecret);
-            }
-            connection.setRequestProperty("Metadata", "true");
+                connection.setRequestMethod("GET");
+                if (msiSecret != null) {
+                    connection.setRequestProperty("Secret", msiSecret);
+                }
+                connection.setRequestProperty("Metadata", "true");
 
-            connection.connect();
+                connection.connect();
 
-            Scanner s = new Scanner(connection.getInputStream(), StandardCharsets.UTF_8.name()).useDelimiter("\\A");
-            String result = s.hasNext() ? s.next() : "";
+                Scanner s = new Scanner(connection.getInputStream(), StandardCharsets.UTF_8.name())
+                        .useDelimiter("\\A");
+                String result = s.hasNext() ? s.next() : "";
 
-            return Mono.just(SERIALIZER_ADAPTER.deserialize(result, MSIToken.class, SerializerEncoding.JSON));
-        } catch (IOException e) {
-            return Mono.error(e);
-        } finally {
-            if (connection != null) {
-                connection.disconnect();
+                MSIToken msiToken = SERIALIZER_ADAPTER.deserialize(result, MSIToken.class, SerializerEncoding.JSON);
+                return new IdentityToken(msiToken.getToken(), msiToken.getExpiresAt(), options);
+            } finally {
+                if (connection != null) {
+                    connection.disconnect();
+                }
             }
-        }
+        });
     }
 
     /**
@@ -446,7 +440,9 @@ public Mono<AccessToken> authenticateToIMDSEndpoint(TokenRequestContext request)
                             .useDelimiter("\\A");
                     String result = s.hasNext() ? s.next() : "";
 
-                    return SERIALIZER_ADAPTER.<MSIToken>deserialize(result, MSIToken.class, SerializerEncoding.JSON);
+                    MSIToken msiToken = SERIALIZER_ADAPTER.deserialize(result,
+                            MSIToken.class, SerializerEncoding.JSON);
+                    return new IdentityToken(msiToken.getToken(), msiToken.getExpiresAt(), options);
                 } catch (IOException exception) {
                     if (connection == null) {
                         throw logger.logExceptionAsError(new RuntimeException(

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

@@ -22,6 +22,7 @@ public final class IdentityClientOptions {
     private Function<Duration, Duration> retryTimeout;
     private ProxyOptions proxyOptions;
     private HttpPipeline httpPipeline;
+    private Duration tokenRefreshOffset = Duration.ofMinutes(2);
     private HttpClient httpClient;
 
     /**
@@ -125,6 +126,30 @@ public IdentityClientOptions setHttpPipeline(HttpPipeline httpPipeline) {
         return this;
     }
 
+    /**
+     * @return how long before the actual token expiry to refresh the token.
+     */
+    public Duration getTokenRefreshOffset() {
+        return tokenRefreshOffset;
+    }
+
+    /**
+     * Sets how long before the actual token expiry to refresh the token. The
+     * 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.
+     *
+     * @param tokenRefreshOffset the duration before the actual expiry of a token to refresh it
+     */
+    public IdentityClientOptions setTokenRefreshOffset(Duration tokenRefreshOffset) {
+        if (tokenRefreshOffset != null) {
+            this.tokenRefreshOffset = tokenRefreshOffset;
+        }
+        return this;
+    }
+
     /**
      * Specifies the HttpClient to send use for requests.
      * @param httpClient the http client to use for requests

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

@@ -0,0 +1,24 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.identity.implementation;
+
+import com.azure.core.credential.AccessToken;
+
+import java.time.OffsetDateTime;
+
+/**
+ * Type representing authentication result from the azure-identity client.
+ */
+public class IdentityToken extends AccessToken {
+    /**
+     * Creates an identity token instance.
+     *
+     * @param token the token string.
+     * @param expiresAt the expiration time.
+     * @param options the identity client options.
+     */
+    public IdentityToken(String token, OffsetDateTime expiresAt, IdentityClientOptions options) {
+        super(token, expiresAt.plusMinutes(2).minus(options.getTokenRefreshOffset()));
+    }
+}

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

@@ -3,7 +3,6 @@
 
 package com.azure.identity.implementation;
 
-import com.azure.core.credential.AccessToken;
 import com.microsoft.aad.msal4j.IAccount;
 import com.microsoft.aad.msal4j.IAuthenticationResult;
 
@@ -13,7 +12,7 @@
 /**
  * Type representing authentication result from the MSAL (Microsoft Authentication Library).
  */
-public final class MsalToken extends AccessToken {
+public final class MsalToken extends IdentityToken {
 
     private IAccount account;
 
@@ -22,9 +21,10 @@ public final class MsalToken extends AccessToken {
      *
      * @param msalResult the raw authentication result returned by MSAL
      */
-    public MsalToken(IAuthenticationResult msalResult) {
-        super(msalResult.accessToken(), OffsetDateTime.ofInstant(msalResult.expiresOnDate().toInstant(),
-            ZoneOffset.UTC));
+    public MsalToken(IAuthenticationResult msalResult, IdentityClientOptions options) {
+        super(msalResult.accessToken(),
+                OffsetDateTime.ofInstant(msalResult.expiresOnDate().toInstant(), ZoneOffset.UTC),
+                options);
         this.account = msalResult.account();
     }
 

sdk/identity/azure-identity/src/test/java/com/azure/identity/ClientSecretCredentialTest.java

@@ -17,6 +17,7 @@
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
+import java.time.Duration;
 import java.time.OffsetDateTime;
 import java.time.ZoneOffset;
 import java.util.UUID;
@@ -61,6 +62,40 @@ public void testValidSecrets() throws Exception {
             .verifyComplete();
     }
 
+    @Test
+    public void testValidSecretsWithTokenRefreshOffset() throws Exception {
+        // setup
+        String secret = "secret";
+        String token1 = "token1";
+        String token2 = "token2";
+        TokenRequestContext request1 = new TokenRequestContext().addScopes("https://management.azure.com");
+        TokenRequestContext request2 = new TokenRequestContext().addScopes("https://vault.azure.net");
+        OffsetDateTime expiresAt = OffsetDateTime.now(ZoneOffset.UTC).plusHours(1);
+        Duration offset = Duration.ofMinutes(10);
+
+        // mock
+        IdentityClient identityClient = PowerMockito.mock(IdentityClient.class);
+        when(identityClient.authenticateWithClientSecret(secret, request1)).thenReturn(TestUtils.getMockAccessToken(token1, expiresAt, offset));
+        when(identityClient.authenticateWithClientSecret(secret, request2)).thenReturn(TestUtils.getMockAccessToken(token2, expiresAt, offset));
+        PowerMockito.whenNew(IdentityClient.class).withAnyArguments().thenReturn(identityClient);
+
+        // test
+        ClientSecretCredential credential = new ClientSecretCredentialBuilder()
+                .tenantId(tenantId)
+                .clientId(clientId)
+                .clientSecret(secret)
+                .tokenRefreshOffset(offset)
+                .build();
+        StepVerifier.create(credential.getToken(request1))
+                .expectNextMatches(accessToken -> token1.equals(accessToken.getToken())
+                        && expiresAt.getSecond() == accessToken.getExpiresAt().getSecond())
+                .verifyComplete();
+        StepVerifier.create(credential.getToken(request2))
+                .expectNextMatches(accessToken -> token2.equals(accessToken.getToken())
+                        && expiresAt.getSecond() == accessToken.getExpiresAt().getSecond())
+                .verifyComplete();
+    }
+
     @Test
     public void testInvalidSecrets() throws Exception {
         // setup