Merge pull request #960 from AzureAD/avdunn/retry-behavior

Refactor retry logic and add special behavior for IMDS

Author: Avery-Dunn

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

@@ -235,8 +235,8 @@ static AadInstanceDiscoveryResponse sendInstanceDiscoveryRequest(URL authorityUr
 
         AadInstanceDiscoveryResponse response = JsonHelper.convertJsonStringToJsonSerializableObject(httpResponse.body(), AadInstanceDiscoveryResponse::fromJson);
 
-        if (httpResponse.statusCode() != HttpHelper.HTTP_STATUS_200) {
-            if (httpResponse.statusCode() == HttpHelper.HTTP_STATUS_400 && response.error().equals("invalid_instance")) {
+        if (httpResponse.statusCode() != HttpStatus.HTTP_OK) {
+            if (httpResponse.statusCode() == HttpStatus.HTTP_BAD_REQUEST && response.error().equals("invalid_instance")) {
                 // instance discovery failed due to an invalid authority, throw an exception.
                 throw MsalServiceExceptionFactory.fromHttpResponse(httpResponse);
             }
@@ -310,7 +310,7 @@ static String discoverRegion(MsalRequest msalRequest, ServiceBundle serviceBundl
             log.info("Starting call to IMDS endpoint.");
             IHttpResponse httpResponse = future.get(IMDS_TIMEOUT, IMDS_TIMEOUT_UNIT);
             //If call to IMDS endpoint was successful, return region from response body
-            if (httpResponse.statusCode() == HttpHelper.HTTP_STATUS_200 && !httpResponse.body().isEmpty()) {
+            if (httpResponse.statusCode() == HttpStatus.HTTP_OK && !httpResponse.body().isEmpty()) {
                 log.info(String.format("Region retrieved from IMDS endpoint: %s", httpResponse.body()));
                 currentRequest.regionSource(RegionTelemetry.REGION_SOURCE_IMDS.telemetryValue);
 

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

@@ -32,6 +32,7 @@ public abstract class AbstractApplicationBase implements IApplicationBase {
     private IHttpClient httpClient;
     private Integer connectTimeoutForDefaultHttpClient;
     private Integer readTimeoutForDefaultHttpClient;
+    private boolean retryDisabled;
     String tenant;
 
     //The following fields are set in only some applications and/or set internally by the library. To avoid excessive
@@ -150,6 +151,10 @@ public Integer readTimeoutForDefaultHttpClient() {
         return this.readTimeoutForDefaultHttpClient;
     }
 
+    boolean isRetryDisabled() {
+        return this.retryDisabled;
+    }
+
     String tenant() {
         return this.tenant;
     }
@@ -190,6 +195,7 @@ public abstract static class Builder<T extends Builder<T>> {
         Boolean onlySendFailureTelemetry = false;
         Integer connectTimeoutForDefaultHttpClient;
         Integer readTimeoutForDefaultHttpClient;
+        boolean disableInternalRetries;
         private String clientId;
         private Authority authenticationAuthority = createDefaultAADAuthority();
 
@@ -319,6 +325,18 @@ public T readTimeoutForDefaultHttpClient(Integer val) {
             return self();
         }
 
+        /**
+         * The library has a number of policies for retrying HTTP calls in different scenarios.
+         * <p>
+         * This will disable all internal retry behavior, allowing customized retry behavior via your own implementation of {@link IHttpClient}
+         *
+         * @return instance of the Builder on which method was called
+         */
+        public T disableInternalRetries() {
+            disableInternalRetries = true;
+            return self();
+        }
+
         T telemetryConsumer(Consumer<List<HashMap<String, String>>> val) {
             validateNotNull("telemetryConsumer", val);
 
@@ -356,5 +374,16 @@ private static Authority createDefaultAADAuthority() {
         readTimeoutForDefaultHttpClient = builder.readTimeoutForDefaultHttpClient;
         authenticationAuthority = builder.authenticationAuthority;
         clientId = builder.clientId;
+        retryDisabled = builder.disableInternalRetries;
+
+        if (builder.httpClient == null) {
+            httpClient = new DefaultHttpClient(
+                    builder.proxy,
+                    builder.sslSocketFactory,
+                    builder.connectTimeoutForDefaultHttpClient,
+                    builder.readTimeoutForDefaultHttpClient);
+        } else {
+            httpClient = builder.httpClient;
+        }
     }
 }

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

@@ -566,9 +566,7 @@ public T correlationId(String val) {
         super.serviceBundle = new ServiceBundle(
                 builder.executorService,
                 new TelemetryManager(telemetryConsumer, builder.onlySendFailureTelemetry),
-                new HttpHelper(builder.httpClient == null ?
-                        new DefaultHttpClient(builder.proxy, builder.sslSocketFactory, builder.connectTimeoutForDefaultHttpClient, builder.readTimeoutForDefaultHttpClient) :
-                        builder.httpClient)
+                new HttpHelper(this, new DefaultRetryPolicy())
         );
 
         if (aadAadInstanceDiscoveryResponse != null) {

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

@@ -41,7 +41,7 @@ class AuthorizationResponseHandler implements HttpHandler {
     public void handle(HttpExchange httpExchange) throws IOException {
         try {
             if (!httpExchange.getRequestURI().getPath().equalsIgnoreCase("/")) {
-                httpExchange.sendResponseHeaders(200, 0);
+                httpExchange.sendResponseHeaders(HttpStatus.HTTP_OK, 0);
                 return;
             }
             String responseBody = new BufferedReader(new InputStreamReader(
@@ -92,13 +92,13 @@ private void sendErrorResponse(HttpExchange httpExchange, String response) throw
     private void send302Response(HttpExchange httpExchange, String redirectUri) throws IOException {
         Headers responseHeaders = httpExchange.getResponseHeaders();
         responseHeaders.set("Location", redirectUri);
-        httpExchange.sendResponseHeaders(302, 0);
+        httpExchange.sendResponseHeaders(HttpStatus.HTTP_FOUND, 0);
     }
 
     private void send200Response(HttpExchange httpExchange, String response) throws IOException {
         byte[] responseBytes = response.getBytes("UTF-8");
         httpExchange.getResponseHeaders().set("Content-Type", "text/html; charset=UTF-8");
-        httpExchange.sendResponseHeaders(200, responseBytes.length);
+        httpExchange.sendResponseHeaders(HttpStatus.HTTP_OK, responseBytes.length);
         OutputStream os = httpExchange.getResponseBody();
         os.write(responseBytes);
         os.close();

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

@@ -0,0 +1,28 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.microsoft.aad.msal4j;
+
+/**
+ * Default retry policy for most MSAL Java flows
+ */
+class DefaultRetryPolicy implements IRetryPolicy {
+    private static final int RETRY_NUM = 1;
+    private static final int RETRY_DELAY_MS = 1000;
+
+    @Override
+    public boolean isRetryable(IHttpResponse httpResponse) {
+        return HttpStatus.isServerError(httpResponse.statusCode()) &&
+                HttpHelper.getRetryAfterHeader(httpResponse) == null;
+    }
+
+    @Override
+    public int getMaxRetryCount(IHttpResponse httpResponse) {
+        return RETRY_NUM;
+    }
+
+    @Override
+    public int getRetryDelayMs(IHttpResponse httpResponse) {
+        return RETRY_DELAY_MS;
+    }
+}

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

@@ -46,7 +46,7 @@ DeviceCode acquireDeviceCode(String url,
                 this.requestContext(),
                 serviceBundle);
 
-        if (response.statusCode() != HttpHelper.HTTP_STATUS_200) {
+        if (response.statusCode() != HttpStatus.HTTP_OK) {
             throw MsalServiceExceptionFactory.fromHttpResponse(response);
         }
 

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

@@ -18,18 +18,20 @@ class HttpHelper implements IHttpHelper {
 
     private static final Logger log = LoggerFactory.getLogger(HttpHelper.class);
     public static final String RETRY_AFTER_HEADER = "Retry-After";
-    public static final int RETRY_NUM = 2;
-    public static final int RETRY_DELAY_MS = 1000;
-
-    public static final int HTTP_STATUS_200 = 200;
-    public static final int HTTP_STATUS_400 = 400;
-    public static final int HTTP_STATUS_429 = 429;
-    public static final int HTTP_STATUS_500 = 500;
 
     private IHttpClient httpClient;
+    private IRetryPolicy retryPolicy;
+    private boolean retryDisabled;
 
-    HttpHelper(IHttpClient httpClient) {
+    HttpHelper(IHttpClient httpClient, IRetryPolicy retryPolicy) {
         this.httpClient = httpClient;
+        this.retryPolicy = retryPolicy != null ? retryPolicy : new DefaultRetryPolicy();
+    }
+
+    HttpHelper(AbstractApplicationBase application, IRetryPolicy retryPolicy) {
+        this.httpClient = application.httpClient();
+        this.retryDisabled = application.isRetryDisabled();
+        this.retryPolicy = retryPolicy != null ? retryPolicy : new DefaultRetryPolicy();
     }
 
     public IHttpResponse executeHttpRequest(HttpRequest httpRequest,
@@ -141,20 +143,23 @@ private String getRequestThumbprint(RequestContext requestContext) {
         return StringHelper.createSha256Hash(sb.toString());
     }
 
-    boolean isRetryable(IHttpResponse httpResponse) {
-        return httpResponse.statusCode() >= HTTP_STATUS_500 &&
-                getRetryAfterHeader(httpResponse) == null;
-    }
-
     IHttpResponse executeHttpRequestWithRetries(HttpRequest httpRequest, IHttpClient httpClient)
             throws Exception {
-        IHttpResponse httpResponse = null;
-        for (int i = 0; i < RETRY_NUM; i++) {
+        IHttpResponse httpResponse = httpClient.send(httpRequest);
+
+        if (retryDisabled) {
+            return httpResponse;
+        }
+
+        int retryCount = 0;
+        int maxRetries = retryPolicy.getMaxRetryCount(httpResponse);
+
+        while (retryPolicy.isRetryable(httpResponse) && retryCount < maxRetries) {
+            Thread.sleep(retryPolicy.getRetryDelayMs(httpResponse));
+
+            retryCount++;
+
             httpResponse = httpClient.send(httpRequest);
-            if (!isRetryable(httpResponse)) {
-                break;
-            }
-            Thread.sleep(RETRY_DELAY_MS);
         }
 
         return httpResponse;
@@ -180,8 +185,8 @@ private void processThrottlingInstructions(IHttpResponse httpResponse, RequestCo
             Integer retryAfterHeaderVal = getRetryAfterHeader(httpResponse);
             if (retryAfterHeaderVal != null) {
                 expirationTimestamp = System.currentTimeMillis() + retryAfterHeaderVal * 1000;
-            } else if (httpResponse.statusCode() == HTTP_STATUS_429 ||
-                    (httpResponse.statusCode() >= HTTP_STATUS_500)) {
+            } else if (httpResponse.statusCode() == HttpStatus.HTTP_TOO_MANY_REQUESTS ||
+                    (httpResponse.statusCode() >= HttpStatus.HTTP_INTERNAL_ERROR)) {
 
                 expirationTimestamp = System.currentTimeMillis() + ThrottlingCache.DEFAULT_THROTTLING_TIME_SEC * 1000;
             }
@@ -191,7 +196,7 @@ private void processThrottlingInstructions(IHttpResponse httpResponse, RequestCo
         }
     }
 
-    private Integer getRetryAfterHeader(IHttpResponse httpResponse) {
+    static Integer getRetryAfterHeader(IHttpResponse httpResponse) {
 
         if (httpResponse.headers() != null) {
             TreeMap<String, List<String>> headers = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
@@ -279,4 +284,8 @@ private static void verifyReturnedCorrelationId(final HttpRequest httpRequest,
             log.info(msg);
         }
     }
+
+    void setRetryPolicy(IRetryPolicy retryPolicy) {
+        this.retryPolicy = retryPolicy;
+    }
 }

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

@@ -0,0 +1,29 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.microsoft.aad.msal4j;
+
+class HttpStatus {
+
+    static final int HTTP_OK = 200;
+    static final int HTTP_FOUND = 302;
+    static final int HTTP_BAD_REQUEST = 400;
+    static final int HTTP_UNAUTHORIZED = 401;
+    static final int HTTP_NOT_FOUND = 404;
+    static final int HTTP_REQUEST_TIMEOUT = 408;
+    static final int HTTP_GONE = 410;
+    static final int HTTP_TOO_MANY_REQUESTS = 429;
+    static final int HTTP_INTERNAL_ERROR = 500;
+    static final int HTTP_UNAVAILABLE = 503;
+    static final int HTTP_GATEWAY_TIMEOUT = 504;
+
+    /**
+     * Determines if the status code represents a server error (5xx).
+     *
+     * @param code The HTTP status code
+     * @return true if the status code is between 500 and 599, inclusive
+     */
+    static boolean isServerError(int code) {
+        return code >= 500 && code < 600;
+    }
+}

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

@@ -36,6 +36,12 @@ public IMDSManagedIdentitySource(MsalRequest msalRequest,
         super(msalRequest, serviceBundle, ManagedIdentitySourceType.IMDS);
         IEnvironmentVariables environmentVariables = getEnvironmentVariables();
 
+        //IMDS uses a different retry policy than the default used in other MI flows
+        IHttpHelper httpHelper = serviceBundle.getHttpHelper();
+        if (httpHelper instanceof HttpHelper) {
+            ((HttpHelper) httpHelper).setRetryPolicy(new IMDSRetryPolicy());
+        }
+
         if (!StringHelper.isNullOrBlank(environmentVariables.getEnvironmentVariable(Constants.AZURE_POD_IDENTITY_AUTHORITY_HOST))){
             LOG.info(String.format("[Managed Identity] Environment variable AZURE_POD_IDENTITY_AUTHORITY_HOST for IMDS returned endpoint: %s", environmentVariables.getEnvironmentVariable(Constants.AZURE_POD_IDENTITY_AUTHORITY_HOST)));
             try {