From 6b8753ac60f4389e82c6ea24e6ab69cf0c4a4afb Mon Sep 17 00:00:00 2001 From: Janak Amarasena Date: Sun, 22 Oct 2023 16:11:11 +0530 Subject: [PATCH] Change sessionDataKey retrieval --- .../framework/AuthenticationService.java | 48 +++++++++---------- .../impl/DefaultRequestCoordinator.java | 6 ++- .../service/AuthServiceRequestWrapper.java | 22 +++++++++ .../service/AuthServiceResponseWrapper.java | 13 ----- .../framework/util/FrameworkConstants.java | 4 ++ .../framework/util/FrameworkUtils.java | 1 + .../framework/AuthenticationServiceTest.java | 2 + 7 files changed, 58 insertions(+), 38 deletions(-) diff --git a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/AuthenticationService.java b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/AuthenticationService.java index e807f426ab52..e631c9a04f2d 100644 --- a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/AuthenticationService.java +++ b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/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, response, authServiceResponse); + handleSuccessAuthResponse(request, authServiceResponse); } else if (isAuthFlowFailed(request, response)) { handleFailedAuthResponse(request, response, authServiceResponse); } else if (isAuthFlowIncomplete(request)) { @@ -107,7 +107,7 @@ private AuthServiceResponse processCommonAuthResponse(AuthServiceRequestWrapper private void handleIntermediateAuthResponse(AuthServiceRequestWrapper request, AuthServiceResponseWrapper response, AuthServiceResponse authServiceResponse) throws AuthServiceException { - authServiceResponse.setSessionDataKey(response.getSessionDataKey()); + authServiceResponse.setSessionDataKey(request.getSessionDataKey()); authServiceResponse.setFlowStatus(AuthServiceConstants.FlowStatus.INCOMPLETE); AuthServiceResponseData responseData = new AuthServiceResponseData(); boolean isMultiOptionsResponse = request.isMultiOptionsResponse(); @@ -124,10 +124,9 @@ private void handleIntermediateAuthResponse(AuthServiceRequestWrapper request, A authServiceResponse.setData(responseData); } - private void handleSuccessAuthResponse(AuthServiceRequestWrapper request, AuthServiceResponseWrapper response, - AuthServiceResponse authServiceResponse) throws AuthServiceException { + private void handleSuccessAuthResponse(AuthServiceRequestWrapper request, AuthServiceResponse authServiceResponse) { - authServiceResponse.setSessionDataKey(getFlowCompletionSessionDataKey(request, response)); + authServiceResponse.setSessionDataKey(request.getSessionDataKey()); authServiceResponse.setFlowStatus(AuthServiceConstants.FlowStatus.SUCCESS_COMPLETED); } @@ -137,15 +136,15 @@ private void handleFailedAuthResponse(AuthServiceRequestWrapper request, AuthSer String errorCode = null; String errorMessage = null; if (request.isAuthFlowConcluded()) { - authServiceResponse.setSessionDataKey(getFlowCompletionSessionDataKey(request, response)); + authServiceResponse.setSessionDataKey(request.getSessionDataKey()); authServiceResponse.setFlowStatus(AuthServiceConstants.FlowStatus.FAIL_COMPLETED); - AuthenticationResult authenticationResult = getAuthenticationResult(request, response); + AuthenticationResult authenticationResult = getAuthenticationResult(request); if (authenticationResult != null) { errorCode = (String) authenticationResult.getProperty(FrameworkConstants.AUTH_ERROR_CODE); errorMessage = (String) authenticationResult.getProperty(FrameworkConstants.AUTH_ERROR_MSG); } } else { - authServiceResponse.setSessionDataKey(response.getSessionDataKey()); + authServiceResponse.setSessionDataKey(request.getSessionDataKey()); authServiceResponse.setFlowStatus(AuthServiceConstants.FlowStatus.FAIL_INCOMPLETE); List authenticatorDataList = request.getAuthInitiationData(); AuthServiceResponseData responseData = new AuthServiceResponseData(authenticatorDataList); @@ -242,7 +241,8 @@ private boolean isAuthFlowSuccessful(AuthServiceRequestWrapper request) { private boolean isAuthFlowFailed(AuthServiceRequestWrapper request, AuthServiceResponseWrapper response) throws AuthServiceException { - return AuthenticatorFlowStatus.FAIL_COMPLETED == request.getAuthFlowStatus() || response.isErrorResponse(); + return AuthenticatorFlowStatus.FAIL_COMPLETED == request.getAuthFlowStatus() || response.isErrorResponse() || + isSentToRetryPageOnMissingContext(request, response); } private boolean isAuthFlowIncomplete(AuthServiceRequestWrapper request) { @@ -250,30 +250,30 @@ private boolean isAuthFlowIncomplete(AuthServiceRequestWrapper request) { return AuthenticatorFlowStatus.INCOMPLETE == request.getAuthFlowStatus(); } - 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; - } - - private AuthenticationResult getAuthenticationResult(AuthServiceRequestWrapper request, - AuthServiceResponseWrapper response) - throws AuthServiceException { + private AuthenticationResult getAuthenticationResult(AuthServiceRequestWrapper request) { AuthenticationResult authenticationResult = (AuthenticationResult) request.getAttribute(FrameworkConstants.RequestAttribute.AUTH_RESULT); if (authenticationResult == null) { AuthenticationResultCacheEntry authenticationResultCacheEntry = - FrameworkUtils.getAuthenticationResultFromCache(getFlowCompletionSessionDataKey(request, response)); + FrameworkUtils.getAuthenticationResultFromCache(request.getSessionDataKey()); if (authenticationResultCacheEntry != null) { authenticationResult = authenticationResultCacheEntry.getResult(); } } return authenticationResult; } + + private boolean isSentToRetryPageOnMissingContext(AuthServiceRequestWrapper request, + AuthServiceResponseWrapper response) throws AuthServiceException { + + // If it's a retry due to context being null there is nothing to retry again the flow should be restarted. + if (AuthenticatorFlowStatus.INCOMPLETE == request.getAuthFlowStatus() && + Boolean.TRUE.equals(request.getAttribute(FrameworkConstants.IS_SENT_TO_RETRY))) { + Map queryParams = AuthServiceUtils.extractQueryParams(response.getRedirectURL()); + return StringUtils.equals(queryParams.get(FrameworkConstants.STATUS_PARAM), + FrameworkConstants.ERROR_STATUS_AUTH_CONTEXT_NULL); + } + return false; + } } diff --git a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/handler/request/impl/DefaultRequestCoordinator.java b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/handler/request/impl/DefaultRequestCoordinator.java index f7bf8d17594a..fa02bd950bb5 100644 --- a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/handler/request/impl/DefaultRequestCoordinator.java +++ b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/handler/request/impl/DefaultRequestCoordinator.java @@ -240,6 +240,9 @@ public void handle(HttpServletRequest request, HttpServletResponse response) thr } if (context != null) { + // Adding the context identifier(sessionDataKey) to the request to be used when the context + // is not available. + request.setAttribute(FrameworkConstants.CONTEXT_IDENTIFIER, context.getContextIdentifier()); if (StringUtils.isNotBlank(context.getServiceProviderName())) { MDC.put(SERVICE_PROVIDER_QUERY_KEY, context.getServiceProviderName()); } @@ -372,7 +375,8 @@ public void handle(HttpServletRequest request, HttpServletResponse response) thr log.error("Context does not exist. Probably due to invalidated cache. " + message); FrameworkUtils.sendToRetryPage(request, responseWrapper, context, - "authentication.context.null", "authentication.context.null.description"); + FrameworkConstants.ERROR_STATUS_AUTH_CONTEXT_NULL, + FrameworkConstants.ERROR_DESCRIPTION_AUTH_CONTEXT_NULL); } } catch (JsFailureException e) { if (log.isDebugEnabled()) { diff --git a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/model/auth/service/AuthServiceRequestWrapper.java b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/model/auth/service/AuthServiceRequestWrapper.java index fbb3acd82882..fde5a7264b8d 100644 --- a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/model/auth/service/AuthServiceRequestWrapper.java +++ b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/model/auth/service/AuthServiceRequestWrapper.java @@ -147,4 +147,26 @@ private void skipNonceCookieValidation() { this.setAttribute(FrameworkConstants.SKIP_NONCE_COOKIE_VALIDATION, true); } + + /** + * Get the session data key. + * + * @return String of session data key. + */ + public String getSessionDataKey() { + + if (this.parameters.containsKey(FrameworkConstants.SESSION_DATA_KEY)) { + String[] sessionDataKeyParam = this.parameters.get(FrameworkConstants.SESSION_DATA_KEY); + if (sessionDataKeyParam != null && sessionDataKeyParam.length > 0) { + return sessionDataKeyParam[0]; + } + } + + Object contextIdentifierAttr = getAttribute(FrameworkConstants.CONTEXT_IDENTIFIER); + if (contextIdentifierAttr != null) { + return contextIdentifierAttr.toString(); + } + + return null; + } } diff --git a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/model/auth/service/AuthServiceResponseWrapper.java b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/model/auth/service/AuthServiceResponseWrapper.java index b9a19621435d..fde174136a7a 100644 --- a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/model/auth/service/AuthServiceResponseWrapper.java +++ b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/model/auth/service/AuthServiceResponseWrapper.java @@ -20,7 +20,6 @@ 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; @@ -52,18 +51,6 @@ public String getAuthenticators() throws AuthServiceException { return queryParams.get(AuthServiceConstants.AUTHENTICATORS); } - /** - * Get the sessionDataKey related to the authentication flow. - * - * @return String of sessionDataKey. - * @throws AuthServiceException - */ - public String getSessionDataKey() throws AuthServiceException { - - Map queryParams = AuthServiceUtils.extractQueryParams(getRedirectURL()); - return queryParams.get(FrameworkConstants.SESSION_DATA_KEY); - } - /** * Check if the response is an error response. * This is determined by checking the existence and the value of the diff --git a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/util/FrameworkConstants.java b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/util/FrameworkConstants.java index 3ee2a77b982d..227178b66f32 100644 --- a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/util/FrameworkConstants.java +++ b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/util/FrameworkConstants.java @@ -232,6 +232,10 @@ public abstract class FrameworkConstants { public static final String BLOCKED_USERSTORE_DOMAINS_SEPARATOR = ","; public static final String IS_USER_RESOLVED = "isUserResolved"; + public static final String ERROR_STATUS_AUTH_CONTEXT_NULL = "authentication.context.null"; + public static final String ERROR_DESCRIPTION_AUTH_CONTEXT_NULL = "authentication.context.null.description"; + public static final String IS_SENT_TO_RETRY = "isSentToRetry"; + public static final String CONTEXT_IDENTIFIER = "contextIdentifier"; private FrameworkConstants() { diff --git a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/util/FrameworkUtils.java b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/util/FrameworkUtils.java index 444d02313506..f03dc0b98358 100644 --- a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/util/FrameworkUtils.java +++ b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/util/FrameworkUtils.java @@ -650,6 +650,7 @@ public static void sendToRetryPage(HttpServletRequest request, HttpServletRespon } } request.setAttribute(FrameworkConstants.RequestParams.FLOW_STATUS, AuthenticatorFlowStatus.INCOMPLETE); + request.setAttribute(FrameworkConstants.IS_SENT_TO_RETRY, true); if (context != null) { if (IdentityTenantUtil.isTenantedSessionsEnabled()) { uriBuilder.addParameter(USER_TENANT_DOMAIN_HINT, context.getUserTenantDomain()); diff --git a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/test/java/org/wso2/carbon/identity/application/authentication/framework/AuthenticationServiceTest.java b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/test/java/org/wso2/carbon/identity/application/authentication/framework/AuthenticationServiceTest.java index 19ec6aeddef5..6c78f095eb1a 100644 --- a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/test/java/org/wso2/carbon/identity/application/authentication/framework/AuthenticationServiceTest.java +++ b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/test/java/org/wso2/carbon/identity/application/authentication/framework/AuthenticationServiceTest.java @@ -122,6 +122,7 @@ public void testHandleAuthentication(boolean isMultiOpsResponse, String redirect AuthServiceRequest authServiceRequest = new AuthServiceRequest(request, response); when(request.getAttribute(FrameworkConstants.IS_MULTI_OPS_RESPONSE)).thenReturn(isMultiOpsResponse); when(request.getAttribute(FrameworkConstants.RequestParams.FLOW_STATUS)).thenReturn(authenticatorFlowStatus); + when(request.getAttribute(FrameworkConstants.SESSION_DATA_KEY)).thenReturn(sessionDataKey); if (isMultiOpsResponse) { List authenticatorDataMap = getMultiOpsAuthenticatorData(authenticatorList); for (AuthenticatorData authenticatorData : authenticatorDataMap) { @@ -179,6 +180,7 @@ public void testNegativeHandleAuthentication(String redirectUrl, Object authenti AuthServiceRequest authServiceRequest = new AuthServiceRequest(request, response); when(request.getAttribute(FrameworkConstants.IS_MULTI_OPS_RESPONSE)).thenReturn(false); when(request.getAttribute(FrameworkConstants.RequestParams.FLOW_STATUS)).thenReturn(authenticatorFlowStatus); + when(request.getAttribute(FrameworkConstants.SESSION_DATA_KEY)).thenReturn(sessionDataKey); List expected = getAuthenticatorData(authenticatorList); if (AuthenticatorFlowStatus.INCOMPLETE == authenticatorFlowStatus) { when(request.getAttribute(AuthServiceConstants.AUTH_SERVICE_AUTH_INITIATION_DATA)).thenReturn(expected);