keycloak-aplcache

diff --git a/services/src/main/java/org/keycloak/services/managers/AuthenticationManager.java b/services/src/main/java/org/keycloak/services/managers/AuthenticationManager.java
index 1e8a53d..38c43da 100755
--- a/services/src/main/java/org/keycloak/services/managers/AuthenticationManager.java
+++ b/services/src/main/java/org/keycloak/services/managers/AuthenticationManager.java
@@ -30,6 +30,7 @@ import org.keycloak.authentication.actiontoken.DefaultActionTokenKey;
 import org.keycloak.broker.provider.IdentityProvider;
 import org.keycloak.common.ClientConnection;
 import org.keycloak.common.VerificationException;
+import org.keycloak.common.util.Base64Url;
 import org.keycloak.common.util.Time;
 import org.keycloak.events.Details;
 import org.keycloak.events.Errors;
@@ -477,7 +478,7 @@ public class AuthenticationManager {
     }
 
 
-    public static AccessToken createIdentityToken(RealmModel realm, UserModel user, UserSessionModel session, String issuer) {
+    public static AccessToken createIdentityToken(KeycloakSession keycloakSession, RealmModel realm, UserModel user, UserSessionModel session, String issuer) {
         AccessToken token = new AccessToken();
         token.id(KeycloakModelUtils.generateId());
         token.issuedNow();
@@ -489,13 +490,21 @@ public class AuthenticationManager {
         if (realm.getSsoSessionMaxLifespan() > 0) {
             token.expiration(Time.currentTime() + realm.getSsoSessionMaxLifespan());
         }
+
+        String stateChecker = (String) keycloakSession.getAttribute("state_checker");
+        if (stateChecker == null) {
+            stateChecker = Base64Url.encode(KeycloakModelUtils.generateSecret());
+            keycloakSession.setAttribute("state_checker", stateChecker);
+        }
+        token.getOtherClaims().put("state_checker", stateChecker);
+
         return token;
     }
 
     public static void createLoginCookie(KeycloakSession keycloakSession, RealmModel realm, UserModel user, UserSessionModel session, UriInfo uriInfo, ClientConnection connection) {
         String cookiePath = getIdentityCookiePath(realm, uriInfo);
         String issuer = Urls.realmIssuer(uriInfo.getBaseUri(), realm.getName());
-        AccessToken identityToken = createIdentityToken(realm, user, session, issuer);
+        AccessToken identityToken = createIdentityToken(keycloakSession, realm, user, session, issuer);
         String encoded = encodeToken(keycloakSession, realm, identityToken);
         boolean secureOnly = realm.getSslRequired().isRequired(connection);
         int maxAge = NewCookie.DEFAULT_MAX_AGE;
@@ -1044,6 +1053,8 @@ public class AuthenticationManager {
                 return null;
             }
 
+            session.setAttribute("state_checker", token.getOtherClaims().get("state_checker"));
+
             return new AuthResult(user, userSession, token);
         } catch (VerificationException e) {
             logger.debugf("Failed to verify identity token: %s", e.getMessage());

diff --git a/services/src/main/java/org/keycloak/services/resources/AbstractSecuredLocalService.java b/services/src/main/java/org/keycloak/services/resources/AbstractSecuredLocalService.java
index bcc99b6..5d2eab5 100755
--- a/services/src/main/java/org/keycloak/services/resources/AbstractSecuredLocalService.java
+++ b/services/src/main/java/org/keycloak/services/resources/AbstractSecuredLocalService.java
@@ -58,8 +58,6 @@ import java.util.Set;
 public abstract class AbstractSecuredLocalService {
     private static final Logger logger = Logger.getLogger(AbstractSecuredLocalService.class);
 
-    private static final String KEYCLOAK_STATE_CHECKER = "KEYCLOAK_STATE_CHECKER";
-
     protected final ClientModel client;
     protected RealmModel realm;
 
@@ -127,55 +125,8 @@ public abstract class AbstractSecuredLocalService {
         }
     }
 
-    protected void updateCsrfChecks() {
-        stateChecker = getStateChecker();
-        if (stateChecker == null) {
-            stateChecker = Base64Url.encode(KeycloakModelUtils.generateSecret());
-
-            StringBuilder sb = new StringBuilder();
-            sb.append(auth.getSession().getId());
-            sb.append("/");
-            sb.append(stateChecker);
-
-            String sessionCookieValue = sb.toString();
-
-            String cookiePath = AuthenticationManager.getAccountCookiePath(realm, uriInfo);
-            boolean secureOnly = realm.getSslRequired().isRequired(clientConnection);
-            CookieHelper.addCookie(KEYCLOAK_STATE_CHECKER, sessionCookieValue, cookiePath, null, null, -1, secureOnly, true);
-        }
-    }
-
     protected abstract Set<String> getValidPaths();
 
-    /**
-     * Check to see if form post has sessionId hidden field and match it against the session id.
-     *
-     * @param formData
-     */
-    protected void csrfCheck(final MultivaluedMap<String, String> formData) {
-        String stateChecker = formData.getFirst("stateChecker");
-        if (stateChecker == null || !stateChecker.equals(getStateChecker())) {
-            throw new ForbiddenException();
-        }
-    }
-
-    protected String getStateChecker() {
-        Cookie cookie = headers.getCookies().get(KEYCLOAK_STATE_CHECKER);
-        if (cookie != null) {
-            stateChecker = cookie.getValue();
-            String[] s = stateChecker.split("/");
-            if (s.length == 2) {
-                String sessionId = s[0];
-                String stateChecker = s[1];
-
-                if (auth.getSession().getId().equals(sessionId)) {
-                    return stateChecker;
-                }
-            }
-        }
-        return null;
-    }
-
     protected abstract URI getBaseRedirectUri();
 
     protected Response login(String path) {

diff --git a/services/src/main/java/org/keycloak/services/resources/account/AccountFormService.java b/services/src/main/java/org/keycloak/services/resources/account/AccountFormService.java
index 3cfb6d3..06b3758 100755
--- a/services/src/main/java/org/keycloak/services/resources/account/AccountFormService.java
+++ b/services/src/main/java/org/keycloak/services/resources/account/AccountFormService.java
@@ -120,8 +120,8 @@ public class AccountFormService extends AbstractSecuredLocalService {
 
         AuthenticationManager.AuthResult authResult = authManager.authenticateIdentityCookie(session, realm);
         if (authResult != null) {
+            stateChecker = (String) session.getAttribute("state_checker");
             auth = new Auth(realm, authResult.getToken(), authResult.getUser(), client, authResult.getSession(), true);
-            updateCsrfChecks();
             account.setStateChecker(stateChecker);
         }
 
@@ -798,4 +798,12 @@ public class AccountFormService extends AbstractSecuredLocalService {
         }
     }
 
+    private void csrfCheck(final MultivaluedMap<String, String> formData) {
+        String formStateChecker = formData.getFirst("stateChecker");
+        if (formStateChecker == null || !formStateChecker.equals(this.stateChecker)) {
+            throw new ForbiddenException();
+        }
+    }
+
+
 }