Merge pull request #6021 from RushanNanayakkara/fix_branding_not_appl…

…ied_to_mfa Fix branding not applied to mfa
wso2 · Oct 14, 2024 · 3a79f0b · 3a79f0b
2 parents c891b68 + 2b6aacb
commit 3a79f0b
Showing 1 changed file with 38 additions and 0 deletions.
diff --git a/.../application/authentication/framework/handler/request/impl/DefaultRequestCoordinator.java b/.../application/authentication/framework/handler/request/impl/DefaultRequestCoordinator.java
@@ -75,9 +75,11 @@
 import java.io.IOException;
 import java.io.Serializable;
 import java.io.UnsupportedEncodingException;
+import java.net.URI;
 import java.net.URISyntaxException;
 import java.net.URLDecoder;
 import java.net.URLEncoder;
+import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Enumeration;
@@ -96,6 +98,7 @@
 import static org.wso2.carbon.identity.application.authentication.framework.util.FrameworkConstants.ACCOUNT_DISABLED_CLAIM_URI;
 import static org.wso2.carbon.identity.application.authentication.framework.util.FrameworkConstants.ACCOUNT_LOCKED_CLAIM_URI;
 import static org.wso2.carbon.identity.application.authentication.framework.util.FrameworkConstants.ACCOUNT_UNLOCK_TIME_CLAIM;
+import static org.wso2.carbon.identity.application.authentication.framework.util.FrameworkConstants.AUTHENTICATOR;
 import static org.wso2.carbon.identity.application.authentication.framework.util.FrameworkConstants.AnalyticsAttributes.SESSION_ID;
 import static org.wso2.carbon.identity.application.authentication.framework.util.FrameworkConstants.BACK_TO_FIRST_STEP;
 import static org.wso2.carbon.identity.application.authentication.framework.util.FrameworkConstants.ERROR_DESCRIPTION_APP_DISABLED;
@@ -394,6 +397,10 @@ && isStepHasMultiOption(context)
 
                 if (!context.isLogoutRequest()) {
                     FrameworkUtils.getAuthenticationRequestHandler().handle(request, responseWrapper, context);
+                    // Adding spIp param to the redirect URL if the authenticator is not the organization authenticator.
+                    if (!ORGANIZATION_AUTHENTICATOR.equals(request.getParameter(AUTHENTICATOR))) {
+                        addServiceProviderIdToRedirectUrl(responseWrapper, context);
+                    }
                 } else {
                     FrameworkUtils.getLogoutRequestHandler().handle(request, responseWrapper, context);
                 }
@@ -1326,4 +1333,35 @@ private boolean isStepHasMultiOption(AuthenticationContext context) {
         }
         return false;
     }
+
+    private void addServiceProviderIdToRedirectUrl(CommonAuthResponseWrapper responseWrapper,
+                                                   AuthenticationContext context) {
+
+        if (responseWrapper == null || context == null) {
+            return;
+        }
+        try {
+            String redirectURL = responseWrapper.getRedirectURL();
+            String serviceProviderID = context.getServiceProviderResourceId();
+            if (StringUtils.isNotBlank(redirectURL) && StringUtils.isNotBlank(serviceProviderID)) {
+                URI uri = new URI(redirectURL);
+                String query = uri.getRawQuery();
+                if (StringUtils.isNotBlank(query)) {
+                    if (!query.contains(FrameworkConstants.REQUEST_PARAM_SP_UUID + "=")) {
+                        redirectURL = redirectURL + "&" + FrameworkConstants.REQUEST_PARAM_SP_UUID
+                                + "=" + URLEncoder.encode(serviceProviderID,
+                                StandardCharsets.UTF_8.name());
+                    }
+                } else {
+                    redirectURL = redirectURL + "?" + FrameworkConstants.REQUEST_PARAM_SP_UUID
+                            + "=" + URLEncoder.encode(serviceProviderID, StandardCharsets.UTF_8.name());
+                }
+                responseWrapper.sendRedirect(redirectURL);
+            }
+        } catch (URISyntaxException | IOException e) {
+            // No need to break the flow due to this error since added spId to redirect URL is used only
+            // for branding purposes.
+            log.debug("Error while adding spId to redirect URL.");
+        }
+    }
 }