Get callerSessionDataKey for success request

janakamarasena · Oct 24, 2023 · 7bb0b7f · 7bb0b7f
1 parent 67f1aea
commit 7bb0b7f
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 4 deletions.
diff --git a/.../org/wso2/carbon/identity/application/authentication/framework/AuthenticationService.java b/.../org/wso2/carbon/identity/application/authentication/framework/AuthenticationService.java
@@ -92,7 +92,7 @@ private AuthServiceResponse processCommonAuthResponse(AuthServiceRequestWrapper
  AuthServiceResponse authServiceResponse = new AuthServiceResponse();
 
  if (isAuthFlowSuccessful(request)) {
- handleSuccessAuthResponse(request, authServiceResponse);
+ handleSuccessAuthResponse(request, response, authServiceResponse);
  } else if (isAuthFlowFailed(request, response)) {
  handleFailedAuthResponse(request, response, authServiceResponse);
  } else if (isAuthFlowIncomplete(request)) {
@@ -124,9 +124,10 @@ private void handleIntermediateAuthResponse(AuthServiceRequestWrapper request, A
  authServiceResponse.setData(responseData);
  }
 
- private void handleSuccessAuthResponse(AuthServiceRequestWrapper request, AuthServiceResponse authServiceResponse) {
+ private void handleSuccessAuthResponse(AuthServiceRequestWrapper request, AuthServiceResponseWrapper response,
+ AuthServiceResponse authServiceResponse) throws AuthServiceException {
 
- authServiceResponse.setSessionDataKey(request.getSessionDataKey());
+ authServiceResponse.setSessionDataKey(getFlowCompletionSessionDataKey(request, response));
  authServiceResponse.setFlowStatus(AuthServiceConstants.FlowStatus.SUCCESS_COMPLETED);
  }
 
@@ -272,8 +273,19 @@ private boolean isSentToRetryPageOnMissingContext(AuthServiceRequestWrapper requ
  Boolean.TRUE.equals(request.getAttribute(FrameworkConstants.IS_SENT_TO_RETRY))) {
  Map<String, String> queryParams = AuthServiceUtils.extractQueryParams(response.getRedirectURL());
  return StringUtils.equals(queryParams.get(FrameworkConstants.STATUS_PARAM),
-  FrameworkConstants.ERROR_STATUS_AUTH_CONTEXT_NULL);
+ FrameworkConstants.ERROR_STATUS_AUTH_CONTEXT_NULL);
  }
  return false;
  }
+
+ private String getFlowCompletionSessionDataKey(AuthServiceRequestWrapper request,
+ AuthServiceResponseWrapper response) throws AuthServiceException {
+
+ String completionSessionDataKey = (String) request.getAttribute(FrameworkConstants.SESSION_DATA_KEY);
+ if (StringUtils.isBlank(completionSessionDataKey)) {
+ completionSessionDataKey = response.getSessionDataKey();
+ }
+
+ return completionSessionDataKey;
+ }
 }
diff --git a/...y/application/authentication/framework/model/auth/service/AuthServiceResponseWrapper.java b/...y/application/authentication/framework/model/auth/service/AuthServiceResponseWrapper.java
@@ -20,6 +20,7 @@
 
 import org.wso2.carbon.identity.application.authentication.framework.exception.auth.service.AuthServiceException;
 import org.wso2.carbon.identity.application.authentication.framework.model.CommonAuthResponseWrapper;
+import org.wso2.carbon.identity.application.authentication.framework.util.FrameworkConstants;
 import org.wso2.carbon.identity.application.authentication.framework.util.auth.service.AuthServiceConstants;
 import org.wso2.carbon.identity.application.authentication.framework.util.auth.service.AuthServiceUtils;
 
@@ -65,6 +66,18 @@ public boolean isErrorResponse() throws AuthServiceException {
  return Boolean.parseBoolean(queryParams.get(AuthServiceConstants.AUTH_FAILURE_PARAM));
  }
 
+ /**
+ * Get the sessionDataKey related to the authentication flow.
+ *
+ * @return String of sessionDataKey.
+ * @throws AuthServiceException
+ */
+ public String getSessionDataKey() throws AuthServiceException {
+
+ Map<String, String> queryParams = AuthServiceUtils.extractQueryParams(getRedirectURL());
+ return queryParams.get(FrameworkConstants.SESSION_DATA_KEY);
+ }
+
  @Override
  public String getRedirectURL() {