Fix for issue #3864

authorization/msi-auth-token-provider-jar/src/main/java/com/microsoft/azure/msiAuthTokenProvider/MSIConfigurationForVirtualMachine.java

@@ -17,6 +17,7 @@ public class MSIConfigurationForVirtualMachine {
     private String clientId;
     private String identityId;
     private int maxRetry = -1;
+    private int customTimeout = -1;
 
     /**
      * Creates MSIConfigurationForVirtualMachine.
@@ -85,6 +86,11 @@ public int maxRetry() {
         return this.maxRetry;
     }
 
+    /**
+     * @return the custom timeout (in milliseconds) when getting the token from IMDS
+     */
+    public int customTimeout() { return this.customTimeout; }
+
     /**
      * Specifies the token retrieval source.
      *
@@ -153,6 +159,17 @@ public MSIConfigurationForVirtualMachine withMaxRetry(int maxRetry) {
         return this;
     }
 
+    /**
+     * Specifies the custom timeout (in milliseconds) to be used for IMDS retries.
+     *
+     * @param timeoutInMs the total timeout value in milliseconds;
+     * @return MSIConfigurationForVirtualMachine
+     */
+    public MSIConfigurationForVirtualMachine withCustomTimeout(int timeoutInMs) {
+        this.customTimeout = timeoutInMs;
+        return this;
+    }
+
     @Override
     public MSIConfigurationForVirtualMachine clone() {
         MSIConfigurationForVirtualMachine copy = new MSIConfigurationForVirtualMachine(this.managementEndpoint);
@@ -172,6 +189,7 @@ public MSIConfigurationForVirtualMachine clone() {
             copy.withTokenSource(this.tokenSource());
         }
         copy.withMaxRetry(this.maxRetry());
+        copy.withCustomTimeout(this.customTimeout());
         return copy;
     }
 
@@ -185,4 +203,4 @@ public enum MSITokenSource {
          */
         IMDS_ENDPOINT
     }
-}
+}

authorization/msi-auth-token-provider-jar/src/main/java/com/microsoft/azure/msiAuthTokenProvider/MSICredentials.java

@@ -34,6 +34,7 @@ public final class MSICredentials{
     private final MSIConfigurationForAppService configForAppService;
     private final HostType hostType;
     private final int maxRetry;
+    private final int customTimeout;
     private static final int MAX_RETRY_DEFAULT_LIMIT = 20;
 
 
@@ -121,9 +122,10 @@ private MSICredentials(MSIConfigurationForVirtualMachine config) {
         this.configForAppService = null;
         this.hostType = HostType.VIRTUAL_MACHINE;
         this.maxRetry = config.maxRetry() < 0 ? MAX_RETRY_DEFAULT_LIMIT : config.maxRetry();
+        this.customTimeout = config.customTimeout();
         // Simplified variant of https://en.wikipedia.org/wiki/Exponential_backoff
         for (int x = 0; x < this.maxRetry; x++) {
-            this.retrySlots.add(500 * ((2 << 1) - 1) / 1000);
+            this.retrySlots.add(500 * ((2 << x) - 1) / 1000);
         }
     }
 
@@ -132,6 +134,7 @@ private MSICredentials(MSIConfigurationForAppService config) {
         this.configForVM = null;
         this.hostType = HostType.APP_SERVICE;
         this.maxRetry = -1;
+        this.customTimeout = -1;
     }
 
     /**
@@ -271,6 +274,7 @@ private MSIToken getTokenForVirtualMachineFromIMDSEndpoint(String tokenAudience)
     private MSIToken retrieveTokenFromIDMSWithRetry(String tokenAudience) throws AzureMSICredentialException, IOException {
         StringBuilder payload = new StringBuilder();
         final int imdsUpgradeTimeInMs = 70 * 1000;
+        boolean hasTimedout = false;
 
         //
         try {
@@ -301,12 +305,15 @@ private MSIToken retrieveTokenFromIDMSWithRetry(String tokenAudience) throws Azu
             throw new AzureMSICredentialException(exception);
         }
 
+        //A 0 custom timeout implies only 1 try... no more
+        hasTimedout = this.customTimeout == 0;
         int retry = 1;
         while (retry <= maxRetry) {
             URL url = new URL(String.format("http://169.254.169.254/metadata/identity/oauth2/token?%s", payload.toString()));
             //
             HttpURLConnection connection = null;
             //
+            long startTime = Calendar.getInstance().getTime().getTime();
             try {
                 connection = (HttpURLConnection) url.openConnection();
                 connection.setRequestMethod("GET");
@@ -319,6 +326,9 @@ private MSIToken retrieveTokenFromIDMSWithRetry(String tokenAudience) throws Azu
             } catch (Exception exception) {
                 int responseCode = connection.getResponseCode();
                 if (responseCode == 410 || responseCode == 429 || responseCode == 404 || (responseCode >= 500 && responseCode <= 599)) {
+                    if (hasTimedout) {
+                        throw new AzureMSICredentialException("Couldn't acquire access token from IMDS within the specified timeout : " + this.customTimeout + " milliseconds");
+                    }
                     int retryTimeoutInMs = retrySlots.get(new Random().nextInt(retry));
                     // Error code 410 indicates IMDS upgrade is in progress, which can take up to 70s
                     //
@@ -327,7 +337,7 @@ private MSIToken retrieveTokenFromIDMSWithRetry(String tokenAudience) throws Azu
                     if (retry > maxRetry) {
                         break;
                     } else {
-                        sleep(retryTimeoutInMs);
+                        hasTimedout = sleep(retryTimeoutInMs, startTime);
                     }
                 } else {
                     throw new AzureMSICredentialException("Couldn't acquire access token from IMDS, verify your objectId, clientId or msiResourceId", exception);
@@ -345,6 +355,28 @@ private MSIToken retrieveTokenFromIDMSWithRetry(String tokenAudience) throws Azu
         return null;
     }
 
+    /**
+     * Sleep for timeToWait or time remaining until timeout reached.
+     * @param timeToWaitinMs Time to wait in milliseconds
+     * @param startTime Abcolute tim in milliseconds
+     * @return true if we used the custom timeout.
+     */
+    private boolean sleep(int timeToWaitinMs, long startTime) {
+        long timeToSleep = 0;
+
+        if (this.customTimeout > -1) {
+            //timeToSleep = ;
+            long timeRemainingToTimeout = (startTime + this.customTimeout - Calendar.getInstance().getTime().getTime());
+            timeRemainingToTimeout = (timeToWaitinMs < timeRemainingToTimeout) ? timeToWaitinMs : timeRemainingToTimeout;
+            timeToSleep = (timeRemainingToTimeout > 0) ? timeRemainingToTimeout : 0;
+        } else {
+            timeToSleep = timeToWaitinMs;
+        }
+
+        sleep(timeToSleep);
+        return (timeToSleep != timeToWaitinMs);
+    }
+
     private static MSIToken getMsiTokenFromResult(String result, HostType hostType) throws AzureMSICredentialException{
         try {
             return new MSIToken(getTokenFromResult(result), getTokenTypeFromResult(result), getExpiryTimeFromResult(result, hostType));
@@ -401,7 +433,7 @@ private static Date getExpiryTimeFromResult(String result, HostType hostType) th
         return cal.getTime();
     }
 
-    private static void sleep(int millis) {
+    private static void sleep(long millis) {
         try {
             Thread.sleep(millis);
         } catch (InterruptedException ex) {
@@ -422,4 +454,4 @@ private enum HostType {
          */
         APP_SERVICE
     }
-}
+}