Update the handling of invite tokens for new user registration to work with the base level oauth flows and implicit grants

Signed-off-by: Alice W <105500542+alice-wondered@users.noreply.github.com>

server-spi-private/src/main/java/org/keycloak/models/Constants.java

@@ -80,7 +80,6 @@ public final class Constants {
     public static final String EXECUTION = "execution";
     public static final String CLIENT_ID = "client_id";
     public static final String ORG_TOKEN = "org_token";
-    public static final String ORG_INVITE = "org_invite";
     public static final String TAB_ID = "tab_id";
     public static final String CLIENT_DATA = "client_data";
 

services/src/main/java/org/keycloak/authentication/AuthenticationProcessor.java

@@ -18,6 +18,7 @@
 package org.keycloak.authentication;
 
 import org.jboss.logging.Logger;
+import org.keycloak.authentication.actiontoken.inviteorg.InviteOrgActionToken;
 import org.keycloak.http.HttpRequest;
 import org.keycloak.authentication.authenticators.browser.AbstractUsernameFormAuthenticator;
 import org.keycloak.authentication.authenticators.client.ClientAuthUtil;
@@ -101,6 +102,9 @@ public class AuthenticationProcessor {
     protected HttpRequest request;
     protected String flowId;
     protected String flowPath;
+
+
+    protected String orgToken;
     protected boolean browserFlow;
     protected BruteForceProtector protector;
     protected Runnable afterResetListener;
@@ -228,6 +232,15 @@ public class AuthenticationProcessor {
         return this;
     }
 
+    public String getOrgToken() {
+        return orgToken;
+    }
+
+    public AuthenticationProcessor setOrgToken(String orgToken) {
+        this.orgToken = orgToken;
+        return this;
+    }
+
     public AuthenticationProcessor setForwardedErrorMessage(FormMessage forwardedErrorMessage) {
         this.forwardedErrorMessageStore.setForwardedMessage(forwardedErrorMessage);
         return this;

services/src/main/java/org/keycloak/authentication/FormAuthenticationFlow.java

@@ -17,6 +17,7 @@
 
 package org.keycloak.authentication;
 
+import jakarta.ws.rs.core.UriBuilder;
 import org.keycloak.http.HttpRequest;
 import org.keycloak.common.ClientConnection;
 import org.keycloak.events.EventBuilder;
@@ -264,15 +265,19 @@ public class FormAuthenticationFlow implements AuthenticationFlow {
         return null;
     }
 
-    public URI getActionUrl(String executionId, String code) {
+    public URI getActionUrl(String executionId, String code, String token) {
         ClientModel client = processor.getAuthenticationSession().getClient();
-        return LoginActionsService.registrationFormProcessor(processor.getUriInfo())
+        UriBuilder builder = LoginActionsService.registrationFormProcessor(processor.getUriInfo())
                 .queryParam(LoginActionsService.SESSION_CODE, code)
                 .queryParam(Constants.EXECUTION, executionId)
                 .queryParam(Constants.CLIENT_ID, client.getClientId())
                 .queryParam(Constants.TAB_ID, processor.getAuthenticationSession().getTabId())
-                .queryParam(Constants.CLIENT_DATA, AuthenticationProcessor.getClientData(processor.getSession(), processor.getAuthenticationSession()))
+                .queryParam(Constants.CLIENT_DATA, AuthenticationProcessor.getClientData(processor.getSession(), processor.getAuthenticationSession()));
-                .build(processor.getRealm().getName());
+        if (token != null) {
+            builder.queryParam(Constants.ORG_TOKEN, token);
+        }
+
+        return builder.build(processor.getRealm().getName());
     }
 
 
@@ -290,7 +295,7 @@ public class FormAuthenticationFlow implements AuthenticationFlow {
         String executionId = formExecution.getId();
         processor.getAuthenticationSession().setAuthNote(AuthenticationProcessor.CURRENT_AUTHENTICATION_EXECUTION, executionId);
         String code = processor.generateCode();
-        URI actionUrl = getActionUrl(executionId, code);
+        URI actionUrl = getActionUrl(executionId, code, processor.orgToken);
         LoginFormsProvider form = processor.getSession().getProvider(LoginFormsProvider.class)
                 .setAuthenticationSession(processor.getAuthenticationSession())
                 .setActionUri(actionUrl)

services/src/main/java/org/keycloak/protocol/oidc/endpoints/AuthorizationEndpoint.java

@@ -195,7 +195,7 @@ public class AuthorizationEndpoint extends AuthorizationEndpointBase {
         CacheControlUtil.noBackButtonCacheControlHeader(session);
         switch (action) {
             case REGISTER:
-                return buildRegister();
+                return buildRegister(params.getFirst(Constants.ORG_TOKEN));
             case FORGOT_CREDENTIALS:
                 return buildForgotCredential();
             case CODE:
@@ -341,13 +341,16 @@ public class AuthorizationEndpoint extends AuthorizationEndpointBase {
         return handleBrowserAuthenticationRequest(authenticationSession, new OIDCLoginProtocol(session, realm, session.getContext().getUri(), headers, event), TokenUtil.hasPrompt(request.getPrompt(), OIDCLoginProtocol.PROMPT_VALUE_NONE), false);
     }
 
-    private Response buildRegister() {
+    private Response buildRegister(String inviteToken) {
         authManager.expireIdentityCookie(session);
 
         AuthenticationFlowModel flow = realm.getRegistrationFlow();
         String flowId = flow.getId();
 
         AuthenticationProcessor processor = createProcessor(authenticationSession, flowId, LoginActionsService.REGISTRATION_PATH);
+        if (inviteToken != null) {
+            processor.setOrgToken(inviteToken);
+        }
         authenticationSession.setClientNote(APP_INITIATED_FLOW, LoginActionsService.REGISTRATION_PATH);
 
         return processor.authenticate();

services/src/main/java/org/keycloak/services/resources/LoginActionsService.java

@@ -735,6 +735,10 @@ public class LoginActionsService {
         return processFlow(action, execution, authSession, REGISTRATION_PATH, realm.getRegistrationFlow(), errorMessage, new AuthenticationProcessor());
     }
 
+    protected Response processRegistrationWithInviteToken(boolean action, String execution, AuthenticationSessionModel authSession, String errorMessage, String token) {
+        AuthenticationProcessor authenticationProcessor = new AuthenticationProcessor().setOrgToken(token);
+        return processFlow(action, execution, authSession, REGISTRATION_PATH, realm.getRegistrationFlow(), errorMessage, authenticationProcessor);
+    }
 
     /**
      * protocol independent registration page entry point
@@ -746,15 +750,13 @@ public class LoginActionsService {
     @GET
     public Response registerPage(@QueryParam(AUTH_SESSION_ID) String authSessionId, // optional, can get from cookie instead
                                  @QueryParam(SESSION_CODE) String code,
-                                 // TODO this is unused but having it here adds it to openapi. What's the better approach?
-                                 // Should this be pulled off the query params and then injected into the flow processor as its own thing?
                                  @QueryParam(Constants.ORG_TOKEN) String orgToken,
                                  @QueryParam(Constants.EXECUTION) String execution,
                                  @QueryParam(Constants.CLIENT_ID) String clientId,
                                  @QueryParam(Constants.CLIENT_DATA) String clientData,
                                  @QueryParam(Constants.TAB_ID) String tabId) {
 
-        return registerRequest(authSessionId, code, execution, clientId,  tabId,clientData);
+        return registerRequest(authSessionId, code, execution, clientId,  tabId,clientData, orgToken);
     }
 
 
@@ -773,15 +775,13 @@ public class LoginActionsService {
                                     @QueryParam(Constants.CLIENT_ID) String clientId,
                                     @QueryParam(Constants.CLIENT_DATA) String clientData,
                                     @QueryParam(Constants.TAB_ID) String tabId) {
-        return registerRequest(authSessionId, code, execution, clientId,  tabId,clientData);
+        return registerRequest(authSessionId, code, execution, clientId,  tabId,clientData, orgToken);
     }
 
 
-    private Response registerRequest(String authSessionId, String code, String execution, String clientId, String tabId, String clientData) {
+    private Response registerRequest(String authSessionId, String code, String execution, String clientId, String tabId, String clientData, String orgToken) {
         event.event(EventType.REGISTER);
 
-        // TODO if we parse the org token here and then pass in the already decoded token we can save ourselves some duplicated work
-
         if (!realm.isRegistrationAllowed()) {
             event.error(Errors.REGISTRATION_DISABLED);
             return ErrorPage.error(session, null, Response.Status.BAD_REQUEST, Messages.REGISTRATION_NOT_ALLOWED);
@@ -798,7 +798,7 @@ public class LoginActionsService {
 
         AuthenticationManager.expireIdentityCookie(session);
 
-        return processRegistration(checks.isActionRequest(), execution, authSession, null);
+        return processRegistrationWithInviteToken(checks.isActionRequest(), execution, authSession, null, orgToken);
     }