Region discovery support (#343)

Avery-Dunn · web-flow · commit a9d3ad3bf05f · 2021-01-29T16:05:12.000-08:00
* Add Azure regional support

* Refactor

* Add logs for success/failure to find regional info

* Extra log
diff --git a/src/main/java/com/microsoft/aad/msal4j/AadInstanceDiscoveryProvider.java b/src/main/java/com/microsoft/aad/msal4j/AadInstanceDiscoveryProvider.java
@@ -3,22 +3,36 @@
 
 package com.microsoft.aad.msal4j;
 
+import com.nimbusds.oauth2.sdk.http.HTTPResponse;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import java.net.URL;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Set;
 import java.util.TreeSet;
+import java.util.Map;
+import java.util.HashMap;
 import java.util.concurrent.ConcurrentHashMap;
 
 class AadInstanceDiscoveryProvider {
 
     private final static String DEFAULT_TRUSTED_HOST = "login.microsoftonline.com";
     private final static String AUTHORIZE_ENDPOINT_TEMPLATE = "https://{host}/{tenant}/oauth2/v2.0/authorize";
     private final static String INSTANCE_DISCOVERY_ENDPOINT_TEMPLATE = "https://{host}/common/discovery/instance";
+    private final static String INSTANCE_DISCOVERY_ENDPOINT_TEMPLATE_WITH_REGION = "https://{region}.{host}/common/discovery/instance";
     private final static String INSTANCE_DISCOVERY_REQUEST_PARAMETERS_TEMPLATE = "?api-version=1.1&authorization_endpoint={authorizeEndpoint}";
+    private final static String REGION_NAME = "REGION_NAME";
+    // For information of the current api-version refer: https://docs.microsoft.com/en-us/azure/virtual-machines/windows/instance-metadata-service#versioning
+    private final static String DEFAULT_API_VERSION = "2020-06-01";
+    private final static String IMDS_ENDPOINT = "https://169.254.169.254/metadata/instance/compute/location?" + DEFAULT_API_VERSION + "&format=text";
 
     final static TreeSet<String> TRUSTED_HOSTS_SET = new TreeSet<>(String.CASE_INSENSITIVE_ORDER);
 
+    private final static Logger log = LoggerFactory.getLogger(HttpHelper.class);
+
     static ConcurrentHashMap<String, InstanceDiscoveryMetadataEntry> cache = new ConcurrentHashMap<>();
 
     static {
@@ -102,26 +116,98 @@ private static String getInstanceDiscoveryEndpoint(String host) {
                 replace("{host}", discoveryHost);
     }
 
+    private static String getInstanceDiscoveryEndpointWithRegion(String host, String region) {
+
+        String discoveryHost = TRUSTED_HOSTS_SET.contains(host) ? host : DEFAULT_TRUSTED_HOST;
+
+        return INSTANCE_DISCOVERY_ENDPOINT_TEMPLATE_WITH_REGION.
+                replace("{region}", region).
+                replace("{host}", discoveryHost);
+    }
+
     private static AadInstanceDiscoveryResponse sendInstanceDiscoveryRequest(URL authorityUrl,
                                                                              MsalRequest msalRequest,
                                                                              ServiceBundle serviceBundle) {
 
-        String instanceDiscoveryRequestUrl = getInstanceDiscoveryEndpoint(authorityUrl.getAuthority()) +
-                INSTANCE_DISCOVERY_REQUEST_PARAMETERS_TEMPLATE.replace("{authorizeEndpoint}",
-                        getAuthorizeEndpoint(authorityUrl.getAuthority(),
-                                Authority.getTenant(authorityUrl, Authority.detectAuthorityType(authorityUrl))));
+        String region = StringHelper.EMPTY_STRING;
+        IHttpResponse httpResponse = null;
+
+        //If the autoDetectRegion parameter in the request is set, attempt to discover the region
+        if (msalRequest.application().autoDetectRegion()) {
+            region = discoverRegion(msalRequest, serviceBundle);
+        }
+
+        //If the region is known, attempt to make instance discovery request with region endpoint
+        if (!region.isEmpty()) {
+            String instanceDiscoveryRequestUrl = getInstanceDiscoveryEndpointWithRegion(authorityUrl.getAuthority(), region) +
+                    formInstanceDiscoveryParameters(authorityUrl);
+
+            httpResponse = httpRequest(instanceDiscoveryRequestUrl, msalRequest.headers().getReadonlyHeaderMap(), msalRequest, serviceBundle);
+        }
+
+        //If the region is unknown or the instance discovery failed at the region endpoint, try the global endpoint
+        if (region.isEmpty() || httpResponse == null || httpResponse.statusCode() != HTTPResponse.SC_OK) {
+            if (!region.isEmpty()) {
+                log.warn("Could not retrieve regional instance discovery metadata, falling back to global endpoint");
+            }
+
+            String instanceDiscoveryRequestUrl = getInstanceDiscoveryEndpoint(authorityUrl.getAuthority()) +
+                    formInstanceDiscoveryParameters(authorityUrl);
+
+            httpResponse = httpRequest(instanceDiscoveryRequestUrl, msalRequest.headers().getReadonlyHeaderMap(), msalRequest, serviceBundle);
+        }
 
+        return JsonHelper.convertJsonToObject(httpResponse.body(), AadInstanceDiscoveryResponse.class);
+    }
+
+    private static String formInstanceDiscoveryParameters(URL authorityUrl) {
+        return INSTANCE_DISCOVERY_REQUEST_PARAMETERS_TEMPLATE.replace("{authorizeEndpoint}",
+                getAuthorizeEndpoint(authorityUrl.getAuthority(),
+                        Authority.getTenant(authorityUrl, Authority.detectAuthorityType(authorityUrl))));
+    }
+
+    private static IHttpResponse httpRequest(String requestUrl, Map<String, String> headers, MsalRequest msalRequest, ServiceBundle serviceBundle) {
         HttpRequest httpRequest = new HttpRequest(
                 HttpMethod.GET,
-                instanceDiscoveryRequestUrl,
-                msalRequest.headers().getReadonlyHeaderMap());
+                requestUrl,
+                headers);
 
-        IHttpResponse httpResponse= HttpHelper.executeHttpRequest(
+        return HttpHelper.executeHttpRequest(
                 httpRequest,
                 msalRequest.requestContext(),
                 serviceBundle);
+    }
 
-        return JsonHelper.convertJsonToObject(httpResponse.body(), AadInstanceDiscoveryResponse.class);
+    private static String discoverRegion(MsalRequest msalRequest, ServiceBundle serviceBundle) {
+
+        //Check if the REGION_NAME environment variable has a value for the region
+        if (System.getenv(REGION_NAME) != null) {
+            log.info("Region found in environment variable: " + System.getenv(REGION_NAME));
+
+            return System.getenv(REGION_NAME);
+        }
+
+        try {
+            //Check the IMDS endpoint to retrieve current region (will only work if application is running in an Azure VM)
+            Map<String, String> headers = new HashMap<>();
+            headers.put("Metadata", "true");
+            IHttpResponse httpResponse = httpRequest(IMDS_ENDPOINT, headers, msalRequest, serviceBundle);
+
+            //If call to IMDS endpoint was successful, return region from response body
+            if (httpResponse.statusCode() == HTTPResponse.SC_OK && !httpResponse.body().isEmpty()) {
+                log.info("Region retrieved from IMDS endpoint: " + httpResponse.body());
+
+                return httpResponse.body();
+            }
+
+            log.warn(String.format("Call to local IMDS failed with status code: %s, or response was empty", httpResponse.statusCode()));
+
+            return StringHelper.EMPTY_STRING;
+        } catch (Exception e) {
+            //IMDS call failed, cannot find region
+            log.warn(String.format("Exception during call to local IMDS endpoint: %s", e.getMessage()));
+            return StringHelper.EMPTY_STRING;
+        }
     }
 
     private static void doInstanceDiscoveryAndCache(URL authorityUrl,
diff --git a/src/main/java/com/microsoft/aad/msal4j/AbstractClientApplicationBase.java b/src/main/java/com/microsoft/aad/msal4j/AbstractClientApplicationBase.java
@@ -96,6 +96,10 @@ abstract class AbstractClientApplicationBase implements IClientApplicationBase {
     @Getter
     private String clientCapabilities;
 
+    @Accessors(fluent = true)
+    @Getter
+    private boolean autoDetectRegion;
+
     @Override
     public CompletableFuture<IAuthenticationResult> acquireToken(AuthorizationCodeParameters parameters) {
 
@@ -292,6 +296,7 @@ abstract static class Builder<T extends Builder<T>> {
         private ITokenCacheAccessAspect tokenCacheAccessAspect;
         private AadInstanceDiscoveryResponse aadInstanceDiscoveryResponse;
         private String clientCapabilities;
+        private boolean autoDetectRegion;
         private Integer connectTimeoutForDefaultHttpClient;
         private Integer readTimeoutForDefaultHttpClient;
 
@@ -573,6 +578,22 @@ public T clientCapabilities(Set<String> capabilities) {
             return self();
         }
 
+        /**
+         * Indicates that the library should attempt to discover the Azure region the application is running in when
+         *  fetching the instance discovery metadata.
+         *
+         * If the region is found, token requests will be sent to the regional ESTS endpoint rather than the global endpoint.
+         * If region information could not be found, the library will fall back to using the global endpoint, which is also
+         *  the default behavior if this value is not set.
+         *
+         * @param val boolean (default is false)
+         * @return instance of the Builder on which method was called
+         */
+        public T autoDetectRegion(boolean val) {
+            autoDetectRegion = val;
+            return self();
+        }
+
         abstract AbstractClientApplicationBase build();
     }
 
@@ -599,6 +620,7 @@ public T clientCapabilities(Set<String> capabilities) {
         tokenCache = new TokenCache(builder.tokenCacheAccessAspect);
         aadAadInstanceDiscoveryResponse = builder.aadInstanceDiscoveryResponse;
         clientCapabilities = builder.clientCapabilities;
+        autoDetectRegion = builder.autoDetectRegion;
 
         if(aadAadInstanceDiscoveryResponse != null){
             AadInstanceDiscoveryProvider.cacheInstanceDiscoveryMetadata(
diff --git a/src/test/java/com/microsoft/aad/msal4j/AadInstanceDiscoveryTest.java b/src/test/java/com/microsoft/aad/msal4j/AadInstanceDiscoveryTest.java
@@ -135,4 +135,55 @@ public void aadInstanceDiscoveryTest_responseSetByDeveloper_invalidJson() throws
                 .aadInstanceDiscoveryResponse(instanceDiscoveryResponse)
                 .build();
     }
+
+    @Test()
+    public void aadInstanceDiscoveryTest_AutoDetectRegion_NoRegionDetected() throws Exception{
+
+        String instanceDiscoveryResponse = TestHelper.readResource(
+                this.getClass(),
+                "/instance_discovery_data/aad_instance_discovery_response_valid.json");
+
+        PublicClientApplication app = PublicClientApplication.builder("client_id")
+                .aadInstanceDiscoveryResponse(instanceDiscoveryResponse)
+                .autoDetectRegion(true)
+                .build();
+
+        AuthorizationCodeParameters parameters = AuthorizationCodeParameters.builder(
+                "code", new URI("http://my.redirect.com")).build();
+
+        MsalRequest msalRequest = new AuthorizationCodeRequest(
+                parameters,
+                app,
+                new RequestContext(app, PublicApi.ACQUIRE_TOKEN_BY_AUTHORIZATION_CODE, parameters));
+
+        URL authority = new URL(app.authority());
+
+        PowerMock.mockStaticPartial(
+                AadInstanceDiscoveryProvider.class,
+                "discoverRegion");
+
+        PowerMock.expectPrivate(
+                AadInstanceDiscoveryProvider.class,
+                "discoverRegion",
+                msalRequest,
+                app.getServiceBundle()).andThrow(new AssertionError()).anyTimes();
+
+        PowerMock.replay(AadInstanceDiscoveryProvider.class);
+
+        InstanceDiscoveryMetadataEntry entry = AadInstanceDiscoveryProvider.getMetadataEntry(
+                authority,
+                false,
+                msalRequest,
+                app.getServiceBundle());
+
+        //Region detection will have been performed in the expected discoverRegion method, but these tests (likely) aren't
+        // being run in an Azure VM and nstance discovery will fall back to the global endpoint (login.microsoftonline.com)
+        Assert.assertEquals(entry.preferredNetwork(), "login.microsoftonline.com");
+        Assert.assertEquals(entry.preferredCache(), "login.windows.net");
+        Assert.assertEquals(entry.aliases().size(), 4);
+        Assert.assertTrue(entry.aliases().contains("login.microsoftonline.com"));
+        Assert.assertTrue(entry.aliases().contains("login.windows.net"));
+        Assert.assertTrue(entry.aliases().contains("login.microsoft.com"));
+        Assert.assertTrue(entry.aliases().contains("sts.windows.net"));
+    }
 }
