Callers need to indicate if cookies need to be set at the end of the transaction

Closes #17141

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

@@ -49,6 +49,7 @@ import org.keycloak.services.clientpolicy.context.PreAuthorizationRequestContext
 import org.keycloak.services.messages.Messages;
 import org.keycloak.services.resources.LoginActionsService;
 import org.keycloak.services.util.CacheControlUtil;
+import org.keycloak.services.util.CookieHelper;
 import org.keycloak.sessions.AuthenticationSessionModel;
 import org.keycloak.util.TokenUtil;
 
@@ -132,6 +133,7 @@ public class AuthorizationEndpoint extends AuthorizationEndpointBase {
         return KeycloakModelUtils.runJobInRetriableTransaction(session.getKeycloakSessionFactory(), new ResponseSessionTask(session) {
             @Override
             public Response runInternal(KeycloakSession session) {
+                CookieHelper.addCookiesAtEndOfTransaction(session);
                 // create another instance of the endpoint to isolate each run.
                 AuthorizationEndpoint other = new AuthorizationEndpoint(session,
                         new EventBuilder(session.getContext().getRealm(), session, clientConnection), action);

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

@@ -81,6 +81,7 @@ import org.keycloak.services.messages.Messages;
 import org.keycloak.services.util.AuthenticationFlowURLHelper;
 import org.keycloak.services.util.BrowserHistoryHelper;
 import org.keycloak.services.util.CacheControlUtil;
+import org.keycloak.services.util.CookieHelper;
 import org.keycloak.services.util.LocaleUtil;
 import org.keycloak.sessions.AuthenticationSessionCompoundId;
 import org.keycloak.sessions.AuthenticationSessionModel;
@@ -260,6 +261,7 @@ public class LoginActionsService {
             @Override
             public Response runInternal(KeycloakSession session) {
                 // create another instance of the endpoint to isolate each run.
+                CookieHelper.addCookiesAtEndOfTransaction(session);
                 LoginActionsService other = new LoginActionsService(session, new EventBuilder(session.getContext().getRealm(), session, clientConnection));
                 // process the request in the created instance.
                 return other.authenticateInternal(authSessionId, code, execution, clientId, tabId);

services/src/main/java/org/keycloak/services/util/CookieHelper.java

@@ -29,6 +29,7 @@ import javax.ws.rs.core.HttpHeaders;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 
 import static org.keycloak.common.util.ServerCookie.SameSiteAttributeValue;
@@ -43,6 +44,7 @@ public class CookieHelper {
     public static final String LEGACY_COOKIE = "_LEGACY";
 
     private static final Logger logger = Logger.getLogger(CookieHelper.class);
+    private static final String ADD_COOKIES_AT_END_OF_TRANSACTION = CookieHelper.class.getName() + "_ADD_COOKIES_AT_END_OF_TRANSACTION";
 
     /**
      * Set a response cookie.  This solely exists because JAX-RS 1.1 does not support setting HttpOnly cookies
@@ -70,7 +72,11 @@ public class CookieHelper {
         StringBuffer cookieBuf = new StringBuffer();
         ServerCookie.appendCookieValue(cookieBuf, 1, name, value, path, domain, comment, maxAge, secure_sameSite, httpOnly, sameSite);
         String cookie = cookieBuf.toString();
+        if (shouldAddCookiesAtEndOfTransaction(session)) {
             session.getTransactionManager().enlistAfterCompletion(new CookieTransaction(response, cookie));
+        } else {
+            response.addHeader(HttpHeaders.SET_COOKIE, cookie);
+        }
 
         // a workaround for browser in older Apple OSs – browsers ignore cookies with SameSite=None
         if (sameSiteParam == SameSiteAttributeValue.NONE) {
@@ -78,6 +84,21 @@ public class CookieHelper {
         }
     }
 
+    private static boolean shouldAddCookiesAtEndOfTransaction(KeycloakSession session) {
+        return Objects.equals(session.getAttribute(ADD_COOKIES_AT_END_OF_TRANSACTION), Boolean.TRUE);
+    }
+
+    /**
+     * Adding cookies at the end of the transaction helps when retrying a transaction might add the
+     * cookie multiple times. In some scenarios it must not be added at the end of the transaction,
+     * as at that time the response has already been sent to the caller ("committed"), so the code
+     * needs to make a choice. As retrying transactions is the exception, adding cookies at the end
+     * of the transaction is also the exception and needs to be switched on where necessary.
+     */
+    public static void addCookiesAtEndOfTransaction(KeycloakSession session) {
+        session.setAttribute(ADD_COOKIES_AT_END_OF_TRANSACTION, Boolean.TRUE);
+    }
+
     /**
      * Set a response cookie avoiding SameSite parameter
      * @param name