expose instanceDiscovery flag

siddhijain · siddhijain · commit c4a4b3bfa83b · 2022-11-23T12:29:55.000-06:00
diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AadInstanceDiscoveryProvider.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AadInstanceDiscoveryProvider.java
@@ -18,22 +18,22 @@
 
 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}:{port}/common/discovery/instance";
-    private final static String INSTANCE_DISCOVERY_REQUEST_PARAMETERS_TEMPLATE = "?api-version=1.1&authorization_endpoint={authorizeEndpoint}";
-    private final static String HOST_TEMPLATE_WITH_REGION = "{region}.r.{host}";
-    private final static String SOVEREIGN_HOST_TEMPLATE_WITH_REGION = "{region}.{host}";
-    private final static String REGION_NAME = "REGION_NAME";
-    private final static int PORT_NOT_SET = -1;
+    private static final String DEFAULT_TRUSTED_HOST = "login.microsoftonline.com";
+    private static final String AUTHORIZE_ENDPOINT_TEMPLATE = "https://{host}/{tenant}/oauth2/v2.0/authorize";
+    private static final String INSTANCE_DISCOVERY_ENDPOINT_TEMPLATE = "https://{host}:{port}/common/discovery/instance";
+    private static final String INSTANCE_DISCOVERY_REQUEST_PARAMETERS_TEMPLATE = "?api-version=1.1&authorization_endpoint={authorizeEndpoint}";
+    private static final String HOST_TEMPLATE_WITH_REGION = "{region}.r.{host}";
+    private static final String SOVEREIGN_HOST_TEMPLATE_WITH_REGION = "{region}.{host}";
+    private static final String REGION_NAME = "REGION_NAME";
+    private static final int PORT_NOT_SET = -1;
     // 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";
+    private static final String DEFAULT_API_VERSION = "2020-06-01";
+    private static final 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);
-    final static TreeSet<String> TRUSTED_SOVEREIGN_HOSTS_SET = new TreeSet<>(String.CASE_INSENSITIVE_ORDER);
+    static final TreeSet<String> TRUSTED_HOSTS_SET = new TreeSet<>(String.CASE_INSENSITIVE_ORDER);
+    static final TreeSet<String> TRUSTED_SOVEREIGN_HOSTS_SET = new TreeSet<>(String.CASE_INSENSITIVE_ORDER);
 
-    private final static Logger log = LoggerFactory.getLogger(HttpHelper.class);
+    private static final Logger log = LoggerFactory.getLogger(AadInstanceDiscoveryProvider.class);
 
     static ConcurrentHashMap<String, InstanceDiscoveryMetadataEntry> cache = new ConcurrentHashMap<>();
 
@@ -67,10 +67,9 @@ static InstanceDiscoveryMetadataEntry getMetadataEntry(URL authorityUrl,
 
             //If region autodetection is enabled and a specific region not already set,
             // set the application's region to the discovered region so that future requests can skip the IMDS endpoint call
-            if (msalRequest.application().azureRegion() == null && msalRequest.application().autoDetectRegion()) {
-                if (detectedRegion != null) {
+            if (msalRequest.application().azureRegion() == null && msalRequest.application().autoDetectRegion()
+                        && null != detectedRegion) {
                     msalRequest.application().azureRegion = detectedRegion;
-                }
             }
             cacheRegionInstanceMetadata(authorityUrl.getHost(), msalRequest.application().azureRegion());
             serviceBundle.getServerSideTelemetry().getCurrentRequest().regionOutcome(
@@ -80,7 +79,16 @@ static InstanceDiscoveryMetadataEntry getMetadataEntry(URL authorityUrl,
         InstanceDiscoveryMetadataEntry result = cache.get(host);
 
         if (result == null) {
-            doInstanceDiscoveryAndCache(authorityUrl, validateAuthority, msalRequest, serviceBundle);
+            if(msalRequest.application().instanceDiscovery()){
+                doInstanceDiscoveryAndCache(authorityUrl, validateAuthority, msalRequest, serviceBundle);
+            } else {
+                // instanceDiscovery flag is set to False. Do no perform instanceDiscovery.
+                cache.putIfAbsent(host, InstanceDiscoveryMetadataEntry.builder().
+                        preferredCache(host).
+                        preferredNetwork(host).
+                        aliases(Collections.singleton(host)).
+                        build());
+            }
         }
 
         return cache.get(host);
diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AbstractClientApplicationBase.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AbstractClientApplicationBase.java
@@ -104,6 +104,10 @@ public abstract class AbstractClientApplicationBase implements IClientApplicatio
     @Getter
     protected String azureRegion;
 
+    @Accessors(fluent = true)
+    @Getter
+    private boolean instanceDiscovery;
+
     @Override
     public CompletableFuture<IAuthenticationResult> acquireToken(AuthorizationCodeParameters parameters) {
 
@@ -325,6 +329,7 @@ public abstract static class Builder<T extends Builder<T>> {
         private String azureRegion;
         private Integer connectTimeoutForDefaultHttpClient;
         private Integer readTimeoutForDefaultHttpClient;
+        private boolean instanceDiscovery = true;
 
         /**
          * Constructor to create instance of Builder of client application
@@ -643,6 +648,30 @@ public T azureRegion(String val) {
             return self();
         }
 
+        /** Historically, MSAL would connect to a central endpoint located at
+            ``https://login.microsoftonline.com`` to acquire some metadata, especially when using an unfamiliar authority.
+        This behavior is known as Instance Discovery.
+        This parameter defaults to true, which enables the Instance Discovery.
+        If you know some authorities which you allow MSAL to operate with as-is,
+        without involving any Instance Discovery, the recommended pattern is::
+        knownAuthorities = frozenset([  # Treat your known authorities as const
+                "https://contoso.com/adfs", "https://login.azs/foo"])
+                ...
+        authority = "https://contoso.com/adfs"  # Assuming your app will use this
+        app1 = PublicClientApplication(
+                    "client_id",
+                    authority=authority,
+                    # Conditionally disable Instance Discovery for known authorities
+                            instance_discovery=authority not in known_authorities,
+                    )
+        If you do not know some authorities beforehand,
+        yet still want MSAL to accept any authority that you will provide,
+        you can use a ``False`` to unconditionally disable Instance Discovery. */
+        public T instanceDiscovery(boolean val) {
+            instanceDiscovery = val;
+            return self();
+        }
+
         abstract AbstractClientApplicationBase build();
     }
 
@@ -671,6 +700,7 @@ public T azureRegion(String val) {
         clientCapabilities = builder.clientCapabilities;
         autoDetectRegion = builder.autoDetectRegion;
         azureRegion = builder.azureRegion;
+        instanceDiscovery = builder.instanceDiscovery;
 
         if (aadAadInstanceDiscoveryResponse != null) {
             AadInstanceDiscoveryProvider.cacheInstanceDiscoveryMetadata(
diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/InstanceDiscoveryMetadataEntry.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/InstanceDiscoveryMetadataEntry.java
@@ -11,6 +11,7 @@
 
 @Accessors(fluent = true)
 @Getter(AccessLevel.PACKAGE)
+@Setter
 @Builder
 @NoArgsConstructor
 @AllArgsConstructor
diff --git a/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/AadInstanceDiscoveryTest.java b/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/AadInstanceDiscoveryTest.java
@@ -8,10 +8,12 @@
 import org.powermock.modules.testng.PowerMockTestCase;
 import org.testng.Assert;
 import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
 import java.net.URI;
 import java.net.URL;
+import java.util.Collections;
 
 @PrepareForTest(AadInstanceDiscoveryProvider.class)
 public class AadInstanceDiscoveryTest extends PowerMockTestCase {
@@ -186,4 +188,137 @@ public void aadInstanceDiscoveryTest_AutoDetectRegion_NoRegionDetected() throws
         Assert.assertTrue(entry.aliases().contains("login.microsoft.com"));
         Assert.assertTrue(entry.aliases().contains("sts.windows.net"));
     }
+
+    @DataProvider(name = "aadClouds")
+    private static Object[][] getAadClouds(){
+        return new Object[][] {{"https://login.microsoftonline.com/common"} , // #Known to Microsoft
+                {"https://private.cloud/foo"}//Private Cloud
+                };
+    }
+
+    @DataProvider(name = "b2cAdfsClouds")
+    private static Object[][] getNonAadClouds(){
+        return new Object[][] {{"https://contoso.com/adfs"}//ADFS
+//                {"https://login.b2clogin.com/contoso/b2c_policy"},//B2C
+                };
+    }
+
+    /**
+     * when instance_discovery flag is set to true (by default), an instance_discovery is performed for authorityType = AAD and
+     *     hence, an exception is thrown while making a call to getMetaDataEntry() if instanceDiscoveryResponse is not mocked.
+    */
+    @Test(  dataProvider = "aadClouds",
+            expectedExceptions = StringIndexOutOfBoundsException.class)
+    public void aad_instance_discovery_true(String authority) throws Exception {
+
+        PublicClientApplication app = PublicClientApplication.builder("client_id")
+                .authority(authority)
+                .build();
+
+        AuthorizationCodeParameters parameters = AuthorizationCodeParameters.builder(
+                "code", new URI("http://my.redirect.com"))
+                .scopes(Collections.singleton("scope")).build();
+
+        MsalRequest msalRequest = new AuthorizationCodeRequest(
+                parameters,
+                app,
+                new RequestContext(app, PublicApi.ACQUIRE_TOKEN_BY_AUTHORIZATION_CODE, parameters));
+
+        URL authorityURL = new URL(authority);
+
+       AadInstanceDiscoveryProvider.getMetadataEntry(
+                authorityURL,
+                false,
+                msalRequest,
+                app.getServiceBundle());
+
+    }
+
+    /**
+     * when instance_discovery flag is set to true (by default), an instance_discovery is NOT performed for b2c.
+     */
+    @Test(  dataProvider = "b2cAdfsClouds")
+    public void b2c_adfs_instance_discovery_true(String authority) throws Exception {
+
+        PublicClientApplication app = PublicClientApplication.builder("client_id")
+                .authority(authority)
+                .build();
+
+        AuthorizationCodeParameters parameters = AuthorizationCodeParameters.builder(
+                        "code", new URI("http://my.redirect.com"))
+                .scopes(Collections.singleton("scope")).build();
+
+        MsalRequest msalRequest = new AuthorizationCodeRequest(
+                parameters,
+                app,
+                new RequestContext(app, PublicApi.ACQUIRE_TOKEN_BY_AUTHORIZATION_CODE, parameters));
+
+        URL authorityURL = new URL(authority);
+
+        AadInstanceDiscoveryProvider.getMetadataEntry(
+                authorityURL,
+                false,
+                msalRequest,
+                app.getServiceBundle());
+    }
+
+    @Test (dataProvider = "aadClouds")
+    /**
+     * when instance_discovery flag is set to false, instance_discovery is not performed and hence,
+     * no exception is thrown while making a call to getMetaDataEntry() even when instanceDiscoveryResponse is not mocked.
+     */
+    public void aad_instance_discovery_false(String authority) throws Exception{
+
+        PublicClientApplication app = PublicClientApplication.builder("client_id")
+                .authority(authority)
+                .instanceDiscovery(false)
+                .build();
+
+        AuthorizationCodeParameters parameters = AuthorizationCodeParameters.builder(
+                        "code", new URI("http://my.redirect.com"))
+                .scopes(Collections.singleton("scope")).build();
+
+        MsalRequest msalRequest = new AuthorizationCodeRequest(
+                parameters,
+                app,
+                new RequestContext(app, PublicApi.ACQUIRE_TOKEN_BY_AUTHORIZATION_CODE, parameters));
+
+        URL authorityURL = new URL(authority);
+
+        AadInstanceDiscoveryProvider.getMetadataEntry(
+                authorityURL,
+                false,
+                msalRequest,
+                app.getServiceBundle());
+    }
+
+    @Test (dataProvider = "b2cAdfsClouds")
+    /**
+     * when instance_discovery flag is set to true, instance_discovery is not performed and hence,
+     * no exception is thrown while making a call to getMetaDataEntry() even when instanceDiscoveryResponse is not mocked.
+     */
+    public void b2c_adfs_instance_discovery_false(String authority) throws Exception{
+
+        PublicClientApplication app = PublicClientApplication.builder("client_id")
+                .authority(authority)
+                .instanceDiscovery(false)
+                .build();
+
+        AuthorizationCodeParameters parameters = AuthorizationCodeParameters.builder(
+                        "code", new URI("http://my.redirect.com"))
+                .scopes(Collections.singleton("scope")).build();
+
+        MsalRequest msalRequest = new AuthorizationCodeRequest(
+                parameters,
+                app,
+                new RequestContext(app, PublicApi.ACQUIRE_TOKEN_BY_AUTHORIZATION_CODE, parameters));
+
+        URL authorityURL = new URL(authority);
+
+        AadInstanceDiscoveryProvider.getMetadataEntry(
+                authorityURL,
+                false,
+                msalRequest,
+                app.getServiceBundle());
+    }
 }
