Add disable instance discovery (Azure#33225)

billwert · web-flow · commit 6ddb120f8644 · 2023-02-01T20:34:33.000Z
Add method for disabling AAD instance discovery. Fixes Azure#27888 # All SDK Contribution checklist: - [x] **The pull request does not introduce [breaking changes]** - [x] **CHANGELOG is updated for new features, bug fixes or other significant changes.** - [x] **I have read the [contribution guidelines](https://github.com/Azure/azure-sdk-for-java/blob/main/CONTRIBUTING.md).** ## [General Guidelines and Best Practices](https://github.com/Azure/azure-sdk-for-java/blob/main/CONTRIBUTING.md#developer-guide) - [x] Title of the pull request is clear and informative. - [x] There are a small number of commits, each of which have an informative message. This means that previously merged commits do not appear in the history of the PR. For more information on cleaning up the commits in your PR, [see this page](https://github.com/Azure/azure-powershell/blob/master/documentation/development-docs/cleaning-up-commits.md). ### [Testing Guidelines](https://github.com/Azure/azure-sdk-for-java/blob/main/CONTRIBUTING.md#building-and-unit-testing) - [x] Pull request includes test coverage for the included changes.
diff --git a/eng/versioning/external_dependencies.txt b/eng/versioning/external_dependencies.txt
@@ -197,7 +197,7 @@ com.microsoft.azure:azure-mgmt-resources;1.3.0
 com.microsoft.azure:azure-mgmt-search;1.24.1
 com.microsoft.azure:azure-mgmt-storage;1.3.0
 com.microsoft.azure:azure-storage;8.0.0
-com.microsoft.azure:msal4j;1.13.3
+com.microsoft.azure:msal4j;1.13.4
 com.microsoft.azure:msal4j-persistence-extension;1.1.0
 com.sun.activation:jakarta.activation;1.2.2
 io.opentelemetry:opentelemetry-api;1.20.0
diff --git a/sdk/eventhubs/microsoft-azure-eventhubs-eph/pom.xml b/sdk/eventhubs/microsoft-azure-eventhubs-eph/pom.xml
@@ -64,7 +64,7 @@
     <dependency>
       <groupId>com.microsoft.azure</groupId>
       <artifactId>msal4j</artifactId>
-      <version>1.13.3</version> <!-- {x-version-update;com.microsoft.azure:msal4j;external_dependency} -->
+      <version>1.13.4</version> <!-- {x-version-update;com.microsoft.azure:msal4j;external_dependency} -->
       <scope>test</scope>
     </dependency>
     <dependency>
diff --git a/sdk/eventhubs/microsoft-azure-eventhubs-extensions/pom.xml b/sdk/eventhubs/microsoft-azure-eventhubs-extensions/pom.xml
@@ -68,7 +68,7 @@
       <dependency>
         <groupId>com.microsoft.azure</groupId>
         <artifactId>msal4j</artifactId>
-        <version>1.13.3</version> <!-- {x-version-update;com.microsoft.azure:msal4j;external_dependency} -->
+        <version>1.13.4</version> <!-- {x-version-update;com.microsoft.azure:msal4j;external_dependency} -->
         <scope>test</scope>
       </dependency>
       <dependency>
diff --git a/sdk/eventhubs/microsoft-azure-eventhubs/pom.xml b/sdk/eventhubs/microsoft-azure-eventhubs/pom.xml
@@ -77,7 +77,7 @@
     <dependency>
       <groupId>com.microsoft.azure</groupId>
       <artifactId>msal4j</artifactId>
-      <version>1.13.3</version> <!-- {x-version-update;com.microsoft.azure:msal4j;external_dependency} -->
+      <version>1.13.4</version> <!-- {x-version-update;com.microsoft.azure:msal4j;external_dependency} -->
       <scope>test</scope>
     </dependency>
     <dependency>
diff --git a/sdk/identity/azure-identity/CHANGELOG.md b/sdk/identity/azure-identity/CHANGELOG.md
@@ -16,6 +16,7 @@
 
 ### Features Added
 - Added support to configure `clientOptions`, `httpLogOptions`, `retryPolicy`, `retryOptions` and `addPolicy` on Identity credentials.
+- Added support to disable instance discovery on AAD credentials.
 
 ## 1.7.3 (2023-01-06)
 
diff --git a/sdk/identity/azure-identity/pom.xml b/sdk/identity/azure-identity/pom.xml
@@ -43,7 +43,7 @@
     <dependency>
       <groupId>com.microsoft.azure</groupId>
       <artifactId>msal4j</artifactId>
-      <version>1.13.3</version> <!-- {x-version-update;com.microsoft.azure:msal4j;external_dependency} -->
+      <version>1.13.4</version> <!-- {x-version-update;com.microsoft.azure:msal4j;external_dependency} -->
     </dependency>
     <dependency>
       <groupId>com.microsoft.azure</groupId>
@@ -122,7 +122,7 @@
           <rules>
             <bannedDependencies>
               <includes>
-                <include>com.microsoft.azure:msal4j:[1.13.3]</include> <!-- {x-include-update;com.microsoft.azure:msal4j;external_dependency} -->
+                <include>com.microsoft.azure:msal4j:[1.13.4]</include> <!-- {x-include-update;com.microsoft.azure:msal4j;external_dependency} -->
                 <include>com.microsoft.azure:msal4j-persistence-extension:[1.1.0]</include> <!-- {x-include-update;com.microsoft.azure:msal4j-persistence-extension;external_dependency} -->
                 <include>net.java.dev.jna:jna-platform:[5.6.0]</include> <!-- {x-include-update;net.java.dev.jna:jna-platform;external_dependency} -->
                 <include>org.linguafranca.pwdb:KeePassJava2:[2.1.4]</include> <!-- {x-include-update;org.linguafranca.pwdb:KeePassJava2;external_dependency} -->
diff --git a/sdk/identity/azure-identity/src/main/java/com/azure/identity/AadCredentialBuilderBase.java b/sdk/identity/azure-identity/src/main/java/com/azure/identity/AadCredentialBuilderBase.java
@@ -107,4 +107,16 @@ public T additionallyAllowedTenants(List<String> additionallyAllowedTenants) {
         identityClientOptions.setAdditionallyAllowedTenants(IdentityUtil.resolveAdditionalTenants(additionallyAllowedTenants));
         return (T) this;
     }
+
+    /**
+     * Disables instance discovery.
+     *
+     * @return An updated instance of this builder with instance discovery disabled.
+     */
+    @SuppressWarnings("unchecked")
+
+    public T disableInstanceDiscovery() {
+        this.identityClientOptions.disableInstanceDisovery();
+        return (T) this;
+    }
 }
diff --git a/sdk/identity/azure-identity/src/main/java/com/azure/identity/implementation/IdentityClient.java b/sdk/identity/azure-identity/src/main/java/com/azure/identity/implementation/IdentityClient.java
@@ -218,7 +218,8 @@ public Mono<MsalToken> authenticateWithIntelliJ(TokenRequestContext request) {
                     ConfidentialClientApplication.Builder applicationBuilder =
                         ConfidentialClientApplication.builder(spDetails.get("client"),
                             ClientCredentialFactory.createFromSecret(spDetails.get("key")))
-                            .authority(authorityUrl);
+                            .authority(authorityUrl)
+                            .instanceDiscovery(options.getInstanceDiscovery());
 
                     // If http pipeline is available, then it should override the proxy options if any configured.
                     if (httpPipelineAdapter != null) {
diff --git a/sdk/identity/azure-identity/src/main/java/com/azure/identity/implementation/IdentityClientBase.java b/sdk/identity/azure-identity/src/main/java/com/azure/identity/implementation/IdentityClientBase.java
@@ -206,7 +206,7 @@ ConfidentialClientApplication getConfidentialClient() {
         ConfidentialClientApplication.Builder applicationBuilder =
             ConfidentialClientApplication.builder(clientId, credential);
         try {
-            applicationBuilder = applicationBuilder.authority(authorityUrl);
+            applicationBuilder = applicationBuilder.authority(authorityUrl).instanceDiscovery(options.getInstanceDiscovery());
         } catch (MalformedURLException e) {
             throw LOGGER.logExceptionAsWarning(new IllegalStateException(e));
         }
@@ -260,7 +260,7 @@ PublicClientApplication getPublicClient(boolean sharedTokenCacheCredential) {
             + tenantId;
         PublicClientApplication.Builder builder = PublicClientApplication.builder(clientId);
         try {
-            builder = builder.authority(authorityUrl);
+            builder = builder.authority(authorityUrl).instanceDiscovery(options.getInstanceDiscovery());
         } catch (MalformedURLException e) {
             throw LOGGER.logExceptionAsWarning(new IllegalStateException(e));
         }
diff --git a/sdk/identity/azure-identity/src/main/java/com/azure/identity/implementation/IdentityClientOptions.java b/sdk/identity/azure-identity/src/main/java/com/azure/identity/implementation/IdentityClientOptions.java
@@ -68,6 +68,7 @@ public final class IdentityClientOptions implements Cloneable {
     private RetryPolicy retryPolicy;
     private List<HttpPipelinePolicy> perCallPolicies;
     private List<HttpPipelinePolicy> perRetryPolicies;
+    private boolean instanceDiscovery;
 
     /**
      * Creates an instance of IdentityClientOptions with default settings.
@@ -83,6 +84,7 @@ public IdentityClientOptions() {
         additionallyAllowedTenants = new HashSet<>();
         regionalAuthority = RegionalAuthority.fromString(
             configuration.get(Configuration.PROPERTY_AZURE_REGIONAL_AUTHORITY_NAME));
+        instanceDiscovery = true;
     }
 
     /**
@@ -654,6 +656,25 @@ IdentityClientOptions setPerRetryPolicies(List<HttpPipelinePolicy> perRetryPolic
         return this;
     }
 
+    /**
+     * Disable instance discovery. Instance discovery is acquiring metadata about an authority from https://login.microsoft.com
+     * to validate that authority. This may need to be disabled in private cloud or ADFS scenarios.
+     *
+     * @return the updated client options
+     */
+    public IdentityClientOptions disableInstanceDisovery() {
+        this.instanceDiscovery = false;
+        return this;
+    }
+
+    /**
+     * Gets the instance discovery policy.
+     * @return boolean indicating if instance discovery is enabled.
+     */
+    public boolean getInstanceDiscovery() {
+        return this.instanceDiscovery;
+    }
+
     /**
      * Loads the details from the specified Configuration Store.
      */
@@ -669,7 +690,7 @@ private void loadFromConfiguration(Configuration configuration) {
     }
 
     public IdentityClientOptions clone() {
-        return new IdentityClientOptions()
+        IdentityClientOptions clone =  new IdentityClientOptions()
             .setAdditionallyAllowedTenants(this.additionallyAllowedTenants)
             .setAllowUnencryptedCache(this.allowUnencryptedCache)
             .setHttpClient(this.httpClient)
@@ -697,5 +718,9 @@ public IdentityClientOptions clone() {
             .setRetryPolicy(this.retryPolicy)
             .setPerCallPolicies(this.perCallPolicies)
             .setPerRetryPolicies(this.perRetryPolicies);
+        if (!getInstanceDiscovery()) {
+            clone.disableInstanceDisovery();
+        }
+        return clone;
     }
 }
diff --git a/sdk/identity/azure-identity/src/test/java/com/azure/identity/implementation/IdentityClientOptionsTest.java b/sdk/identity/azure-identity/src/test/java/com/azure/identity/implementation/IdentityClientOptionsTest.java
@@ -38,4 +38,11 @@ public void testCustomAuthorityHost() {
         identityClientOptions.setAuthorityHost(authorityHost);
         Assert.assertEquals(authorityHost, identityClientOptions.getAuthorityHost());
     }
+
+    @Test
+    public void testDisableInstanceDiscovery() {
+        IdentityClientOptions identityClientOptions = new IdentityClientOptions();
+        identityClientOptions.disableInstanceDisovery();
+        Assert.assertFalse(identityClientOptions.getInstanceDiscovery());
+    }
 }
diff --git a/sdk/identity/azure-identity/src/test/java/com/azure/identity/implementation/IdentityClientTests.java b/sdk/identity/azure-identity/src/test/java/com/azure/identity/implementation/IdentityClientTests.java
@@ -392,7 +392,7 @@ public void testBrowserAuthenicationCodeFlow() throws Exception {
         IdentityClientOptions options = new IdentityClientOptions();
         IdentityClient client = new IdentityClientBuilder().tenantId(TENANT_ID).clientId(CLIENT_ID).identityClientOptions(options).build();
         // mock
-        mocForBrowserAuthenticationCodeFlow(token, request, expiresOn, () -> {
+        mockForBrowserAuthenticationCodeFlow(token, request, expiresOn, () -> {
             // test
             StepVerifier.create(client.authenticateWithBrowserInteraction(request, 4567, null, null))
                 .expectNextMatches(accessToken -> token.equals(accessToken.getToken())
@@ -476,6 +476,7 @@ private void mockForClientSecret(String secret, TokenRequestContext request, Str
         try (MockedStatic<ConfidentialClientApplication> staticConfidentialClientApplicationMock = mockStatic(ConfidentialClientApplication.class); MockedConstruction<ConfidentialClientApplication.Builder> confidentialClientApplicationBuilderMock = mockConstruction(ConfidentialClientApplication.Builder.class, (builder, context) -> {
 
             when(builder.authority(any())).thenReturn(builder);
+            when(builder.instanceDiscovery(anyBoolean())).thenReturn(builder);
             when(builder.httpClient(any())).thenReturn(builder);
             ConfidentialClientApplication application = Mockito.mock(ConfidentialClientApplication.class);
             when(application.acquireToken(any(ClientCredentialParameters.class))).thenAnswer(invocation -> {
@@ -504,6 +505,7 @@ private void mockForClientCertificate(TokenRequestContext request, String access
 
         try (MockedStatic<ConfidentialClientApplication> staticConfidentialClientApplicationMock = mockStatic(ConfidentialClientApplication.class); MockedConstruction<ConfidentialClientApplication.Builder> confidentialClientApplicationBuilderMock = mockConstruction(ConfidentialClientApplication.Builder.class, (builder, context) -> {
             when(builder.authority(any())).thenReturn(builder);
+            when(builder.instanceDiscovery(anyBoolean())).thenReturn(builder);
             when(builder.httpClient(any())).thenReturn(builder);
             ConfidentialClientApplication application = Mockito.mock(ConfidentialClientApplication.class);
             when(application.acquireToken(any(ClientCredentialParameters.class))).thenAnswer(invocation -> {
@@ -546,6 +548,7 @@ private void mockForDeviceCodeFlow(TokenRequestContext request, String accessTok
                 return TestUtils.getMockAuthenticationResult(accessToken, expiresOn);
             });
             when(builder.build()).thenReturn(application);
+            when(builder.instanceDiscovery(anyBoolean())).thenReturn(builder);
         })) {
             test.run();
             Assert.assertNotNull(publicClientApplicationMock);
@@ -571,6 +574,7 @@ private void mockForClientPemCertificate(String accessToken, TokenRequestContext
                 });
                 when(builder.build()).thenReturn(application);
                 when(builder.authority(any())).thenReturn(builder);
+                when(builder.instanceDiscovery(anyBoolean())).thenReturn(builder);
                 when(builder.httpClient(any())).thenReturn(builder);
             })
         )  {
@@ -647,7 +651,7 @@ private void mockForIMDSCodeFlow(String endpoint, String tokenJson, Runnable tes
         }
     }
 
-    private void mocForBrowserAuthenticationCodeFlow(String token, TokenRequestContext request, OffsetDateTime expiresOn, Runnable test) throws Exception {
+    private void mockForBrowserAuthenticationCodeFlow(String token, TokenRequestContext request, OffsetDateTime expiresOn, Runnable test) throws Exception {
         try (MockedConstruction<PublicClientApplication.Builder> publicClientApplicationMock = mockConstruction(PublicClientApplication.Builder.class, (builder, context) -> {
             PublicClientApplication application = Mockito.mock(PublicClientApplication.class);
             when(application.acquireToken(any(InteractiveRequestParameters.class))).thenAnswer(invocation -> {
@@ -660,6 +664,7 @@ private void mocForBrowserAuthenticationCodeFlow(String token, TokenRequestConte
             });
             when(builder.build()).thenReturn(application);
             when(builder.authority(any())).thenReturn(builder);
+            when(builder.instanceDiscovery(anyBoolean())).thenReturn(builder);
             when(builder.httpClient(any())).thenReturn(builder);
         })) {
             test.run();
@@ -691,6 +696,7 @@ private void mockForAuthorizationCodeFlow(String token1, TokenRequestContext req
             });
             when(builder.build()).thenReturn(application);
             when(builder.authority(any())).thenReturn(builder);
+            when(builder.instanceDiscovery(anyBoolean())).thenReturn(builder);
             when(builder.httpClient(any())).thenReturn(builder);
         })) {
             test.run();
@@ -711,6 +717,7 @@ private void mockForUsernamePasswordCodeFlow(String token, TokenRequestContext r
             });
             when(builder.build()).thenReturn(application);
             when(builder.authority(any())).thenReturn(builder);
+            when(builder.instanceDiscovery(anyBoolean())).thenReturn(builder);
             when(builder.httpClient(any())).thenReturn(builder);
         })) {
             test.run();
@@ -731,6 +738,7 @@ private void mockForUserRefreshTokenFlow(String token, TokenRequestContext reque
             });
             when(builder.build()).thenReturn(application);
             when(builder.authority(any())).thenReturn(builder);
+            when(builder.instanceDiscovery(anyBoolean())).thenReturn(builder);
             when(builder.httpClient(any())).thenReturn(builder);
         })) {
             test.run();
