PR feedback, and correctly adjust parameters in ADFS username/password scenarios

Author: Avery-Dunn

msal4j-sdk/src/integrationtest/java/infrastructure/SeleniumExtensions.java

@@ -40,7 +40,7 @@ public static WebDriver createDefaultWebDriver() {
         //No visual rendering, remove to see browser window when debugging
         options.addArguments("--headless");
         //Add to avoid issues if your real browser's history/cookies are affecting tests, should not be needed in ADO pipelines
-        //options.addArguments("--incognito");
+        options.addArguments("--incognito");
 
         System.setProperty("webdriver.chrome.driver", "C:/Windows/chromedriver.exe");
         ChromeDriver driver = new ChromeDriver(options);

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

@@ -39,8 +39,7 @@ AuthenticationResult execute() throws Exception {
         }
 
         if (authGrant instanceof OAuthAuthorizationGrant) {
-            msalRequest.msalAuthorizationGrant =
-                    processPasswordGrant((OAuthAuthorizationGrant) authGrant);
+            processPasswordGrant((OAuthAuthorizationGrant) authGrant);
         }
 
         if (authGrant instanceof IntegratedWindowsAuthorizationGrant) {
@@ -74,39 +73,32 @@ private boolean IsUiRequiredCacheSupported() {
                 clientApplication instanceof PublicClientApplication;
     }
 
-    private OAuthAuthorizationGrant processPasswordGrant(
-            OAuthAuthorizationGrant authGrant) throws Exception {
-
-        if (!(authGrant.toParameters().get(GrantConstants.GRANT_TYPE_PARAMETER).get(0).equals(GrantConstants.PASSWORD))) {
-            return authGrant;
-        }
+    private void processPasswordGrant(OAuthAuthorizationGrant authGrant) throws Exception {
 
-        if (msalRequest.application().authenticationAuthority.authorityType != AuthorityType.AAD) {
-            return authGrant;
+        //Additional processing is only needed if it's a password grant with a non-AAD authority
+        if (!(authGrant.getParamValue(GrantConstants.GRANT_TYPE_PARAMETER).equals(GrantConstants.PASSWORD))
+                || msalRequest.application().authenticationAuthority.authorityType != AuthorityType.AAD) {
+            return;
         }
 
         UserDiscoveryResponse userDiscoveryResponse = UserDiscoveryRequest.execute(
-                this.clientApplication.authenticationAuthority.getUserRealmEndpoint(authGrant.toParameters().get("username").get(0)),
+                this.clientApplication.authenticationAuthority.getUserRealmEndpoint(authGrant.getParamValue(GrantConstants.USERNAME_PARAMETER)),
                 msalRequest.headers().getReadonlyHeaderMap(),
                 msalRequest.requestContext(),
                 this.clientApplication.serviceBundle());
 
         if (userDiscoveryResponse.isAccountFederated()) {
             WSTrustResponse response = WSTrustRequest.execute(
                     userDiscoveryResponse.federationMetadataUrl(),
-                    authGrant.toParameters().get(GrantConstants.USERNAME_PARAMETER).get(0),
-                    authGrant.toParameters().get(GrantConstants.PASSWORD_PARAMETER).get(0),
+                    authGrant.getParamValue(GrantConstants.USERNAME_PARAMETER),
+                    authGrant.getParamValue(GrantConstants.PASSWORD_PARAMETER),
                     userDiscoveryResponse.cloudAudienceUrn(),
                     msalRequest.requestContext(),
                     this.clientApplication.serviceBundle(),
                     this.clientApplication.logPii());
 
-            Map<String, List<String>> params = getSAMLAuthGrantParameters(response);
-            params.putAll(authGrant.toParameters());
-
-            authGrant = new OAuthAuthorizationGrant(params, null);
+            authGrant.addAndReplaceParams(getSAMLAuthGrantParameters(response));
         }
-        return authGrant;
     }
 
     private Map<String, List<String>> getSAMLAuthGrantParameters(WSTrustResponse response) {

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

@@ -37,21 +37,35 @@ class OAuthAuthorizationGrant extends AbstractMsalAuthorizationGrant {
         }
     }
 
-        /**
-         * Constructor to create an OAuthAuthorizationGrant
-         *
-         * @param params parameters relevant for the specific authorization grant type
-         * @param scopes additional scopes which will be added to a default set of common scopes
-         * @param claims optional claims
-         */
-        OAuthAuthorizationGrant(Map<String, List<String>> params, Set<String> scopes, ClaimsRequest claims) {
-            this(params, scopes);
-
-            if (claims != null) {
-                this.claims = claims;
-                this.params.put("claims", Collections.singletonList(claims.formatAsJSONString()));
-            }
+    /**
+     * Constructor to create an OAuthAuthorizationGrant
+     *
+     * @param params parameters relevant for the specific authorization grant type
+     * @param scopes additional scopes which will be added to a default set of common scopes
+     * @param claims optional claims
+     */
+    OAuthAuthorizationGrant(Map<String, List<String>> params, Set<String> scopes, ClaimsRequest claims) {
+        this(params, scopes);
+
+        if (claims != null) {
+            this.claims = claims;
+            this.params.put("claims", Collections.singletonList(claims.formatAsJSONString()));
+        }
+    }
+
+    void addAndReplaceParams(Map<String, List<String>> params) {
+        if (params != null) {
+            //putAll() will overwrite existing values if the key already exists in the map
+            this.params.putAll(params);
         }
+    }
+
+    String getParamValue(String paramKey) {
+        if (this.params.containsKey(paramKey)) {
+            return this.params.get(paramKey).get(0);
+        }
+        return null;
+    }
 
     /**
      * Returns an unmodifiable version of the parameters map