keycloak-aplcache

diff --git a/federation/ldap/src/main/java/org/keycloak/storage/ldap/idm/query/internal/LDAPQuery.java b/federation/ldap/src/main/java/org/keycloak/storage/ldap/idm/query/internal/LDAPQuery.java
index c27b2c6..a3e04c1 100644
--- a/federation/ldap/src/main/java/org/keycloak/storage/ldap/idm/query/internal/LDAPQuery.java
+++ b/federation/ldap/src/main/java/org/keycloak/storage/ldap/idm/query/internal/LDAPQuery.java
@@ -17,6 +17,7 @@
 
 package org.keycloak.storage.ldap.idm.query.internal;
 
+import org.jboss.logging.Logger;
 import org.keycloak.component.ComponentModel;
 import org.keycloak.models.ModelDuplicateException;
 import org.keycloak.models.ModelException;
@@ -26,7 +27,10 @@ import org.keycloak.storage.ldap.idm.query.Condition;
 import org.keycloak.storage.ldap.idm.query.Sort;
 import org.keycloak.storage.ldap.mappers.LDAPStorageMapper;
 
+import javax.naming.NamingException;
 import javax.naming.directory.SearchControls;
+import javax.naming.ldap.LdapContext;
+
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
@@ -39,16 +43,19 @@ import static java.util.Collections.unmodifiableSet;
 /**
  * Default IdentityQuery implementation.
  *
+ * LDAPQuery should be closed after use in case that pagination was used (initPagination was called)
  *
  * @author Shane Bryzak
  */
-public class LDAPQuery {
+public class LDAPQuery implements AutoCloseable{
+
+    private static final Logger logger = Logger.getLogger(LDAPQuery.class);
 
     private final LDAPStorageProvider ldapFedProvider;
 
     private int offset;
     private int limit;
-    private byte[] paginationContext;
+    private PaginationContext paginationContext;
     private String searchDn;
     private final Set<Condition> conditions = new LinkedHashSet<Condition>();
     private final Set<Sort> ordering = new LinkedHashSet<Sort>();
@@ -144,7 +151,7 @@ public class LDAPQuery {
         return offset;
     }
 
-    public byte[] getPaginationContext() {
+    public PaginationContext getPaginationContext() {
         return paginationContext;
     }
 
@@ -197,8 +204,8 @@ public class LDAPQuery {
         return this;
     }
 
-    public LDAPQuery setPaginationContext(byte[] paginationContext) {
-        this.paginationContext = paginationContext;
+    public LDAPQuery initPagination(LdapContext ldapContext) {
+        this.paginationContext = new PaginationContext(ldapContext);
         return this;
     }
 
@@ -210,4 +217,47 @@ public class LDAPQuery {
         return ldapFedProvider;
     }
 
+
+    @Override
+    public void close() {
+        if (paginationContext != null) {
+            try {
+                paginationContext.ldapContext.close();
+            } catch (NamingException ne) {
+                logger.error("Could not close Ldap context.", ne);
+            }
+        }
+    }
+
+
+    public static class PaginationContext {
+
+        private final LdapContext ldapContext;
+        private byte[] cookie;
+
+        private PaginationContext(LdapContext ldapContext) {
+            if (ldapContext == null) {
+                throw new IllegalArgumentException("Bad usage. Ldap context must be not null");
+            }
+            this.ldapContext = ldapContext;
+        }
+
+
+        public LdapContext getLdapContext() {
+            return ldapContext;
+        }
+
+        public byte[] getCookie() {
+            return cookie;
+        }
+
+        public void setCookie(byte[] cookie) {
+            this.cookie = cookie;
+        }
+
+        public boolean hasNextPage() {
+            return this.cookie != null;
+        }
+    }
+
 }

diff --git a/federation/ldap/src/main/java/org/keycloak/storage/ldap/idm/store/ldap/LDAPOperationManager.java b/federation/ldap/src/main/java/org/keycloak/storage/ldap/idm/store/ldap/LDAPOperationManager.java
index f4b8186..2edfd17 100644
--- a/federation/ldap/src/main/java/org/keycloak/storage/ldap/idm/store/ldap/LDAPOperationManager.java
+++ b/federation/ldap/src/main/java/org/keycloak/storage/ldap/idm/store/ldap/LDAPOperationManager.java
@@ -286,13 +286,19 @@ public class LDAPOperationManager {
         final List<SearchResult> result = new ArrayList<SearchResult>();
         final SearchControls cons = getSearchControls(identityQuery.getReturningLdapAttributes(), identityQuery.getSearchScope());
 
+        // Very 1st page. Pagination context is not yet present
+        if (identityQuery.getPaginationContext() == null) {
+            LdapContext ldapContext = createLdapContext();
+            identityQuery.initPagination(ldapContext);
+        }
+
         try {
             return execute(new LdapOperation<List<SearchResult>>() {
 
                 @Override
                 public List<SearchResult> execute(LdapContext context) throws NamingException {
                     try {
-                        byte[] cookie = identityQuery.getPaginationContext();
+                        byte[] cookie = identityQuery.getPaginationContext().getCookie();
                         PagedResultsControl pagedControls = new PagedResultsControl(identityQuery.getLimit(), cookie, Control.CRITICAL);
                         context.setRequestControls(new Control[] { pagedControls });
 
@@ -310,7 +316,7 @@ public class LDAPOperationManager {
                                 if (respControl instanceof PagedResultsResponseControl) {
                                     PagedResultsResponseControl prrc = (PagedResultsResponseControl)respControl;
                                     cookie = prrc.getCookie();
-                                    identityQuery.setPaginationContext(cookie);
+                                    identityQuery.getPaginationContext().setCookie(cookie);
                                 }
                             }
                         }
@@ -335,7 +341,7 @@ public class LDAPOperationManager {
                             .toString();
                 }
 
-            });
+            }, identityQuery.getPaginationContext().getLdapContext(), null);
         } catch (NamingException e) {
             logger.errorf(e, "Could not query server using DN [%s] and filter [%s]", baseDN, filter);
             throw e;
@@ -565,7 +571,7 @@ public class LDAPOperationManager {
                 }
 
 
-            }, decorator);
+            }, null, decorator);
         } catch (NamingException e) {
             throw new ModelException("Could not modify attribute for DN [" + dn + "]", e);
         }
@@ -726,11 +732,13 @@ public class LDAPOperationManager {
     }
 
     private <R> R execute(LdapOperation<R> operation) throws NamingException {
-        return execute(operation, null);
+        return execute(operation, null, null);
     }
 
-    private <R> R execute(LdapOperation<R> operation, LDAPOperationDecorator decorator) throws NamingException {
-        LdapContext context = null;
+    private <R> R execute(LdapOperation<R> operation, LdapContext context, LDAPOperationDecorator decorator) throws NamingException {
+        // We won't manage LDAP context (create and close) in case that existing context was passed as an argument to this method
+        boolean manageContext = context == null;
+
         Long start = null;
 
         try {
@@ -738,14 +746,17 @@ public class LDAPOperationManager {
                 start = Time.currentTimeMillis();
             }
 
-            context = createLdapContext();
+            if (manageContext) {
+                context = createLdapContext();
+            }
+
             if (decorator != null) {
                 decorator.beforeLDAPOperation(context, operation);
             }
 
             return operation.execute(context);
         } finally {
-            if (context != null) {
+            if (context != null && manageContext) {
                 try {
                     context.close();
                 } catch (NamingException ne) {

diff --git a/federation/ldap/src/main/java/org/keycloak/storage/ldap/LDAPStorageProviderFactory.java b/federation/ldap/src/main/java/org/keycloak/storage/ldap/LDAPStorageProviderFactory.java
index b3ed27a..92f3c06 100755
--- a/federation/ldap/src/main/java/org/keycloak/storage/ldap/LDAPStorageProviderFactory.java
+++ b/federation/ldap/src/main/java/org/keycloak/storage/ldap/LDAPStorageProviderFactory.java
@@ -424,13 +424,14 @@ public class LDAPStorageProviderFactory implements UserStorageProviderFactory<LD
 
         logger.infof("Sync all users from LDAP to local store: realm: %s, federation provider: %s", realmId, model.getName());
 
-        LDAPQuery userQuery = createQuery(sessionFactory, realmId, model);
-        SynchronizationResult syncResult = syncImpl(sessionFactory, userQuery, realmId, model);
+        try (LDAPQuery userQuery = createQuery(sessionFactory, realmId, model)) {
+            SynchronizationResult syncResult = syncImpl(sessionFactory, userQuery, realmId, model);
 
-        // TODO: Remove all existing keycloak users, which have federation links, but are not in LDAP. Perhaps don't check users, which were just added or updated during this sync?
+            // TODO: Remove all existing keycloak users, which have federation links, but are not in LDAP. Perhaps don't check users, which were just added or updated during this sync?
 
-        logger.infof("Sync all users finished: %s", syncResult.getStatus());
-        return syncResult;
+            logger.infof("Sync all users finished: %s", syncResult.getStatus());
+            return syncResult;
+        }
     }
 
     @Override
@@ -445,12 +446,13 @@ public class LDAPStorageProviderFactory implements UserStorageProviderFactory<LD
         Condition modifyCondition = conditionsBuilder.greaterThanOrEqualTo(LDAPConstants.MODIFY_TIMESTAMP, lastSync);
         Condition orCondition = conditionsBuilder.orCondition(createCondition, modifyCondition);
 
-        LDAPQuery userQuery = createQuery(sessionFactory, realmId, model);
-        userQuery.addWhereCondition(orCondition);
-        SynchronizationResult result = syncImpl(sessionFactory, userQuery, realmId, model);
+        try (LDAPQuery userQuery = createQuery(sessionFactory, realmId, model)) {
+            userQuery.addWhereCondition(orCondition);
+            SynchronizationResult result = syncImpl(sessionFactory, userQuery, realmId, model);
 
-        logger.infof("Sync changed users finished: %s", result.getStatus());
-        return result;
+            logger.infof("Sync changed users finished: %s", result.getStatus());
+            return result;
+        }
     }
 
     protected void syncMappers(KeycloakSessionFactory sessionFactory, final String realmId, final ComponentModel model) {
@@ -486,7 +488,7 @@ public class LDAPStorageProviderFactory implements UserStorageProviderFactory<LD
             while (nextPage) {
                 userQuery.setLimit(pageSize);
                 final List<LDAPObject> users = userQuery.getResultList();
-                nextPage = userQuery.getPaginationContext() != null;
+                nextPage = userQuery.getPaginationContext().hasNextPage();
                 SynchronizationResult currentPageSync = importLdapUsers(sessionFactory, realmId, fedModel, users);
                 syncResult.add(currentPageSync);
             }

diff --git a/federation/ldap/src/main/java/org/keycloak/storage/ldap/LDAPUtils.java b/federation/ldap/src/main/java/org/keycloak/storage/ldap/LDAPUtils.java
index 4f05ad4..03107f8 100755
--- a/federation/ldap/src/main/java/org/keycloak/storage/ldap/LDAPUtils.java
+++ b/federation/ldap/src/main/java/org/keycloak/storage/ldap/LDAPUtils.java
@@ -250,7 +250,7 @@ public class LDAPUtils {
      * Load all LDAP objects corresponding to given query. We will load them paginated, so we allow to bypass the limitation of 1000
      * maximum loaded objects in single query in MSAD
      *
-     * @param ldapQuery
+     * @param ldapQuery LDAP query to be used. The caller should close it after calling this method
      * @param ldapProvider
      * @return
      */
@@ -268,7 +268,7 @@ public class LDAPUtils {
                 ldapQuery.setLimit(pageSize);
                 final List<LDAPObject> currentPageGroups = ldapQuery.getResultList();
                 result.addAll(currentPageGroups);
-                nextPage = ldapQuery.getPaginationContext() != null;
+                nextPage = ldapQuery.getPaginationContext().hasNextPage();
             }
 
             return result;

diff --git a/federation/ldap/src/main/java/org/keycloak/storage/ldap/mappers/membership/group/GroupLDAPStorageMapper.java b/federation/ldap/src/main/java/org/keycloak/storage/ldap/mappers/membership/group/GroupLDAPStorageMapper.java
index 66ec3c3..d446fbd 100644
--- a/federation/ldap/src/main/java/org/keycloak/storage/ldap/mappers/membership/group/GroupLDAPStorageMapper.java
+++ b/federation/ldap/src/main/java/org/keycloak/storage/ldap/mappers/membership/group/GroupLDAPStorageMapper.java
@@ -350,8 +350,9 @@ public class GroupLDAPStorageMapper extends AbstractLDAPStorageMapper implements
 
     // Send LDAP query to retrieve all groups
     protected List<LDAPObject> getAllLDAPGroups(boolean includeMemberAttribute) {
-        LDAPQuery ldapGroupQuery = createGroupQuery(includeMemberAttribute);
-        return LDAPUtils.loadAllLDAPObjects(ldapGroupQuery, ldapProvider);
+        try (LDAPQuery ldapGroupQuery = createGroupQuery(includeMemberAttribute)) {
+            return LDAPUtils.loadAllLDAPObjects(ldapGroupQuery, ldapProvider);
+        }
     }
 
 

diff --git a/federation/ldap/src/main/java/org/keycloak/storage/ldap/mappers/membership/role/RoleLDAPStorageMapper.java b/federation/ldap/src/main/java/org/keycloak/storage/ldap/mappers/membership/role/RoleLDAPStorageMapper.java
index cee571f..ee77715 100644
--- a/federation/ldap/src/main/java/org/keycloak/storage/ldap/mappers/membership/role/RoleLDAPStorageMapper.java
+++ b/federation/ldap/src/main/java/org/keycloak/storage/ldap/mappers/membership/role/RoleLDAPStorageMapper.java
@@ -124,24 +124,25 @@ public class RoleLDAPStorageMapper extends AbstractLDAPStorageMapper implements 
         logger.debugf("Syncing roles from LDAP into Keycloak DB. Mapper is [%s], LDAP provider is [%s]", mapperModel.getName(), ldapProvider.getModel().getName());
 
         // Send LDAP query to load all roles
-        LDAPQuery ldapRoleQuery = createRoleQuery(false);
-        List<LDAPObject> ldapRoles = LDAPUtils.loadAllLDAPObjects(ldapRoleQuery, ldapProvider);
+        try (LDAPQuery ldapRoleQuery = createRoleQuery(false)) {
+            List<LDAPObject> ldapRoles = LDAPUtils.loadAllLDAPObjects(ldapRoleQuery, ldapProvider);
 
-        RoleContainerModel roleContainer = getTargetRoleContainer(realm);
-        String rolesRdnAttr = config.getRoleNameLdapAttribute();
-        for (LDAPObject ldapRole : ldapRoles) {
-            String roleName = ldapRole.getAttributeAsString(rolesRdnAttr);
+            RoleContainerModel roleContainer = getTargetRoleContainer(realm);
+            String rolesRdnAttr = config.getRoleNameLdapAttribute();
+            for (LDAPObject ldapRole : ldapRoles) {
+                String roleName = ldapRole.getAttributeAsString(rolesRdnAttr);
 
-            if (roleContainer.getRole(roleName) == null) {
-                logger.debugf("Syncing role [%s] from LDAP to keycloak DB", roleName);
-                roleContainer.addRole(roleName);
-                syncResult.increaseAdded();
-            } else {
-                syncResult.increaseUpdated();
+                if (roleContainer.getRole(roleName) == null) {
+                    logger.debugf("Syncing role [%s] from LDAP to keycloak DB", roleName);
+                    roleContainer.addRole(roleName);
+                    syncResult.increaseAdded();
+                } else {
+                    syncResult.increaseUpdated();
+                }
             }
-        }
 
-        return syncResult;
+            return syncResult;
+        }
     }
 
 
@@ -165,32 +166,33 @@ public class RoleLDAPStorageMapper extends AbstractLDAPStorageMapper implements 
         logger.debugf("Syncing roles from Keycloak into LDAP. Mapper is [%s], LDAP provider is [%s]", mapperModel.getName(), ldapProvider.getModel().getName());
 
         // Send LDAP query to see which roles exists there
-        LDAPQuery ldapQuery = createRoleQuery(false);
-        List<LDAPObject> ldapRoles = LDAPUtils.loadAllLDAPObjects(ldapQuery, ldapProvider);
+        try (LDAPQuery ldapQuery = createRoleQuery(false)) {
+            List<LDAPObject> ldapRoles = LDAPUtils.loadAllLDAPObjects(ldapQuery, ldapProvider);
 
-        Set<String> ldapRoleNames = new HashSet<>();
-        String rolesRdnAttr = config.getRoleNameLdapAttribute();
-        for (LDAPObject ldapRole : ldapRoles) {
-            String roleName = ldapRole.getAttributeAsString(rolesRdnAttr);
-            ldapRoleNames.add(roleName);
-        }
+            Set<String> ldapRoleNames = new HashSet<>();
+            String rolesRdnAttr = config.getRoleNameLdapAttribute();
+            for (LDAPObject ldapRole : ldapRoles) {
+                String roleName = ldapRole.getAttributeAsString(rolesRdnAttr);
+                ldapRoleNames.add(roleName);
+            }
 
 
-        RoleContainerModel roleContainer = getTargetRoleContainer(realm);
-        Set<RoleModel> keycloakRoles = roleContainer.getRoles();
+            RoleContainerModel roleContainer = getTargetRoleContainer(realm);
+            Set<RoleModel> keycloakRoles = roleContainer.getRoles();
 
-        for (RoleModel keycloakRole : keycloakRoles) {
-            String roleName = keycloakRole.getName();
-            if (ldapRoleNames.contains(roleName)) {
-                syncResult.increaseUpdated();
-            } else {
-                logger.debugf("Syncing role [%s] from Keycloak to LDAP", roleName);
-                createLDAPRole(roleName);
-                syncResult.increaseAdded();
+            for (RoleModel keycloakRole : keycloakRoles) {
+                String roleName = keycloakRole.getName();
+                if (ldapRoleNames.contains(roleName)) {
+                    syncResult.increaseUpdated();
+                } else {
+                    logger.debugf("Syncing role [%s] from Keycloak to LDAP", roleName);
+                    createLDAPRole(roleName);
+                    syncResult.increaseAdded();
+                }
             }
-        }
 
-        return syncResult;
+            return syncResult;
+        }
     }
 
     // TODO: Possible to merge with GroupMapper and move to common class