Add retry logic to LoginActionsService#authenticate

In addition to that, avoid adding cookies on each retry.

Closes #15849

Co-authored-by: Alexander Schwartz <aschwart@redhat.com>

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

@@ -311,6 +311,7 @@ public final class KeycloakModelUtils {
                 } else {
                     if (retryCount == attemptsCount) {
                         logger.debug("Exhausted all retry attempts for request.");
+                        throw new RuntimeException("retries exceeded", re);
                     }
                     throw re;
                 }

services/src/main/java/org/keycloak/services/HttpResponseImpl.java

@@ -21,7 +21,7 @@ import org.keycloak.http.HttpResponse;
 
 public class HttpResponseImpl implements HttpResponse {
 
-    private org.jboss.resteasy.spi.HttpResponse delegate;
+    private final org.jboss.resteasy.spi.HttpResponse delegate;
 
     public HttpResponseImpl(org.jboss.resteasy.spi.HttpResponse delegate) {
         this.delegate = delegate;
@@ -34,11 +34,25 @@ public class HttpResponseImpl implements HttpResponse {
 
     @Override
     public void addHeader(String name, String value) {
+        checkCommitted();
         delegate.getOutputHeaders().add(name, value);
     }
 
     @Override
     public void setHeader(String name, String value) {
+        checkCommitted();
         delegate.getOutputHeaders().putSingle(name, value);
     }
+
+    /**
+     * Validate that the response has not been committed.
+     * If the response is already committed, the headers and part of the response have been sent already.
+     * Therefore, additional headers including cookies won't be delivered to the caller.
+     */
+    private void checkCommitted() {
+        if (delegate.isCommitted()) {
+            throw new IllegalStateException("response already committed, can't be changed");
+        }
+    }
+
 }

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

@@ -17,6 +17,7 @@
 package org.keycloak.services.resources;
 
 import org.jboss.logging.Logger;
+import org.keycloak.common.util.ResponseSessionTask;
 import org.keycloak.http.HttpRequest;
 import org.keycloak.OAuth2Constants;
 import org.keycloak.TokenVerifier;
@@ -255,6 +256,20 @@ public class LoginActionsService {
                                  @QueryParam(Constants.EXECUTION) String execution,
                                  @QueryParam(Constants.CLIENT_ID) String clientId,
                                  @QueryParam(Constants.TAB_ID) String tabId) {
+        return KeycloakModelUtils.runJobInRetriableTransaction(session.getKeycloakSessionFactory(), new ResponseSessionTask(session) {
+            @Override
+            public Response runInternal(KeycloakSession session) {
+                // create another instance of the endpoint to isolate each run.
+                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);
+            }
+        }, 10, 100);
+
+    }
+
+    private Response authenticateInternal(final String authSessionId, final String code, final String execution,
+                                          final String clientId, final String tabId) {
         event.event(EventType.LOGIN);
 
         SessionCodeChecks checks = checksForCode(authSessionId, code, execution, clientId, tabId, AUTHENTICATE_PATH);

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

@@ -20,9 +20,9 @@ package org.keycloak.services.util;
 import org.jboss.logging.Logger;
 import org.keycloak.http.HttpResponse;
 import org.jboss.resteasy.util.CookieParser;
-import org.keycloak.common.util.Resteasy;
 import org.keycloak.common.util.ServerCookie;
 import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.KeycloakTransaction;
 
 import javax.ws.rs.core.Cookie;
 import javax.ws.rs.core.HttpHeaders;
@@ -70,7 +70,7 @@ 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();
-        response.addHeader(HttpHeaders.SET_COOKIE, cookie);
+        session.getTransactionManager().enlistAfterCompletion(new CookieTransaction(response, cookie));
 
         // a workaround for browser in older Apple OSs – browsers ignore cookies with SameSite=None
         if (sameSiteParam == SameSiteAttributeValue.NONE) {
@@ -151,4 +151,49 @@ public class CookieHelper {
             return cookies.get(legacy);
         }
     }
+
+    /**
+     * Ensure that cookies are only added when the transaction is complete, as otherwise cookies will be set for error pages,
+     * or will be added twice when running retries.
+     */
+    private static class CookieTransaction implements KeycloakTransaction {
+        private final HttpResponse response;
+        private final String cookie;
+        private boolean transactionActive;
+
+        public CookieTransaction(HttpResponse response, String cookie) {
+            this.response = response;
+            this.cookie = cookie;
+        }
+
+        @Override
+        public void begin() {
+            transactionActive = true;
+        }
+
+        @Override
+        public void commit() {
+            response.addHeader(HttpHeaders.SET_COOKIE, cookie);
+            transactionActive = false;
+        }
+
+        @Override
+        public void rollback() {
+            transactionActive = false;
+        }
+
+        @Override
+        public void setRollbackOnly() {
+        }
+
+        @Override
+        public boolean getRollbackOnly() {
+            return false;
+        }
+
+        @Override
+        public boolean isActive() {
+            return transactionActive;
+        }
+    }
 }