Enable Authorization Code grant support for sub organization OAuth2 applications

Author: ShanChathusanda93

components/org.wso2.carbon.identity.oauth.endpoint/src/main/java/org/wso2/carbon/identity/oauth/endpoint/util/AuthzUtil.java

@@ -43,6 +43,7 @@
 import org.json.JSONException;
 import org.json.JSONObject;
 import org.owasp.encoder.Encode;
+import org.wso2.carbon.context.PrivilegedCarbonContext;
 import org.wso2.carbon.identity.application.authentication.framework.AuthenticationService;
 import org.wso2.carbon.identity.application.authentication.framework.AuthenticatorFlowStatus;
 import org.wso2.carbon.identity.application.authentication.framework.CommonAuthenticationHandler;
@@ -51,6 +52,7 @@
 import org.wso2.carbon.identity.application.authentication.framework.config.model.graph.js.JsLogger;
 import org.wso2.carbon.identity.application.authentication.framework.context.AuthHistory;
 import org.wso2.carbon.identity.application.authentication.framework.context.SessionContext;
+import org.wso2.carbon.identity.application.authentication.framework.exception.FrameworkException;
 import org.wso2.carbon.identity.application.authentication.framework.exception.UserIdNotFoundException;
 import org.wso2.carbon.identity.application.authentication.framework.exception.auth.service.AuthServiceClientException;
 import org.wso2.carbon.identity.application.authentication.framework.exception.auth.service.AuthServiceException;
@@ -151,6 +153,8 @@
 import org.wso2.carbon.identity.openidconnect.OIDCRequestObjectUtil;
 import org.wso2.carbon.identity.openidconnect.model.RequestObject;
 import org.wso2.carbon.identity.openidconnect.model.RequestedClaim;
+import org.wso2.carbon.identity.organization.management.service.exception.OrganizationManagementException;
+import org.wso2.carbon.identity.organization.management.service.util.OrganizationManagementUtil;
 import org.wso2.carbon.utils.CarbonUtils;
 import org.wso2.carbon.utils.DiagnosticLog;
 
@@ -1523,7 +1527,7 @@ private static void addToAuthenticationResultDetailsToOAuthMessage(OAuthMessage
         addMappedRemoteClaimsToSessionCache(oAuthMessage, authnResult);
     }
 
-    private static void updateAuthTimeInSessionDataCacheEntry(OAuthMessage oAuthMessage) {
+    private static void updateAuthTimeInSessionDataCacheEntry(OAuthMessage oAuthMessage) throws OAuthSystemException {
 
         String commonAuthIdCookieValue = getCommonAuthCookieString(oAuthMessage.getRequest());
         long authTime = getAuthenticatedTimeFromCommonAuthCookieValue(commonAuthIdCookieValue,
@@ -2429,7 +2433,7 @@ public static void populateValidationResponseWithAppDetail(OAuthMessage oAuthMes
 
         String clientId = oAuthMessage.getRequest().getParameter(CLIENT_ID);
         try {
-            OAuthAppDO appDO = OAuth2Util.getAppInformationByClientId(clientId);
+            OAuthAppDO appDO = OAuth2Util.getAppInformationByClientId(clientId, OAuth2Util.extractTenantDomain());
             if (Boolean.TRUE.equals(oAuthMessage.getRequest().getAttribute(OAuthConstants.PKCE_UNSUPPORTED_FLOW))) {
                 validationResponse.setPkceMandatory(false);
             } else {
@@ -2834,7 +2838,8 @@ private static String getSpTenantDomain(String clientId) throws InvalidRequestEx
         try {
             // At this point we have verified that a valid app exists for the client_id. So we directly get the SP
             // tenantDomain.
-            return OAuth2Util.getTenantDomainOfOauthApp(clientId);
+            String tenantDomain = IdentityTenantUtil.getTenantDomain(IdentityTenantUtil.getLoginTenantId());
+            return OAuth2Util.getTenantDomainOfOauthApp(clientId, tenantDomain);
         } catch (InvalidOAuthClientException | IdentityOAuth2Exception e) {
             throw new InvalidRequestException("Error retrieving Service Provider tenantDomain for client_id: "
                     + clientId, OAuth2ErrorCodes.INVALID_REQUEST, OAuth2ErrorCodes.OAuth2SubErrorCodes
@@ -2852,7 +2857,12 @@ private static String getLoginTenantDomain(OAuthMessage oAuthMessage, String cli
         String loginTenantDomain =
                 oAuthMessage.getRequest().getParameter(FrameworkConstants.RequestParams.LOGIN_TENANT_DOMAIN);
         if (StringUtils.isBlank(loginTenantDomain)) {
-            return EndpointUtil.getSPTenantDomainFromClientId(oAuthMessage.getClientId());
+            try {
+                return EndpointUtil.verifyAndRetrieveTenantDomain(clientId);
+            } catch (OAuthSystemException e) {
+                throw new InvalidRequestException("Error resolving tenant domain for client id: " + clientId,
+                        OAuth2ErrorCodes.INVALID_REQUEST, OAuth2ErrorCodes.OAuth2SubErrorCodes.INVALID_REQUEST);
+            }
         }
         return loginTenantDomain;
     }
@@ -4351,10 +4361,29 @@ private static String appendAuthenticatedIDPs(SessionDataCacheEntry sessionDataC
      * @param resultFromLogin The session context.
      * @param cookieValue The cookie string which contains the commonAuthId value.
      */
-    private static void associateAuthenticationHistory(SessionDataCacheEntry resultFromLogin, String cookieValue) {
+    private static void associateAuthenticationHistory(SessionDataCacheEntry resultFromLogin, String cookieValue)
+            throws OAuthSystemException {
 
-        SessionContext sessionContext = getSessionContext(cookieValue,
-                resultFromLogin.getoAuth2Parameters().getLoginTenantDomain());
+        String tenantDomain = resultFromLogin.getoAuth2Parameters().getLoginTenantDomain();
+        /*
+         If the app is created in an organization, get the tenant domain of the primary organization since the
+         session is stored against the primary organization.
+        */
+        try {
+            String appTenantDomain = OAuth2Util.getTenantDomainOfOauthApp(
+                    resultFromLogin.getoAuth2Parameters().getClientId(), tenantDomain);
+            if (OrganizationManagementUtil.isOrganization(appTenantDomain)) {
+                String appOrgId = OAuth2ServiceComponentHolder.getInstance().getOrganizationManager()
+                        .resolveOrganizationId(appTenantDomain);
+                String primaryOrgId = OAuth2ServiceComponentHolder.getInstance().getOrganizationManager()
+                        .getPrimaryOrganizationId(appOrgId);
+                tenantDomain = OAuth2ServiceComponentHolder.getInstance().getOrganizationManager()
+                        .resolveTenantDomain(primaryOrgId);
+            }
+        } catch (OrganizationManagementException | IdentityOAuth2Exception | InvalidOAuthClientException e) {
+            throw new OAuthSystemException(e);
+        }
+        SessionContext sessionContext = getSessionContext(cookieValue, tenantDomain);
         if (sessionContext != null && sessionContext.getSessionAuthHistory() != null
                 && sessionContext.getSessionAuthHistory().getHistory() != null) {
             List<String> authMethods = new ArrayList<>();

components/org.wso2.carbon.identity.oauth.endpoint/src/main/java/org/wso2/carbon/identity/oauth/endpoint/util/EndpointUtil.java

@@ -1544,6 +1544,29 @@ public static String getSPTenantDomainFromClientId(String clientId) {
         }
     }
 
+    /**
+     * This method verifies the service provider tenant domain using the client ID.
+     *
+     * @param clientId Client id of the application.
+     * @return tenantDomain domain of the service provider.
+     */
+    public static String verifyAndRetrieveTenantDomain(String clientId)
+            throws OAuthSystemException {
+
+        try {
+            String extractedTenantDomain = OAuth2Util.extractTenantDomain();
+            OAuthAppDO oAuthAppDO = OAuth2Util.getAppInformationByClientId(clientId, extractedTenantDomain);
+            String appTenantDomain = OAuth2Util.getTenantDomainOfOauthApp(oAuthAppDO);
+            if (StringUtils.equals(extractedTenantDomain, appTenantDomain)) {
+                return appTenantDomain;
+            }
+            throw new OAuthSystemException("Provided tenant domain: " + extractedTenantDomain + " does not " +
+                    "match with the application's tenant domain: " + appTenantDomain);
+        } catch (IdentityOAuth2Exception | InvalidOAuthClientException e) {
+            throw new OAuthSystemException("Error while getting oauth app for client Id: " + clientId, e);
+        }
+    }
+
     /**
      * Extract information related to the token request and exception and publish the event to listeners.
      *