Re-calculate the organization scope when re-authenticating in the bro…

…wser flow

Closes keycloak#35935
Closes keycloak#35830

Signed-off-by: Pedro Igor <[email protected]>

services/src/main/java/org/keycloak/authentication/authenticators/browser/CookieAuthenticator.java

@@ -26,7 +26,10 @@
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.UserModel;
+import org.keycloak.organization.protocol.mappers.oidc.OrganizationScope;
+import org.keycloak.organization.utils.Organizations;
 import org.keycloak.protocol.LoginProtocol;
+import org.keycloak.protocol.oidc.OIDCLoginProtocol;
 import org.keycloak.services.managers.AuthenticationManager;
 import org.keycloak.services.messages.Messages;
 import org.keycloak.sessions.AuthenticationSessionModel;
@@ -84,7 +87,13 @@ public void authenticate(AuthenticationFlowContext context) {
                     acrStore.setLevelAuthenticatedToCurrentRequest(previouslyAuthenticatedLevel);
                     authSession.setAuthNote(AuthenticationManager.SSO_AUTH, "true");
                     context.attachUserSession(authResult.getSession());
-                    context.success();
+
+                    if (isOrganizationContext(context)) {
+                        // if re-authenticating in the scope of an organization, an organization must be resolved prior to authenticating the user
+                        context.attempted();
+                    } else {
+                        context.success();
+                    }
                 }
             }
         }
@@ -109,4 +118,16 @@ public void setRequiredActions(KeycloakSession session, RealmModel realm, UserMo
     public void close() {
 
     }
+
+    private boolean isOrganizationContext(AuthenticationFlowContext context) {
+        KeycloakSession session = context.getSession();
+
+        if (Organizations.isEnabledAndOrganizationsPresent(session)) {
+            AuthenticationSessionModel authSession = context.getAuthenticationSession();
+            String requestedScopes = authSession.getClientNote(OIDCLoginProtocol.SCOPE_PARAM);
+            return OrganizationScope.valueOfScope(requestedScopes) != null;
+        }
+
+        return false;
+    }
 }

services/src/main/java/org/keycloak/organization/authentication/authenticators/browser/OrganizationAuthenticator.java

@@ -17,6 +17,7 @@
 
 package org.keycloak.organization.authentication.authenticators.browser;
 
+import static org.keycloak.authentication.AuthenticatorUtil.isSSOAuthentication;
 import static org.keycloak.organization.utils.Organizations.getEmailDomain;
 import static org.keycloak.organization.utils.Organizations.isEnabledAndOrganizationsPresent;
 import static org.keycloak.organization.utils.Organizations.resolveHomeBroker;
@@ -121,7 +122,12 @@ public void action(AuthenticationFlowContext context) {
             return;
         }
 
-        context.attempted();
+        if (isSSOAuthentication(context.getAuthenticationSession())) {
+            // if re-authenticating in the scope of an organization
+            context.success();
+        } else {
+            context.attempted();
+        }
     }
 
     @Override
@@ -134,9 +140,17 @@ private OrganizationModel resolveOrganization(UserModel user, String domain) {
         HttpRequest request = context.getHttpRequest();
         MultivaluedMap<String, String> parameters = request.getDecodedFormParameters();
         List<String> alias = parameters.getOrDefault(OrganizationModel.ORGANIZATION_ATTRIBUTE, List.of());
+        AuthenticationSessionModel authSession = context.getAuthenticationSession();
 
         if (alias.isEmpty()) {
-            return Organizations.resolveOrganization(session, user, domain);
+            OrganizationModel organization = Organizations.resolveOrganization(session, user, domain);
+
+            if (organization != null) {
+                // make sure the organization selected by the user is available from the client session when running mappers and issuing tokens
+                authSession.setClientNote(OrganizationModel.ORGANIZATION_ATTRIBUTE, organization.getId());
+            }
+
+            return organization;
         }
 
         OrganizationProvider provider = getOrganizationProvider();
@@ -146,7 +160,6 @@ private OrganizationModel resolveOrganization(UserModel user, String domain) {
             return null;
         }
 
-        AuthenticationSessionModel authSession = context.getAuthenticationSession();
         // make sure the organization selected by the user is available from the client session when running mappers and issuing tokens
         authSession.setClientNote(OrganizationModel.ORGANIZATION_ATTRIBUTE, organization.getId());
 
@@ -163,6 +176,11 @@ private boolean shouldUserSelectOrganization(AuthenticationFlowContext context,
             return false;
         }
 
+        if (authSession.getClientNote(OrganizationModel.ORGANIZATION_ATTRIBUTE) != null) {
+            // organization already selected
+            return false;
+        }
+
         Stream<OrganizationModel> organizations = provider.getByMember(user);
 
         if (organizations.count() > 1) {
@@ -274,20 +292,31 @@ private void unknownUserChallenge(AuthenticationFlowContext context, Organizatio
         context.challenge(form.createLoginUsername());
     }
 
-    private void initialChallenge(AuthenticationFlowContext context){
-        // the default challenge won't show any broker but just the identity-first login page and the option to try a different authentication mechanism
-        LoginFormsProvider form = context.form()
-                .setAttributeMapper(attributes -> {
-                    attributes.computeIfPresent("social",
-                            (key, bean) -> new OrganizationAwareIdentityProviderBean((IdentityProviderBean) bean, false, true)
-                    );
-                    attributes.computeIfPresent("auth",
-                            (key, bean) -> new OrganizationAwareAuthenticationContextBean((AuthenticationContextBean) bean, false)
-                    );
-                    return attributes;
-                });
+    private void initialChallenge(AuthenticationFlowContext context) {
+        UserModel user = context.getUser();
 
-        context.challenge(form.createLoginUsername());
+        if (user == null) {
+            // the default challenge won't show any broker but just the identity-first login page and the option to try a different authentication mechanism
+            LoginFormsProvider form = context.form()
+                    .setAttributeMapper(attributes -> {
+                        attributes.computeIfPresent("social",
+                                (key, bean) -> new OrganizationAwareIdentityProviderBean((IdentityProviderBean) bean, false, true)
+                        );
+                        attributes.computeIfPresent("auth",
+                                (key, bean) -> new OrganizationAwareAuthenticationContextBean((AuthenticationContextBean) bean, false)
+                        );
+                        return attributes;
+                    });
+
+            context.challenge(form.createLoginUsername());
+        } else if (isSSOAuthentication(context.getAuthenticationSession())) {
+            if (shouldUserSelectOrganization(context, user)) {
+                return;
+            }
+
+            // user is re-authenticating and there are no organizations to select
+            context.success();
+        }
     }
 
     private boolean hasPublicBrokers(OrganizationModel organization) {

services/src/main/java/org/keycloak/organization/protocol/mappers/oidc/OrganizationMembershipMapper.java

@@ -110,18 +110,11 @@ public String getHelpText() {
 
     @Override
     protected void setClaim(IDToken token, ProtocolMapperModel model, UserSessionModel userSession, KeycloakSession session, ClientSessionContext clientSessionCtx) {
-        String rawScopes = clientSessionCtx.getScopeString();
-        OrganizationScope scope = OrganizationScope.valueOfScope(rawScopes);
-
-        if (scope == null) {
-            return;
-        }
-
         String orgId = clientSessionCtx.getClientSession().getNote(OrganizationModel.ORGANIZATION_ATTRIBUTE);
         Stream<OrganizationModel> organizations;
 
         if (orgId == null) {
-            organizations = scope.resolveOrganizations(userSession.getUser(), rawScopes, session);
+            organizations = resolveFromRequestedScopes(session, userSession, clientSessionCtx);
         } else {
             organizations = Stream.of(session.getProvider(OrganizationProvider.class).getById(orgId));
         }
@@ -139,6 +132,18 @@ protected void setClaim(IDToken token, ProtocolMapperModel model, UserSessionMod
         OIDCAttributeMapperHelper.mapClaim(token, effectiveModel, claim);
     }
 
+    private Stream<OrganizationModel> resolveFromRequestedScopes(KeycloakSession session, UserSessionModel userSession, ClientSessionContext context) {
+        String rawScopes = context.getScopeString();
+        OrganizationScope scope = OrganizationScope.valueOfScope(rawScopes);
+
+        if (scope == null) {
+            return Stream.empty();
+        }
+
+        return scope.resolveOrganizations(userSession.getUser(), rawScopes, session);
+
+    }
+
     private Object resolveValue(ProtocolMapperModel model, List<OrganizationModel> organizations) {
         if (organizations.isEmpty()) {
             return null;