Set refreshOn in MI flows to half of token lifetime

Author: Avery-Dunn

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenByManagedIdentitySupplier.java

@@ -13,6 +13,8 @@ class AcquireTokenByManagedIdentitySupplier extends AuthenticationResultSupplier
 
     private static final Logger LOG = LoggerFactory.getLogger(AcquireTokenByManagedIdentitySupplier.class);
 
+    private static final int TWO_HOURS = 2*3600;
+
     private ManagedIdentityParameters managedIdentityParameters;
 
     AcquireTokenByManagedIdentitySupplier(ManagedIdentityApplication managedIdentityApplication, MsalRequest msalRequest) {
@@ -93,15 +95,27 @@ private AuthenticationResult fetchNewAccessTokenAndSaveToCache(TokenRequestExecu
     }
 
     private AuthenticationResult createFromManagedIdentityResponse(ManagedIdentityResponse managedIdentityResponse) {
-        long expiresOn = Long.valueOf(managedIdentityResponse.expiresOn);
-        long refreshOn = expiresOn > 2 * 3600 ? (expiresOn / 2) : 0L;
+        long expiresOn = Long.parseLong(managedIdentityResponse.expiresOn);
+        long refreshOn = calculateRefreshOn(expiresOn);
+        AuthenticationResultMetadata metadata = AuthenticationResultMetadata.builder()
+                .refreshOn(refreshOn)
+                .build();
 
         return AuthenticationResult.builder()
                 .accessToken(managedIdentityResponse.getAccessToken())
                 .scopes(managedIdentityParameters.resource())
                 .expiresOn(expiresOn)
                 .extExpiresOn(0)
                 .refreshOn(refreshOn)
+                .metadata(metadata)
                 .build();
     }
+
+    private long calculateRefreshOn(long expiresOn){
+        long timestampSeconds = System.currentTimeMillis() / 1000;
+        long expiresIn = expiresOn - timestampSeconds;
+
+        //The refreshOn value should be half the value of the token lifetime, if the lifetime is greater than two hours
+        return expiresIn > TWO_HOURS ? (expiresIn / 2) + timestampSeconds : 0;
+    }
 }

msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/ManagedIdentityTests.java

@@ -44,7 +44,7 @@ class ManagedIdentityTests {
     private static ManagedIdentityApplication miApp;
 
     private String getSuccessfulResponse(String resource) {
-        long expiresOn = Instant.now().plus(1, ChronoUnit.HOURS).getEpochSecond();
+        long expiresOn = (System.currentTimeMillis() / 1000) + (24 * 3600);//A long-lived, 24 hour token
         return "{\"access_token\":\"accesstoken\",\"expires_on\":\"" + expiresOn + "\",\"resource\":\"" + resource + "\",\"token_type\":" +
                 "\"Bearer\",\"client_id\":\"client_id\"}";
     }
@@ -222,6 +222,34 @@ void managedIdentityTest_UserAssigned_SuccessfulResponse(ManagedIdentitySourceTy
         verify(httpClientMock, times(1)).send(any());
     }
 
+    @Test
+    void managedIdentityTest_RefreshOnHalfOfExpiresOn() throws Exception {
+        //All managed identity flows use the same AcquireTokenByManagedIdentitySupplier where refreshOn is set,
+        //  so any of the MI options should let us verify that it's being set correctly
+        IEnvironmentVariables environmentVariables = new EnvironmentVariablesHelper(ManagedIdentitySourceType.APP_SERVICE, appServiceEndpoint);
+        ManagedIdentityApplication.setEnvironmentVariables(environmentVariables);
+        ManagedIdentityClient.resetManagedIdentitySourceType();
+        DefaultHttpClient httpClientMock = mock(DefaultHttpClient.class);
+
+        when(httpClientMock.send(expectedRequest(ManagedIdentitySourceType.APP_SERVICE, resource))).thenReturn(expectedResponse(200, getSuccessfulResponse(resource)));
+
+        miApp = ManagedIdentityApplication
+                .builder(ManagedIdentityId.systemAssigned())
+                .httpClient(httpClientMock)
+                .build();
+
+        AuthenticationResult result = (AuthenticationResult) miApp.acquireTokenForManagedIdentity(
+                ManagedIdentityParameters.builder(resource)
+                        .build()).get();
+
+        long timestampSeconds = (System.currentTimeMillis() / 1000);
+
+        assertNotNull(result.accessToken());
+        assertEquals((result.expiresOn() - timestampSeconds)/2, result.refreshOn() - timestampSeconds);
+
+        verify(httpClientMock, times(1)).send(any());
+    }
+
     @ParameterizedTest
     @MethodSource("com.microsoft.aad.msal4j.ManagedIdentityTestDataProvider#createDataUserAssignedNotSupported")
     void managedIdentityTest_UserAssigned_NotSupported(ManagedIdentitySourceType source, String endpoint, ManagedIdentityId id) throws Exception {