keycloak-aplcache

KEYCLOAK-1070 add setProtocolMappers/getProtocolMappers

4/22/2015 12:06:14 PM

Changes

Details

diff --git a/connections/jpa/src/main/resources/META-INF/persistence.xml b/connections/jpa/src/main/resources/META-INF/persistence.xml
index b4d93a7..83eadb5 100755
--- a/connections/jpa/src/main/resources/META-INF/persistence.xml
+++ b/connections/jpa/src/main/resources/META-INF/persistence.xml
@@ -27,6 +27,7 @@
         <!-- JpaUserSessionProvider -->
         <class>org.keycloak.models.sessions.jpa.entities.ClientSessionEntity</class>
         <class>org.keycloak.models.sessions.jpa.entities.ClientSessionRoleEntity</class>
+        <class>org.keycloak.models.sessions.jpa.entities.ClientSessionProtocolMapperEntity</class>
         <class>org.keycloak.models.sessions.jpa.entities.ClientSessionNoteEntity</class>
         <class>org.keycloak.models.sessions.jpa.entities.UserSessionNoteEntity</class>
         <class>org.keycloak.models.sessions.jpa.entities.UserSessionEntity</class>
diff --git a/connections/jpa-liquibase/src/main/resources/META-INF/jpa-changelog-1.2.0.RC1.xml b/connections/jpa-liquibase/src/main/resources/META-INF/jpa-changelog-1.2.0.RC1.xml
index 7b8c877..bc808bc 100755
--- a/connections/jpa-liquibase/src/main/resources/META-INF/jpa-changelog-1.2.0.RC1.xml
+++ b/connections/jpa-liquibase/src/main/resources/META-INF/jpa-changelog-1.2.0.RC1.xml
@@ -62,16 +62,26 @@
                 <constraints nullable="false"/>
             </column>
         </createTable>
+        <createTable tableName="CLIENT_SESSION_PROT_MAPPER">
+            <column name="PROTOCOL_MAPPER_ID" type="VARCHAR(36)">
+                <constraints nullable="false"/>
+            </column>
+            <column name="CLIENT_SESSION" type="VARCHAR(36)">
+                <constraints nullable="false"/>
+            </column>
+        </createTable>
         <addPrimaryKey columnNames="ID" constraintName="CONSTRAINT_IDPM" tableName="IDENTITY_PROVIDER_MAPPER"/>
         <addPrimaryKey columnNames="IDP_MAPPER_ID, NAME" constraintName="CONSTRAINT_IDPMConfig" tableName="IDP_MAPPER_CONFIG"/>
         <addPrimaryKey columnNames="ID" constraintName="CONSTRAINT_GRNTCSNT_PM" tableName="USER_CONSENT"/>
         <addPrimaryKey columnNames="USER_CONSENT_ID, ROLE_ID" constraintName="CONSTRAINT_GRNTCSNT_ROLE_PM" tableName="USER_CONSENT_ROLE"/>
         <addPrimaryKey columnNames="USER_CONSENT_ID, PROTOCOL_MAPPER_ID" constraintName="CONSTRAINT_GRNTCSNT_PRM_PM" tableName="USER_CONSENT_PROT_MAPPER"/>
+        <addPrimaryKey columnNames="CLIENT_SESSION, PROTOCOL_MAPPER_ID" constraintName="CONSTRAINT_CS_PMP_PK" tableName="CLIENT_SESSION_PROT_MAPPER"/>
         <addForeignKeyConstraint baseColumnNames="REALM_ID" baseTableName="IDENTITY_PROVIDER_MAPPER" constraintName="FK_IDPM_REALM" referencedColumnNames="ID" referencedTableName="REALM"/>
         <addForeignKeyConstraint baseColumnNames="IDP_MAPPER_ID" baseTableName="IDP_MAPPER_CONFIG" constraintName="FK_IDPMConfig" referencedColumnNames="ID" referencedTableName="IDENTITY_PROVIDER_MAPPER"/>
         <addForeignKeyConstraint baseColumnNames="USER_ID" baseTableName="USER_CONSENT" constraintName="FK_GRNTCSNT_USER" referencedColumnNames="ID" referencedTableName="USER_ENTITY"/>
         <addForeignKeyConstraint baseColumnNames="USER_CONSENT_ID" baseTableName="USER_CONSENT_ROLE" constraintName="FK_GRNTCSNT_ROLE_GR" referencedColumnNames="ID" referencedTableName="USER_CONSENT"/>
         <addForeignKeyConstraint baseColumnNames="USER_CONSENT_ID" baseTableName="USER_CONSENT_PROT_MAPPER" constraintName="FK_GRNTCSNT_PRM_GR" referencedColumnNames="ID" referencedTableName="USER_CONSENT"/>
+        <addForeignKeyConstraint baseColumnNames="CLIENT_SESSION" baseTableName="CLIENT_SESSION_PROT_MAPPER" constraintName="FK_33A8SGQW18I532811V7O2DK89" referencedColumnNames="ID" referencedTableName="CLIENT_SESSION"/>
 
         <addColumn tableName="CLIENT">
             <column name="CONSENT_REQUIRED" type="BOOLEAN" defaultValueBoolean="false">
diff --git a/connections/mongo-update/src/main/java/org/keycloak/connections/mongo/updater/impl/updates/Update1_2_0_RC1.java b/connections/mongo-update/src/main/java/org/keycloak/connections/mongo/updater/impl/updates/Update1_2_0_RC1.java
index 8ee4db3..811039c 100644
--- a/connections/mongo-update/src/main/java/org/keycloak/connections/mongo/updater/impl/updates/Update1_2_0_RC1.java
+++ b/connections/mongo-update/src/main/java/org/keycloak/connections/mongo/updater/impl/updates/Update1_2_0_RC1.java
@@ -17,6 +17,9 @@ public class Update1_2_0_RC1 extends Update {
 
     @Override
     public void update(KeycloakSession session) {
+        deleteEntries("clientSessions");
+        deleteEntries("sessions");
+
         convertApplicationsToClients();
         convertOAuthClientsToClients();
 
diff --git a/model/api/src/main/java/org/keycloak/models/ClientSessionModel.java b/model/api/src/main/java/org/keycloak/models/ClientSessionModel.java
index 059afc9..7c197a3 100755
--- a/model/api/src/main/java/org/keycloak/models/ClientSessionModel.java
+++ b/model/api/src/main/java/org/keycloak/models/ClientSessionModel.java
@@ -28,6 +28,9 @@ public interface ClientSessionModel {
     public Set<String> getRoles();
     public void setRoles(Set<String> roles);
 
+    public Set<String> getProtocolMappers();
+    public void setProtocolMappers(Set<String> protocolMappers);
+
     /**
      * Authentication request type, i.e. OAUTH, SAML 2.0, SAML 1.1, etc.
      *
diff --git a/model/api/src/main/java/org/keycloak/models/UserConsentModel.java b/model/api/src/main/java/org/keycloak/models/UserConsentModel.java
index 11f1034..a8e4cc3 100644
--- a/model/api/src/main/java/org/keycloak/models/UserConsentModel.java
+++ b/model/api/src/main/java/org/keycloak/models/UserConsentModel.java
@@ -8,31 +8,20 @@ import java.util.Set;
  */
 public class UserConsentModel {
 
-    private final RealmModel realm;
     private final ClientModel client;
     private Set<ProtocolMapperModel> protocolMappers = new HashSet<ProtocolMapperModel>();
     private Set<RoleModel> roles = new HashSet<RoleModel>();
 
-    public UserConsentModel(RealmModel realm, String clientId) {
-        this.realm = realm;
-        this.client = realm.getClientById(clientId);
-
-        if (client == null) {
-            throw new ModelException("Client with id [" + clientId + "] is not available");
-        }
+    public UserConsentModel(ClientModel client) {
+        this.client = client;
     }
 
     public ClientModel getClient() {
         return client;
     }
 
-    public void addGrantedRole(String roleId) {
-        RoleModel role = realm.getRoleById(roleId);
-
-        // Chance that role was already deleted by other transaction and is not available anymore
-        if (role != null) {
-            roles.add(role);
-        }
+    public void addGrantedRole(RoleModel role) {
+        roles.add(role);
     }
 
     public Set<RoleModel> getGrantedRoles() {
@@ -46,13 +35,8 @@ public class UserConsentModel {
         return false;
     }
 
-    public void addGrantedProtocolMapper(String protocolMapperId) {
-        ProtocolMapperModel protocolMapper = client.getProtocolMapperById(protocolMapperId);
-
-        // Chance that protocolMapper was already deleted by other transaction and is not available anymore
-        if (protocolMapper != null) {
-            protocolMappers.add(protocolMapper);
-        }
+    public void addGrantedProtocolMapper(ProtocolMapperModel protocolMapper) {
+        protocolMappers.add(protocolMapper);
     }
 
     public Set<ProtocolMapperModel> getGrantedProtocolMappers() {
diff --git a/model/api/src/main/java/org/keycloak/models/utils/RepresentationToModel.java b/model/api/src/main/java/org/keycloak/models/utils/RepresentationToModel.java
index 282d376..7043611 100755
--- a/model/api/src/main/java/org/keycloak/models/utils/RepresentationToModel.java
+++ b/model/api/src/main/java/org/keycloak/models/utils/RepresentationToModel.java
@@ -792,23 +792,25 @@ public class RepresentationToModel {
                     throw new RuntimeException("Unable to find client consent mappings for client: " + entry.getKey());
                 }
 
-                UserConsentModel consentModel = new UserConsentModel(newRealm, client.getId());
+                UserConsentModel consentModel = new UserConsentModel(client);
 
                 UserConsentRepresentation consentRep = entry.getValue();
                 if (consentRep.getGrantedRoles() != null) {
                     for (String roleId : consentRep.getGrantedRoles()) {
-                        if (newRealm.getRoleById(roleId) == null) {
+                        RoleModel role = newRealm.getRoleById(roleId);
+                        if (role == null) {
                             throw new RuntimeException("Unable to find realm role referenced in consent mappings of user " + user.getUsername() + ". Role ID: " + roleId);
                         }
-                        consentModel.addGrantedRole(roleId);
+                        consentModel.addGrantedRole(role);
                     }
                 }
                 if (consentRep.getGrantedProtocolMappers() != null) {
                     for (String mapperId : consentRep.getGrantedProtocolMappers()) {
-                        if (client.getProtocolMapperById(mapperId) == null) {
+                        ProtocolMapperModel protocolMapper = client.getProtocolMapperById(mapperId);
+                        if (protocolMapper == null) {
                             throw new RuntimeException("Unable to find protocol mapper referenced in consent mappings of user " + user.getUsername() + ". Protocol mapper ID: " + mapperId);
                         }
-                        consentModel.addGrantedProtocolMapper(mapperId);;
+                        consentModel.addGrantedProtocolMapper(protocolMapper);
                     }
                 }
                 user.addConsent(consentModel);
diff --git a/model/jpa/src/main/java/org/keycloak/models/jpa/UserAdapter.java b/model/jpa/src/main/java/org/keycloak/models/jpa/UserAdapter.java
index 8cc29f7..5f92104 100755
--- a/model/jpa/src/main/java/org/keycloak/models/jpa/UserAdapter.java
+++ b/model/jpa/src/main/java/org/keycloak/models/jpa/UserAdapter.java
@@ -561,19 +561,27 @@ public class UserAdapter implements UserModel {
             return null;
         }
 
-        UserConsentModel model = new UserConsentModel(realm, entity.getClientId());
+        ClientModel client = realm.getClientById(entity.getClientId());
+        if (client == null) {
+            throw new ModelException("Client with id " + entity.getClientId() + " is not available");
+        }
+        UserConsentModel model = new UserConsentModel(client);
 
         Collection<UserConsentRoleEntity> grantedRoleEntities = entity.getGrantedRoles();
         if (grantedRoleEntities != null) {
             for (UserConsentRoleEntity grantedRole : grantedRoleEntities) {
-                model.addGrantedRole(grantedRole.getRoleId());
+                RoleModel grantedRoleModel = realm.getRoleById(grantedRole.getRoleId());
+                if (grantedRoleModel != null) {
+                    model.addGrantedRole(grantedRoleModel);
+                }
             }
         }
 
         Collection<UserConsentProtocolMapperEntity> grantedProtocolMapperEntities = entity.getGrantedProtocolMappers();
         if (grantedProtocolMapperEntities != null) {
             for (UserConsentProtocolMapperEntity grantedProtMapper : grantedProtocolMapperEntities) {
-                model.addGrantedProtocolMapper(grantedProtMapper.getProtocolMapperId());
+                ProtocolMapperModel protocolMapper = client.getProtocolMapperById(grantedProtMapper.getProtocolMapperId());
+                model.addGrantedProtocolMapper(protocolMapper );
             }
         }
 
diff --git a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/UserAdapter.java b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/UserAdapter.java
index b054793..1a4a9b5 100755
--- a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/UserAdapter.java
+++ b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/UserAdapter.java
@@ -462,12 +462,22 @@ public class UserAdapter extends AbstractMongoAdapter<MongoUserEntity> implement
     }
 
     private UserConsentModel toConsentModel(UserConsentEntity entity) {
-        UserConsentModel model = new UserConsentModel(realm, entity.getClientId());
+        ClientModel client = realm.getClientById(entity.getClientId());
+        if (client == null) {
+            throw new ModelException("Client with id " + entity.getClientId() + " is not available");
+        }
+        UserConsentModel model = new UserConsentModel(client);
+
         for (String roleId : entity.getGrantedRoles()) {
-            model.addGrantedRole(roleId);
+            RoleModel roleModel = realm.getRoleById(roleId);
+            if (roleModel != null) {
+                model.addGrantedRole(roleModel);
+            }
         }
+
         for (String protMapperId : entity.getGrantedProtocolMappers()) {
-            model.addGrantedProtocolMapper(protMapperId);
+            ProtocolMapperModel protocolMapper = client.getProtocolMapperById(protMapperId);
+            model.addGrantedProtocolMapper(protocolMapper);
         }
         return model;
     }
diff --git a/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/ClientSessionAdapter.java b/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/ClientSessionAdapter.java
index e380441..cc2c8cc 100644
--- a/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/ClientSessionAdapter.java
+++ b/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/ClientSessionAdapter.java
@@ -119,6 +119,17 @@ public class ClientSessionAdapter implements ClientSessionModel {
     }
 
     @Override
+    public Set<String> getProtocolMappers() {
+        return entity.getProtocolMappers();
+    }
+
+    @Override
+    public void setProtocolMappers(Set<String> protocolMappers) {
+        entity.setProtocolMappers(protocolMappers);
+        update();
+    }
+
+    @Override
     public String getAuthMethod() {
         return entity.getAuthMethod();
     }
diff --git a/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/ClientSessionEntity.java b/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/ClientSessionEntity.java
index 9c085fc..38d2ea5 100755
--- a/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/ClientSessionEntity.java
+++ b/model/sessions-infinispan/src/main/java/org/keycloak/models/sessions/infinispan/entities/ClientSessionEntity.java
@@ -25,6 +25,7 @@ public class ClientSessionEntity extends SessionEntity {
     private ClientSessionModel.Action action;
 
     private Set<String> roles;
+    private Set<String> protocolMappers;
     private Map<String, String> notes;
 
     public String getClient() {
@@ -91,6 +92,14 @@ public class ClientSessionEntity extends SessionEntity {
         this.roles = roles;
     }
 
+    public Set<String> getProtocolMappers() {
+        return protocolMappers;
+    }
+
+    public void setProtocolMappers(Set<String> protocolMappers) {
+        this.protocolMappers = protocolMappers;
+    }
+
     public Map<String, String> getNotes() {
         return notes;
     }
diff --git a/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/ClientSessionAdapter.java b/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/ClientSessionAdapter.java
index ebb47b6..74e795f 100755
--- a/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/ClientSessionAdapter.java
+++ b/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/ClientSessionAdapter.java
@@ -7,6 +7,7 @@ import org.keycloak.models.RealmModel;
 import org.keycloak.models.UserSessionModel;
 import org.keycloak.models.sessions.jpa.entities.ClientSessionEntity;
 import org.keycloak.models.sessions.jpa.entities.ClientSessionNoteEntity;
+import org.keycloak.models.sessions.jpa.entities.ClientSessionProtocolMapperEntity;
 import org.keycloak.models.sessions.jpa.entities.ClientSessionRoleEntity;
 import org.keycloak.models.sessions.jpa.entities.UserSessionEntity;
 
@@ -177,4 +178,36 @@ public class ClientSessionAdapter implements ClientSessionModel {
         }
         return roles;
     }
+
+    @Override
+    public Set<String> getProtocolMappers() {
+        Set<String> protMappers = new HashSet<String>();
+        if (entity.getProtocolMappers() != null) {
+            for (ClientSessionProtocolMapperEntity e : entity.getProtocolMappers()) {
+                protMappers.add(e.getProtocolMapperId());
+            }
+        }
+        return protMappers;
+    }
+
+    @Override
+    public void setProtocolMappers(Set<String> protocolMappers) {
+        if (protocolMappers != null) {
+            for (String pm : protocolMappers) {
+                ClientSessionProtocolMapperEntity protMapperEntity = new ClientSessionProtocolMapperEntity();
+                protMapperEntity.setClientSession(entity);
+                protMapperEntity.setProtocolMapperId(pm);
+                em.persist(protMapperEntity);
+
+                entity.getProtocolMappers().add(protMapperEntity);
+            }
+        } else {
+            if (entity.getProtocolMappers() != null) {
+                for (ClientSessionProtocolMapperEntity pm : entity.getProtocolMappers()) {
+                    em.remove(pm);
+                }
+                entity.getProtocolMappers().clear();
+            }
+        }
+    }
 }
diff --git a/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/entities/ClientSessionEntity.java b/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/entities/ClientSessionEntity.java
index c4bbe06..6f0c535 100755
--- a/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/entities/ClientSessionEntity.java
+++ b/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/entities/ClientSessionEntity.java
@@ -61,6 +61,9 @@ public class ClientSessionEntity {
     protected Collection<ClientSessionRoleEntity> roles = new ArrayList<ClientSessionRoleEntity>();
 
     @OneToMany(cascade = CascadeType.REMOVE, orphanRemoval = true, mappedBy="clientSession")
+    protected Collection<ClientSessionProtocolMapperEntity> protocolMappers = new ArrayList<ClientSessionProtocolMapperEntity>();
+
+    @OneToMany(cascade = CascadeType.REMOVE, orphanRemoval = true, mappedBy="clientSession")
     protected Collection<ClientSessionNoteEntity> notes = new ArrayList<ClientSessionNoteEntity>();
 
     public String getId() {
@@ -127,6 +130,14 @@ public class ClientSessionEntity {
         this.roles = roles;
     }
 
+    public Collection<ClientSessionProtocolMapperEntity> getProtocolMappers() {
+        return protocolMappers;
+    }
+
+    public void setProtocolMappers(Collection<ClientSessionProtocolMapperEntity> protocolMappers) {
+        this.protocolMappers = protocolMappers;
+    }
+
     public Collection<ClientSessionNoteEntity> getNotes() {
         return notes;
     }
diff --git a/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/entities/ClientSessionProtocolMapperEntity.java b/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/entities/ClientSessionProtocolMapperEntity.java
new file mode 100644
index 0000000..628e0f3
--- /dev/null
+++ b/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/entities/ClientSessionProtocolMapperEntity.java
@@ -0,0 +1,98 @@
+package org.keycloak.models.sessions.jpa.entities;
+
+import java.io.Serializable;
+
+import javax.persistence.Column;
+import javax.persistence.Entity;
+import javax.persistence.FetchType;
+import javax.persistence.Id;
+import javax.persistence.IdClass;
+import javax.persistence.JoinColumn;
+import javax.persistence.ManyToOne;
+import javax.persistence.NamedQueries;
+import javax.persistence.NamedQuery;
+import javax.persistence.Table;
+
+/**
+ * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
+ */
+@NamedQueries({
+        @NamedQuery(name = "removeClientSessionProtMapperByUser", query="delete from ClientSessionProtocolMapperEntity pm where pm.clientSession IN (select c from ClientSessionEntity c where c.session IN (select s from UserSessionEntity s where s.realmId = :realmId and s.userId = :userId))"),
+        @NamedQuery(name = "removeClientSessionProtMapperByClient", query="delete from ClientSessionProtocolMapperEntity pm where pm.clientSession IN (select c from ClientSessionEntity c where c.clientId = :clientId and c.realmId = :realmId)"),
+        @NamedQuery(name = "removeClientSessionProtMapperByRealm", query="delete from ClientSessionProtocolMapperEntity pm where pm.clientSession IN (select c from ClientSessionEntity c where c.realmId = :realmId)"),
+        @NamedQuery(name = "removeClientSessionProtMapperByExpired", query = "delete from ClientSessionProtocolMapperEntity pm where pm.clientSession IN (select c from ClientSessionEntity c where c.session IN (select s from UserSessionEntity s where s.realmId = :realmId and (s.started < :maxTime or s.lastSessionRefresh < :idleTime)))"),
+        @NamedQuery(name = "removeDetachedClientSessionProtMapperByExpired", query = "delete from ClientSessionProtocolMapperEntity pm where pm.clientSession IN (select c from ClientSessionEntity c where c.session IS NULL and c.realmId = :realmId and c.timestamp < :maxTime )")
+})
+@Table(name="CLIENT_SESSION_PROT_MAPPER")
+@Entity
+@IdClass(ClientSessionProtocolMapperEntity.Key.class)
+public class ClientSessionProtocolMapperEntity {
+
+    @Id
+    @ManyToOne(fetch= FetchType.LAZY)
+    @JoinColumn(name="CLIENT_SESSION")
+    protected ClientSessionEntity clientSession;
+
+    @Id
+    @Column(name="PROTOCOL_MAPPER_ID")
+    protected String protocolMapperId;
+
+    public ClientSessionEntity getClientSession() {
+        return clientSession;
+    }
+
+    public void setClientSession(ClientSessionEntity clientSession) {
+        this.clientSession = clientSession;
+    }
+
+    public String getProtocolMapperId() {
+        return protocolMapperId;
+    }
+
+    public void setProtocolMapperId(String protocolMapperId) {
+        this.protocolMapperId = protocolMapperId;
+    }
+
+    public static class Key implements Serializable {
+
+        protected ClientSessionEntity clientSession;
+
+        protected String protocolMapperId;
+
+        public Key() {
+        }
+
+        public Key(ClientSessionEntity clientSession, String protocolMapperId) {
+            this.clientSession = clientSession;
+            this.protocolMapperId = protocolMapperId;
+        }
+
+        public ClientSessionEntity getClientSession() {
+            return clientSession;
+        }
+
+        public String getProtocolMapperId() {
+            return protocolMapperId;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+
+            Key key = (Key) o;
+
+            if (!protocolMapperId.equals(key.protocolMapperId)) return false;
+            if (!clientSession.getId().equals(key.clientSession.getId())) return false;
+
+            return true;
+        }
+
+        @Override
+        public int hashCode() {
+            int result = clientSession.getId().hashCode();
+            result = 31 * result + protocolMapperId.hashCode();
+            return result;
+        }
+    }
+}
diff --git a/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/JpaUserSessionProvider.java b/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/JpaUserSessionProvider.java
index 0731de3..23086d9 100755
--- a/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/JpaUserSessionProvider.java
+++ b/model/sessions-jpa/src/main/java/org/keycloak/models/sessions/jpa/JpaUserSessionProvider.java
@@ -214,6 +214,10 @@ public class JpaUserSessionProvider implements UserSessionProvider {
                 .setParameter("realmId", realm.getId())
                 .setParameter("userId", user.getId())
                 .executeUpdate();
+        em.createNamedQuery("removeClientSessionProtMapperByUser")
+                .setParameter("realmId", realm.getId())
+                .setParameter("userId", user.getId())
+                .executeUpdate();
         em.createNamedQuery("removeClientSessionByUser")
                 .setParameter("realmId", realm.getId())
                 .setParameter("userId", user.getId())
@@ -238,6 +242,10 @@ public class JpaUserSessionProvider implements UserSessionProvider {
                 .setParameter("realmId", realm.getId())
                 .setParameter("maxTime", dettachedClientSessionExpired)
                 .executeUpdate();
+        em.createNamedQuery("removeDetachedClientSessionProtMapperByExpired")
+                .setParameter("realmId", realm.getId())
+                .setParameter("maxTime", dettachedClientSessionExpired)
+                .executeUpdate();
         em.createNamedQuery("removeDetachedClientSessionNoteByExpired")
                 .setParameter("realmId", realm.getId())
                 .setParameter("maxTime", dettachedClientSessionExpired)
@@ -251,6 +259,11 @@ public class JpaUserSessionProvider implements UserSessionProvider {
                 .setParameter("maxTime", maxTime)
                 .setParameter("idleTime", idleTime)
                 .executeUpdate();
+        em.createNamedQuery("removeClientSessionProtMapperByExpired")
+                .setParameter("realmId", realm.getId())
+                .setParameter("maxTime", maxTime)
+                .setParameter("idleTime", idleTime)
+                .executeUpdate();
         em.createNamedQuery("removeClientSessionNoteByExpired")
                 .setParameter("realmId", realm.getId())
                 .setParameter("maxTime", maxTime)
@@ -277,6 +290,7 @@ public class JpaUserSessionProvider implements UserSessionProvider {
     public void removeUserSessions(RealmModel realm) {
         em.createNamedQuery("removeClientSessionNoteByRealm").setParameter("realmId", realm.getId()).executeUpdate();
         em.createNamedQuery("removeClientSessionRoleByRealm").setParameter("realmId", realm.getId()).executeUpdate();
+        em.createNamedQuery("removeClientSessionProtMapperByRealm").setParameter("realmId", realm.getId()).executeUpdate();
         em.createNamedQuery("removeClientSessionByRealm").setParameter("realmId", realm.getId()).executeUpdate();
         em.createNamedQuery("removeUserSessionNoteByRealm").setParameter("realmId", realm.getId()).executeUpdate();
         em.createNamedQuery("removeUserSessionByRealm").setParameter("realmId", realm.getId()).executeUpdate();
@@ -292,6 +306,7 @@ public class JpaUserSessionProvider implements UserSessionProvider {
     public void onClientRemoved(RealmModel realm, ClientModel client) {
         em.createNamedQuery("removeClientSessionNoteByClient").setParameter("realmId", realm.getId()).setParameter("clientId", client.getId()).executeUpdate();
         em.createNamedQuery("removeClientSessionRoleByClient").setParameter("realmId", realm.getId()).setParameter("clientId", client.getId()).executeUpdate();
+        em.createNamedQuery("removeClientSessionProtMapperByClient").setParameter("realmId", realm.getId()).setParameter("clientId", client.getId()).executeUpdate();
         em.createNamedQuery("removeClientSessionByClient").setParameter("realmId", realm.getId()).setParameter("clientId", client.getId()).executeUpdate();
     }
 
diff --git a/model/sessions-mem/src/main/java/org/keycloak/models/sessions/mem/ClientSessionAdapter.java b/model/sessions-mem/src/main/java/org/keycloak/models/sessions/mem/ClientSessionAdapter.java
index 62a3246..daa6e3f 100755
--- a/model/sessions-mem/src/main/java/org/keycloak/models/sessions/mem/ClientSessionAdapter.java
+++ b/model/sessions-mem/src/main/java/org/keycloak/models/sessions/mem/ClientSessionAdapter.java
@@ -106,6 +106,16 @@ public class ClientSessionAdapter implements ClientSessionModel {
     }
 
     @Override
+    public Set<String> getProtocolMappers() {
+        return entity.getProtocolMappers();
+    }
+
+    @Override
+    public void setProtocolMappers(Set<String> protocolMappers) {
+        entity.setProtocolMappers(protocolMappers);
+    }
+
+    @Override
     public String getNote(String name) {
         return entity.getNotes().get(name);
     }
diff --git a/model/sessions-mem/src/main/java/org/keycloak/models/sessions/mem/entities/ClientSessionEntity.java b/model/sessions-mem/src/main/java/org/keycloak/models/sessions/mem/entities/ClientSessionEntity.java
index dd945de..ca30496 100755
--- a/model/sessions-mem/src/main/java/org/keycloak/models/sessions/mem/entities/ClientSessionEntity.java
+++ b/model/sessions-mem/src/main/java/org/keycloak/models/sessions/mem/entities/ClientSessionEntity.java
@@ -23,6 +23,7 @@ public class ClientSessionEntity {
     private int timestamp;
     private ClientSessionModel.Action action;
     private Set<String> roles;
+    private Set<String> protocolMappers;
     private Map<String, String> notes = new HashMap<String, String>();
 
     public String getId() {
@@ -89,6 +90,14 @@ public class ClientSessionEntity {
         this.roles = roles;
     }
 
+    public Set<String> getProtocolMappers() {
+        return protocolMappers;
+    }
+
+    public void setProtocolMappers(Set<String> protocolMappers) {
+        this.protocolMappers = protocolMappers;
+    }
+
     public Map<String, String> getNotes() {
         return notes;
     }
diff --git a/model/sessions-mongo/src/main/java/org/keycloak/models/sessions/mongo/ClientSessionAdapter.java b/model/sessions-mongo/src/main/java/org/keycloak/models/sessions/mongo/ClientSessionAdapter.java
index 12efd4a..649ffcf 100755
--- a/model/sessions-mongo/src/main/java/org/keycloak/models/sessions/mongo/ClientSessionAdapter.java
+++ b/model/sessions-mongo/src/main/java/org/keycloak/models/sessions/mongo/ClientSessionAdapter.java
@@ -121,6 +121,23 @@ public class ClientSessionAdapter extends AbstractMongoAdapter<MongoClientSessio
     }
 
     @Override
+    public Set<String> getProtocolMappers() {
+        return entity.getProtocolMappers() != null ? new HashSet<String>(entity.getProtocolMappers()) : null;
+    }
+
+    @Override
+    public void setProtocolMappers(Set<String> protocolMappers) {
+        if (protocolMappers == null) {
+            entity.setProtocolMappers(null);
+        } else {
+            List<String> list = new LinkedList<String>();
+            list.addAll(protocolMappers);
+            entity.setProtocolMappers(list);
+        }
+        updateMongoEntity();
+    }
+
+    @Override
     public String getNote(String name) {
         return entity.getNotes().get(name);
     }
diff --git a/model/sessions-mongo/src/main/java/org/keycloak/models/sessions/mongo/entities/MongoClientSessionEntity.java b/model/sessions-mongo/src/main/java/org/keycloak/models/sessions/mongo/entities/MongoClientSessionEntity.java
index 91ec682..b5abede 100755
--- a/model/sessions-mongo/src/main/java/org/keycloak/models/sessions/mongo/entities/MongoClientSessionEntity.java
+++ b/model/sessions-mongo/src/main/java/org/keycloak/models/sessions/mongo/entities/MongoClientSessionEntity.java
@@ -27,6 +27,7 @@ public class MongoClientSessionEntity extends AbstractIdentifiableEntity impleme
     private int timestamp;
     private ClientSessionModel.Action action;
     private List<String> roles;
+    private List<String> protocolMappers;
     private Map<String, String> notes = new HashMap<String, String>();
 
     public String getId() {
@@ -93,6 +94,14 @@ public class MongoClientSessionEntity extends AbstractIdentifiableEntity impleme
         this.roles = roles;
     }
 
+    public List<String> getProtocolMappers() {
+        return protocolMappers;
+    }
+
+    public void setProtocolMappers(List<String> protocolMappers) {
+        this.protocolMappers = protocolMappers;
+    }
+
     public Map<String, String> getNotes() {
         return notes;
     }
diff --git a/saml/saml-protocol/src/main/java/org/keycloak/protocol/saml/mappers/RoleListMapper.java b/saml/saml-protocol/src/main/java/org/keycloak/protocol/saml/mappers/RoleListMapper.java
index 677b2bf..f4d5bd4 100755
--- a/saml/saml-protocol/src/main/java/org/keycloak/protocol/saml/mappers/RoleListMapper.java
+++ b/saml/saml-protocol/src/main/java/org/keycloak/protocol/saml/mappers/RoleListMapper.java
@@ -11,12 +11,14 @@ import org.keycloak.provider.ProviderConfigProperty;
 import org.keycloak.protocol.saml.SamlProtocol;
 import org.keycloak.dom.saml.v2.assertion.AttributeStatementType;
 import org.keycloak.dom.saml.v2.assertion.AttributeType;
+import org.keycloak.services.managers.ClientSessionCode;
 
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 /**
  * @author <a href="mailto:bill@burkecentral.com">Bill Burke</a>
@@ -96,8 +98,8 @@ public class RoleListMapper extends AbstractSAMLProtocolMapper implements SAMLRo
         List<SamlProtocol.ProtocolMapperProcessor<SAMLRoleNameMapper>> roleNameMappers = new LinkedList<>();
         KeycloakSessionFactory sessionFactory = session.getKeycloakSessionFactory();
         AttributeType singleAttributeType = null;
-        for (ProtocolMapperModel mapping : clientSession.getClient().getProtocolMappers()) {
-            if (!mapping.getProtocol().equals(SamlProtocol.LOGIN_PROTOCOL)) continue;
+        Set<ProtocolMapperModel> requestedProtocolMappers = new ClientSessionCode(clientSession.getRealm(), clientSession).getRequestedProtocolMappers();
+        for (ProtocolMapperModel mapping : requestedProtocolMappers) {
 
             ProtocolMapper mapper = (ProtocolMapper)sessionFactory.getProviderFactory(ProtocolMapper.class, mapping.getProtocolMapper());
             if (mapper == null) continue;
diff --git a/saml/saml-protocol/src/main/java/org/keycloak/protocol/saml/SamlProtocol.java b/saml/saml-protocol/src/main/java/org/keycloak/protocol/saml/SamlProtocol.java
index 747c43c..315e685 100755
--- a/saml/saml-protocol/src/main/java/org/keycloak/protocol/saml/SamlProtocol.java
+++ b/saml/saml-protocol/src/main/java/org/keycloak/protocol/saml/SamlProtocol.java
@@ -282,9 +282,8 @@ public class SamlProtocol implements LoginProtocol {
         List<ProtocolMapperProcessor<SAMLLoginResponseMapper>> loginResponseMappers = new LinkedList<>();
         ProtocolMapperProcessor<SAMLRoleListMapper> roleListMapper = null;
 
-        Set<ProtocolMapperModel> mappings = client.getProtocolMappers();
+        Set<ProtocolMapperModel> mappings = accessCode.getRequestedProtocolMappers();
         for (ProtocolMapperModel mapping : mappings) {
-            if (!mapping.getProtocol().equals(SamlProtocol.LOGIN_PROTOCOL)) continue;
 
             ProtocolMapper mapper = (ProtocolMapper)session.getKeycloakSessionFactory().getProviderFactory(ProtocolMapper.class, mapping.getProtocolMapper());
             if (mapper == null) continue;
diff --git a/services/src/main/java/org/keycloak/protocol/oidc/endpoints/UserInfoEndpoint.java b/services/src/main/java/org/keycloak/protocol/oidc/endpoints/UserInfoEndpoint.java
index f337163..7900d90 100755
--- a/services/src/main/java/org/keycloak/protocol/oidc/endpoints/UserInfoEndpoint.java
+++ b/services/src/main/java/org/keycloak/protocol/oidc/endpoints/UserInfoEndpoint.java
@@ -27,6 +27,7 @@ import org.keycloak.events.Details;
 import org.keycloak.events.EventBuilder;
 import org.keycloak.events.EventType;
 import org.keycloak.models.ClientModel;
+import org.keycloak.models.ClientSessionModel;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.UserModel;
@@ -122,10 +123,11 @@ public class UserInfoEndpoint {
         }
 
         UserSessionModel userSession = session.sessions().getUserSession(realm, token.getSessionState());
+        ClientSessionModel clientSession = session.sessions().getClientSession(token.getClientSession());
         ClientModel clientModel = realm.getClientByClientId(token.getIssuedFor());
         UserModel userModel = userSession.getUser();
         AccessToken userInfo = new AccessToken();
-        tokenManager.transformAccessToken(session, userInfo, realm, clientModel, userModel, userSession, null);
+        tokenManager.transformAccessToken(session, userInfo, realm, clientModel, userModel, userSession, clientSession);
 
         event
             .detail(Details.USERNAME, userModel.getUsername())
diff --git a/services/src/main/java/org/keycloak/protocol/oidc/TokenManager.java b/services/src/main/java/org/keycloak/protocol/oidc/TokenManager.java
index 8d9ab8c..f77fbdd 100755
--- a/services/src/main/java/org/keycloak/protocol/oidc/TokenManager.java
+++ b/services/src/main/java/org/keycloak/protocol/oidc/TokenManager.java
@@ -27,6 +27,7 @@ import org.keycloak.representations.AccessTokenResponse;
 import org.keycloak.representations.IDToken;
 import org.keycloak.representations.RefreshToken;
 import org.keycloak.services.managers.AuthenticationManager;
+import org.keycloak.services.managers.ClientSessionCode;
 import org.keycloak.util.Time;
 
 import javax.ws.rs.core.HttpHeaders;
@@ -208,6 +209,14 @@ public class TokenManager {
             requestedRoles.add(r.getId());
         }
         clientSession.setRoles(requestedRoles);
+
+        Set<String> requestedProtocolMappers = new HashSet<String>();
+        for (ProtocolMapperModel protocolMapper : clientSession.getClient().getProtocolMappers()) {
+            if (protocolMapper.getProtocol().equals(clientSession.getAuthMethod())) {
+                requestedProtocolMappers.add(protocolMapper.getId());
+            }
+        }
+        clientSession.setProtocolMappers(requestedProtocolMappers);
     }
 
     public static void dettachClientSession(UserSessionProvider sessions, RealmModel realm, ClientSessionModel clientSession) {
@@ -218,6 +227,7 @@ public class TokenManager {
 
         clientSession.setUserSession(null);
         clientSession.setRoles(null);
+        clientSession.setProtocolMappers(null);
 
         if (userSession.getClientSessions().isEmpty()) {
             sessions.removeUserSession(realm, userSession);
@@ -274,23 +284,20 @@ public class TokenManager {
 
     public AccessToken transformAccessToken(KeycloakSession session, AccessToken token, RealmModel realm, ClientModel client, UserModel user,
                                             UserSessionModel userSession, ClientSessionModel clientSession) {
-        Set<ProtocolMapperModel> mappings = client.getProtocolMappers();
+        Set<ProtocolMapperModel> mappings = new ClientSessionCode(realm, clientSession).getRequestedProtocolMappers();
         KeycloakSessionFactory sessionFactory = session.getKeycloakSessionFactory();
         for (ProtocolMapperModel mapping : mappings) {
-            if (!mapping.getProtocol().equals(OIDCLoginProtocol.LOGIN_PROTOCOL)) continue;
 
             ProtocolMapper mapper = (ProtocolMapper)sessionFactory.getProviderFactory(ProtocolMapper.class, mapping.getProtocolMapper());
             if (mapper == null || !(mapper instanceof OIDCAccessTokenMapper)) continue;
             token = ((OIDCAccessTokenMapper)mapper).transformAccessToken(token, mapping, session, userSession, clientSession);
 
-
-
         }
         return token;
     }
     public void transformIDToken(KeycloakSession session, IDToken token, RealmModel realm, ClientModel client, UserModel user,
                                       UserSessionModel userSession, ClientSessionModel clientSession) {
-        Set<ProtocolMapperModel> mappings = client.getProtocolMappers();
+        Set<ProtocolMapperModel> mappings = new ClientSessionCode(realm, clientSession).getRequestedProtocolMappers();
         KeycloakSessionFactory sessionFactory = session.getKeycloakSessionFactory();
         for (ProtocolMapperModel mapping : mappings) {
             if (!mapping.getProtocol().equals(OIDCLoginProtocol.LOGIN_PROTOCOL)) continue;
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 7647b36..8e363bb 100755
--- a/services/src/main/java/org/keycloak/services/managers/AuthenticationManager.java
+++ b/services/src/main/java/org/keycloak/services/managers/AuthenticationManager.java
@@ -439,8 +439,8 @@ public class AuthenticationManager {
             }
 
             List<ProtocolMapperModel> protocolMappers = new LinkedList<ProtocolMapperModel>();
-            for (ProtocolMapperModel protocolMapper : client.getProtocolMappers()) {
-                if (protocolMapper.isConsentRequired() && protocolMapper.getProtocol().equals(clientSession.getAuthMethod()) && protocolMapper.getConsentText() != null) {
+            for (ProtocolMapperModel protocolMapper : accessCode.getRequestedProtocolMappers()) {
+                if (protocolMapper.isConsentRequired() && protocolMapper.getConsentText() != null) {
                     if (grantedConsent == null || !grantedConsent.isProtocolMapperGranted(protocolMapper)) {
                         protocolMappers.add(protocolMapper);
                     }
diff --git a/services/src/main/java/org/keycloak/services/managers/ClientSessionCode.java b/services/src/main/java/org/keycloak/services/managers/ClientSessionCode.java
index 9e52c97..99fc5f2 100755
--- a/services/src/main/java/org/keycloak/services/managers/ClientSessionCode.java
+++ b/services/src/main/java/org/keycloak/services/managers/ClientSessionCode.java
@@ -2,6 +2,7 @@ package org.keycloak.services.managers;
 
 import org.keycloak.models.ClientSessionModel;
 import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.ProtocolMapperModel;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.RoleModel;
 import org.keycloak.models.UserModel.RequiredAction;
@@ -112,12 +113,25 @@ public class ClientSessionCode {
         for (String roleId : clientSession.getRoles()) {
             RoleModel role = realm.getRoleById(roleId);
             if (role != null) {
-                requestedRoles.add(realm.getRoleById(roleId));
+                requestedRoles.add(role);
             }
         }
         return requestedRoles;
     }
 
+    public Set<ProtocolMapperModel> getRequestedProtocolMappers() {
+        Set<ProtocolMapperModel> requestedProtocolMappers = new HashSet<ProtocolMapperModel>();
+        if (clientSession.getProtocolMappers() != null) {
+            for (String protocolMapperId : clientSession.getProtocolMappers()) {
+                ProtocolMapperModel protocolMapper = clientSession.getClient().getProtocolMapperById(protocolMapperId);
+                if (protocolMapper != null) {
+                    requestedProtocolMappers.add(protocolMapper);
+                }
+            }
+        }
+        return requestedProtocolMappers;
+    }
+
     public void setAction(ClientSessionModel.Action action) {
         clientSession.setAction(action);
         clientSession.setNote(ACTION_KEY, UUID.randomUUID().toString());
diff --git a/services/src/main/java/org/keycloak/services/resources/LoginActionsService.java b/services/src/main/java/org/keycloak/services/resources/LoginActionsService.java
index d4aa8d7..aeb3e2a 100755
--- a/services/src/main/java/org/keycloak/services/resources/LoginActionsService.java
+++ b/services/src/main/java/org/keycloak/services/resources/LoginActionsService.java
@@ -34,6 +34,7 @@ import org.keycloak.jose.jws.JWSBuilder;
 import org.keycloak.login.LoginFormsProvider;
 import org.keycloak.models.ClientModel;
 import org.keycloak.models.ClientSessionModel;
+import org.keycloak.models.RoleModel;
 import org.keycloak.models.UserConsentModel;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.ModelException;
@@ -611,16 +612,15 @@ public class LoginActionsService {
 
         UserConsentModel grantedConsent = user.getConsentByClient(client.getId());
         if (grantedConsent == null) {
-            grantedConsent = new UserConsentModel(realm, client.getId());
+            grantedConsent = new UserConsentModel(client);
             user.addConsent(grantedConsent);
         }
-        for (String roleId : clientSession.getRoles()) {
-            grantedConsent.addGrantedRole(roleId);
+        for (RoleModel role : accessCode.getRequestedRoles()) {
+            grantedConsent.addGrantedRole(role);
         }
-        // TODO: It's not 100% sure that approved protocolMappers are same like the protocolMappers retrieved here from the client. Maybe clientSession.setProtocolMappers/getProtocolMappers should be added...
-        for (ProtocolMapperModel protocolMapper : client.getProtocolMappers()) {
-            if (protocolMapper.isConsentRequired() && protocolMapper.getProtocol().equals(clientSession.getAuthMethod()) && protocolMapper.getConsentText() != null) {
-                grantedConsent.addGrantedProtocolMapper(protocolMapper.getId());
+        for (ProtocolMapperModel protocolMapper : accessCode.getRequestedProtocolMappers()) {
+            if (protocolMapper.isConsentRequired() && protocolMapper.getConsentText() != null) {
+                grantedConsent.addGrantedProtocolMapper(protocolMapper);
             }
         }
         user.updateConsent(grantedConsent);
diff --git a/testsuite/integration/src/test/java/org/keycloak/testsuite/model/UserConsentModelTest.java b/testsuite/integration/src/test/java/org/keycloak/testsuite/model/UserConsentModelTest.java
index 97cb24b..92ef5dd 100644
--- a/testsuite/integration/src/test/java/org/keycloak/testsuite/model/UserConsentModelTest.java
+++ b/testsuite/integration/src/test/java/org/keycloak/testsuite/model/UserConsentModelTest.java
@@ -46,15 +46,15 @@ public class UserConsentModelTest extends AbstractModelTest {
         UserModel john = session.users().addUser(realm, "john");
         UserModel mary = session.users().addUser(realm, "mary");
 
-        UserConsentModel johnFooGrant = new UserConsentModel(realm, fooClient.getId());
-        johnFooGrant.addGrantedRole(realmRole.getId());
-        johnFooGrant.addGrantedRole(barClientRole.getId());
-        johnFooGrant.addGrantedProtocolMapper(fooMapper.getId());
+        UserConsentModel johnFooGrant = new UserConsentModel(fooClient);
+        johnFooGrant.addGrantedRole(realmRole);
+        johnFooGrant.addGrantedRole(barClientRole);
+        johnFooGrant.addGrantedProtocolMapper(fooMapper);
         john.addConsent(johnFooGrant);
 
-        UserConsentModel johnBarGrant = new UserConsentModel(realm, barClient.getId());
-        johnBarGrant.addGrantedProtocolMapper(barMapper.getId());
-        johnBarGrant.addGrantedRole(realmRole.getId());
+        UserConsentModel johnBarGrant = new UserConsentModel(barClient);
+        johnBarGrant.addGrantedProtocolMapper(barMapper);
+        johnBarGrant.addGrantedRole(realmRole);
 
         // Update should fail as grant doesn't yet exists
         try {
@@ -65,9 +65,9 @@ public class UserConsentModelTest extends AbstractModelTest {
 
         john.addConsent(johnBarGrant);
 
-        UserConsentModel maryFooGrant = new UserConsentModel(realm, fooClient.getId());
-        maryFooGrant.addGrantedRole(realmRole.getId());
-        maryFooGrant.addGrantedProtocolMapper(fooMapper.getId());
+        UserConsentModel maryFooGrant = new UserConsentModel(fooClient);
+        maryFooGrant.addGrantedRole(realmRole);
+        maryFooGrant.addGrantedProtocolMapper(fooMapper);
         mary.addConsent(maryFooGrant);
 
         commit();
@@ -143,7 +143,7 @@ public class UserConsentModelTest extends AbstractModelTest {
         johnConsent.getGrantedRoles().remove(realmRole);
 
         RoleModel newRealmRole = realm.addRole("new-realm-role");
-        johnConsent.addGrantedRole(newRealmRole.getId());
+        johnConsent.addGrantedRole(newRealmRole);
 
         john.updateConsent(johnConsent);
 
diff --git a/testsuite/integration/src/test/java/org/keycloak/testsuite/model/UserSessionProviderTest.java b/testsuite/integration/src/test/java/org/keycloak/testsuite/model/UserSessionProviderTest.java
index b706203..c5b236a 100755
--- a/testsuite/integration/src/test/java/org/keycloak/testsuite/model/UserSessionProviderTest.java
+++ b/testsuite/integration/src/test/java/org/keycloak/testsuite/model/UserSessionProviderTest.java
@@ -101,6 +101,9 @@ public class UserSessionProviderTest {
         assertEquals(2, session1.getRoles().size());
         assertTrue(session1.getRoles().contains("one"));
         assertTrue(session1.getRoles().contains("two"));
+        assertEquals(2, session1.getProtocolMappers().size());
+        assertTrue(session1.getProtocolMappers().contains("mapper-one"));
+        assertTrue(session1.getProtocolMappers().contains("mapper-two"));
     }
 
     @Test
@@ -470,12 +473,13 @@ public class UserSessionProviderTest {
         assertNotNull(session.sessions().getUserLoginFailure(realm, "user2"));
     }
 
-    private ClientSessionModel createClientSession(ClientModel client, UserSessionModel userSession, String redirect, String state, Set<String> roles) {
+    private ClientSessionModel createClientSession(ClientModel client, UserSessionModel userSession, String redirect, String state, Set<String> roles, Set<String> protocolMappers) {
         ClientSessionModel clientSession = session.sessions().createClientSession(realm, client);
         if (userSession != null) clientSession.setUserSession(userSession);
         clientSession.setRedirectUri(redirect);
         if (state != null) clientSession.setNote(OIDCLoginProtocol.STATE_PARAM, state);
         if (roles != null) clientSession.setRoles(roles);
+        if (protocolMappers != null) clientSession.setProtocolMappers(protocolMappers);
         return clientSession;
     }
 
@@ -487,14 +491,18 @@ public class UserSessionProviderTest {
         roles.add("one");
         roles.add("two");
 
-        createClientSession(realm.getClientByClientId("test-app"), sessions[0], "http://redirect", "state", roles);
-        createClientSession(realm.getClientByClientId("third-party"), sessions[0], "http://redirect", "state", new HashSet<String>());
+        Set<String> protocolMappers = new HashSet<String>();
+        protocolMappers.add("mapper-one");
+        protocolMappers.add("mapper-two");
+
+        createClientSession(realm.getClientByClientId("test-app"), sessions[0], "http://redirect", "state", roles, protocolMappers);
+        createClientSession(realm.getClientByClientId("third-party"), sessions[0], "http://redirect", "state", new HashSet<String>(), new HashSet<String>());
 
         sessions[1] = session.sessions().createUserSession(realm, session.users().getUserByUsername("user1", realm), "user1", "127.0.0.2", "form", true, null, null);
-        createClientSession(realm.getClientByClientId("test-app"), sessions[1], "http://redirect", "state", new HashSet<String>());
+        createClientSession(realm.getClientByClientId("test-app"), sessions[1], "http://redirect", "state", new HashSet<String>(), new HashSet<String>());
 
         sessions[2] = session.sessions().createUserSession(realm, session.users().getUserByUsername("user2", realm), "user2", "127.0.0.3", "form", true, null, null);
-        createClientSession(realm.getClientByClientId("test-app"), sessions[2], "http://redirect", "state", new HashSet<String>());
+        createClientSession(realm.getClientByClientId("test-app"), sessions[2], "http://redirect", "state", new HashSet<String>(), new HashSet<String>());
 
         resetSession();