[KEYCLOAK-4902] - Using streams when fetching resources

authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/permission/AbstractPermissionProvider.java

@@ -44,9 +44,10 @@ public abstract class AbstractPermissionProvider implements PolicyProvider {
             Map<Object, Decision.Effect> decisions = decisionCache.computeIfAbsent(associatedPolicy, p -> new HashMap<>());
             Decision.Effect effect = decisions.get(permission);
 
+            defaultEvaluation.setPolicy(associatedPolicy);
+
             if (effect == null) {
                 PolicyProvider policyProvider = authorization.getProvider(associatedPolicy.getType());
-                defaultEvaluation.setPolicy(associatedPolicy);
                 policyProvider.evaluate(defaultEvaluation);
                 evaluation.denyIfNoEffect();
                 decisions.put(permission, defaultEvaluation.getEffect());

authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/permission/ScopePolicyProvider.java

@@ -44,6 +44,7 @@ public class ScopePolicyProvider extends AbstractPermissionProvider {
 
             if (effect != null) {
                 defaultEvaluation.setEffect(effect);
+                return;
             }
         }
 

model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/authorization/StoreFactoryCacheSession.java

@@ -628,6 +628,13 @@ public class StoreFactoryCacheSession implements CachedStoreFactoryProvider {
                     (revision, resources) -> new ResourceListQuery(revision, cacheKey, resources.stream().map(resource -> resource.getId()).collect(Collectors.toSet()), resourceServerId), resourceServerId);
         }
 
+        @Override
+        public void findByOwner(String ownerId, String resourceServerId, Consumer<Resource> consumer) {
+            String cacheKey = getResourceByOwnerCacheKey(ownerId, resourceServerId);
+            cacheQuery(cacheKey, ResourceListQuery.class, () -> getResourceStoreDelegate().findByOwner(ownerId, resourceServerId),
+                    (revision, resources) -> new ResourceListQuery(revision, cacheKey, resources.stream().map(resource -> resource.getId()).collect(Collectors.toSet()), resourceServerId), resourceServerId, consumer);
+        }
+
         @Override
         public List<Resource> findByUri(String uri, String resourceServerId) {
             if (uri == null) return null;
@@ -667,7 +674,19 @@ public class StoreFactoryCacheSession implements CachedStoreFactoryProvider {
                      (revision, resources) -> new ResourceListQuery(revision, cacheKey, resources.stream().map(resource -> resource.getId()).collect(Collectors.toSet()), resourceServerId), resourceServerId);
         }
 
+        @Override
+        public void findByType(String type, String resourceServerId, Consumer<Resource> consumer) {
+            if (type == null) return;
+            String cacheKey = getResourceByTypeCacheKey(type, resourceServerId);
+            cacheQuery(cacheKey, ResourceListQuery.class, () -> getResourceStoreDelegate().findByType(type, resourceServerId),
+                    (revision, resources) -> new ResourceListQuery(revision, cacheKey, resources.stream().map(resource -> resource.getId()).collect(Collectors.toSet()), resourceServerId), resourceServerId, consumer);
+        }
+
         private <R, Q extends ResourceQuery> List<R> cacheQuery(String cacheKey, Class<Q> queryType, Supplier<List<R>> resultSupplier, BiFunction<Long, List<R>, Q> querySupplier, String resourceServerId) {
+            return cacheQuery(cacheKey, queryType, resultSupplier, querySupplier, resourceServerId, null);
+        }
+
+        private <R, Q extends ResourceQuery> List<R> cacheQuery(String cacheKey, Class<Q> queryType, Supplier<List<R>> resultSupplier, BiFunction<Long, List<R>, Q> querySupplier, String resourceServerId, Consumer<R> consumer) {
             Q query = cache.get(cacheKey, queryType);
             if (query != null) {
                 logger.tracev("cache hit for key: {0}", cacheKey);
@@ -679,11 +698,31 @@ public class StoreFactoryCacheSession implements CachedStoreFactoryProvider {
                 if (invalidations.contains(cacheKey)) return model;
                 query = querySupplier.apply(loaded, model);
                 cache.addRevisioned(query, startupRevision);
+                if (consumer != null) {
+                    for (R resource : model) {
+                        consumer.accept(resource);
+                    }
+                }
                 return model;
             } else if (query.isInvalid(invalidations)) {
-                return resultSupplier.get();
+                List<R> result = resultSupplier.get();
+
+                if (consumer != null) {
+                    for (R resource : result) {
+                        consumer.accept(resource);
+                    }
+                }
+
+                return result;
             } else {
-                return query.getResources().stream().map(resourceId -> (R) findById(resourceId, resourceServerId)).collect(Collectors.toList());
+                Set<String> resources = query.getResources();
+
+                if (consumer != null) {
+                    resources.stream().map(resourceId -> (R) findById(resourceId, resourceServerId)).forEach(consumer);
+                    return Collections.emptyList();
+                }
+
+                return resources.stream().map(resourceId -> (R) findById(resourceId, resourceServerId)).collect(Collectors.toList());
             }
         }
     }

model/jpa/src/main/java/org/keycloak/authorization/jpa/store/JPAResourceStore.java

@@ -38,6 +38,8 @@ import java.util.ArrayList;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+import java.util.function.Consumer;
 
 /**
  * @author <a href="mailto:psilva@redhat.com">Pedro Igor</a>
@@ -99,6 +101,15 @@ public class JPAResourceStore implements ResourceStore {
 
     @Override
     public List<Resource> findByOwner(String ownerId, String resourceServerId) {
+        List<Resource> list = new LinkedList<>();
+
+        findByOwner(ownerId, resourceServerId, list::add);
+
+        return list;
+    }
+
+    @Override
+    public void findByOwner(String ownerId, String resourceServerId, Consumer<Resource> consumer) {
         String queryName = "findResourceIdByOwner";
 
         if (resourceServerId == null) {
@@ -114,19 +125,12 @@ public class JPAResourceStore implements ResourceStore {
             query.setParameter("serverId", resourceServerId);
         }
 
-        List<String> result = query.getResultList();
-        List<Resource> list = new LinkedList<>();
         ResourceStore resourceStore = provider.getStoreFactory().getResourceStore();
 
-        for (String id : result) {
+        query.getResultList().stream()
-            Resource resource = resourceStore.findById(id, resourceServerId);
+                .map(id -> resourceStore.findById(id, resourceServerId))
-
+                .filter(Objects::nonNull)
-            if (resource != null) {
+                .forEach(consumer);
-                list.add(resource);
-            }
-        }
-
-        return list;
     }
 
     @Override
@@ -279,6 +283,15 @@ public class JPAResourceStore implements ResourceStore {
 
     @Override
     public List<Resource> findByType(String type, String resourceServerId) {
+        List<Resource> list = new LinkedList<>();
+
+        findByType(type, resourceServerId, list::add);
+
+        return list;
+    }
+
+    @Override
+    public void findByType(String type, String resourceServerId, Consumer<Resource> consumer) {
         TypedQuery<String> query = entityManager.createNamedQuery("findResourceIdByType", String.class);
 
         query.setFlushMode(FlushModeType.COMMIT);
@@ -286,18 +299,11 @@ public class JPAResourceStore implements ResourceStore {
         query.setParameter("ownerId", resourceServerId);
         query.setParameter("serverId", resourceServerId);
 
-        List<String> result = query.getResultList();
-        List<Resource> list = new LinkedList<>();
         ResourceStore resourceStore = provider.getStoreFactory().getResourceStore();
 
-        for (String id : result) {
+        query.getResultList().stream()
-            Resource resource = resourceStore.findById(id, resourceServerId);
+                .map(id -> resourceStore.findById(id, resourceServerId))
-
+                .filter(Objects::nonNull)
-            if (resource != null) {
+                .forEach(consumer);
-                list.add(resource);
-            }
-        }
-
-        return list;
     }
 }

server-spi-private/src/main/java/org/keycloak/authorization/AuthorizationProvider.java

@@ -476,6 +476,11 @@ public final class AuthorizationProvider implements Provider {
                 return delegate.findByOwner(ownerId, resourceServerId);
             }
 
+            @Override
+            public void findByOwner(String ownerId, String resourceServerId, Consumer<Resource> consumer) {
+                delegate.findByOwner(ownerId, resourceServerId, consumer);
+            }
+
             @Override
             public List<Resource> findByUri(String uri, String resourceServerId) {
                 return delegate.findByUri(uri, resourceServerId);
@@ -510,6 +515,11 @@ public final class AuthorizationProvider implements Provider {
             public List<Resource> findByType(String type, String resourceServerId) {
                 return delegate.findByType(type, resourceServerId);
             }
+
+            @Override
+            public void findByType(String type, String resourceServerId, Consumer<Resource> consumer) {
+                delegate.findByType(type, resourceServerId, consumer);
+            }
         };
     }
 }

server-spi-private/src/main/java/org/keycloak/authorization/policy/evaluation/DefaultPolicyEvaluator.java

@@ -18,10 +18,9 @@
 
 package org.keycloak.authorization.policy.evaluation;
 
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
@@ -61,7 +60,7 @@ public class DefaultPolicyEvaluator implements PolicyEvaluator {
             return;
         }
 
-        Set<Policy> verified = new HashSet<>();
+        AtomicBoolean verified = new AtomicBoolean();
         Consumer<Policy> policyConsumer = createPolicyEvaluator(permission, authorizationProvider, executionContext, decision, verified, decisionCache);
         Resource resource = permission.getResource();
 
@@ -85,7 +84,7 @@ public class DefaultPolicyEvaluator implements PolicyEvaluator {
             policyStore.findByScopeIds(scopes.stream().map(Scope::getId).collect(Collectors.toList()), null, resourceServer.getId(), policyConsumer);
         }
 
-        if (!verified.isEmpty()) {
+        if (verified.get()) {
             decision.onComplete(permission);
             return;
         }
@@ -97,12 +96,8 @@ public class DefaultPolicyEvaluator implements PolicyEvaluator {
         }
     }
 
-    private Consumer<Policy> createPolicyEvaluator(ResourcePermission permission, AuthorizationProvider authorizationProvider, EvaluationContext executionContext, Decision decision, Set<Policy> verified, Map<Policy, Map<Object, Decision.Effect>> decisionCache) {
+    private Consumer<Policy> createPolicyEvaluator(ResourcePermission permission, AuthorizationProvider authorizationProvider, EvaluationContext executionContext, Decision decision, AtomicBoolean verified, Map<Policy, Map<Object, Decision.Effect>> decisionCache) {
         return parentPolicy -> {
-            if (!verified.add(parentPolicy)) {
-                return;
-            }
-
             PolicyProvider policyProvider = authorizationProvider.getProvider(parentPolicy.getType());
 
             if (policyProvider == null) {
@@ -110,6 +105,8 @@ public class DefaultPolicyEvaluator implements PolicyEvaluator {
             }
 
             policyProvider.evaluate(new DefaultEvaluation(permission, executionContext, parentPolicy, decision, authorizationProvider, decisionCache));
+
+            verified.compareAndSet(false, true);
         };
     }
 }

server-spi-private/src/main/java/org/keycloak/authorization/store/ResourceStore.java

@@ -22,6 +22,7 @@ import org.keycloak.authorization.model.ResourceServer;
 
 import java.util.List;
 import java.util.Map;
+import java.util.function.Consumer;
 
 /**
  * A {@link ResourceStore} is responsible to manage the persistence of {@link Resource} instances.
@@ -74,6 +75,8 @@ public interface ResourceStore {
      */
     List<Resource> findByOwner(String ownerId, String resourceServerId);
 
+    void findByOwner(String ownerId, String resourceServerId, Consumer<Resource> consumer);
+
     /**
      * Finds all {@link Resource} instances with the given uri.
      *
@@ -133,4 +136,14 @@ public interface ResourceStore {
      * @return a list of resources with the given type
      */
     List<Resource> findByType(String type, String resourceServerId);
+
+    /**
+     * Finds all {@link Resource} with the given type.
+     *
+     * @param type the type of the resource
+     * @param resourceServerId the resource server id
+     * @param consumer the result consumer
+     * @return a list of resources with the given type
+     */
+    void findByType(String type, String resourceServerId, Consumer<Resource> consumer);
 }

services/src/main/java/org/keycloak/authorization/util/Permissions.java

@@ -25,6 +25,7 @@ import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.Collectors;
 
 import javax.ws.rs.core.Response.Status;
@@ -68,17 +69,27 @@ public final class Permissions {
         StoreFactory storeFactory = authorization.getStoreFactory();
         ResourceStore resourceStore = storeFactory.getResourceStore();
         Metadata metadata = request.getMetadata();
-        long limit = Long.MAX_VALUE;
+        final AtomicLong limit;
 
         if (metadata != null && metadata.getLimit() != null) {
-            limit = metadata.getLimit();
+            limit = new AtomicLong(metadata.getLimit());
+        } else {
+            limit = new AtomicLong(Long.MAX_VALUE);
         }
 
         // obtain all resources where owner is the resource server
-        resourceStore.findByOwner(resourceServer.getId(), resourceServer.getId()).stream().limit(limit).forEach(resource -> permissions.add(createResourcePermissionsWithScopes(resource, new LinkedList(resource.getScopes()), authorization, request)));
+        resourceStore.findByOwner(resourceServer.getId(), resourceServer.getId(), resource -> {
+            if (limit.decrementAndGet() >= 0) {
+                permissions.add(createResourcePermissionsWithScopes(resource, new LinkedList(resource.getScopes()), authorization, request));
+            }
+        });
 
         // obtain all resources where owner is the current user
-        resourceStore.findByOwner(identity.getId(), resourceServer.getId()).stream().limit(limit).forEach(resource -> permissions.add(createResourcePermissionsWithScopes(resource, new LinkedList(resource.getScopes()), authorization, request)));
+        resourceStore.findByOwner(identity.getId(), resourceServer.getId(), resource -> {
+            if (limit.decrementAndGet() >= 0) {
+                permissions.add(createResourcePermissionsWithScopes(resource, new LinkedList(resource.getScopes()), authorization, request));
+            }
+        });
 
         // obtain all resources granted to the user via permission tickets (uma)
         List<PermissionTicket> tickets = storeFactory.getPermissionTicketStore().findGranted(identity.getId(), resourceServer.getId());
@@ -91,10 +102,10 @@ public final class Permissions {
 
                 if (permission == null) {
                     userManagedPermissions.put(ticket.getResource().getId(), new ResourcePermission(ticket.getResource(), new ArrayList<>(), resourceServer, request.getClaims()));
-                    limit--;
+                    limit.decrementAndGet();
                 }
 
-                if (--limit <= 0) {
+                if (limit.decrementAndGet() <= 0) {
                     break;
                 }
             }
@@ -117,7 +128,7 @@ public final class Permissions {
             if (type != null && !resource.getOwner().equals(resourceServer.getId())) {
                 StoreFactory storeFactory = authorization.getStoreFactory();
                 ResourceStore resourceStore = storeFactory.getResourceStore();
-                resourceStore.findByType(type, resourceServer.getId()).forEach(resource1 -> {
+                resourceStore.findByType(type, resourceServer.getId(), resource1 -> {
                     if (resource1.getOwner().equals(resourceServer.getId())) {
                         for (Scope typeScope : resource1.getScopes()) {
                             if (!scopes.contains(typeScope)) {
@@ -137,7 +148,7 @@ public final class Permissions {
                 }
 
                 return byName;
-            }).collect(Collectors.toList());
+            }).filter(resource.getScopes()::contains).collect(Collectors.toList());
         }
 
         return new ResourcePermission(resource, scopes, resource.getResourceServer(), request.getClaims());

testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/authz/EntitlementAPITest.java

@@ -639,7 +639,7 @@ public class EntitlementAPITest extends AbstractAuthzTest {
         resource = new ResourceRepresentation();
 
         resource.setName(KeycloakModelUtils.generateId());
-        resource.addScope("sensors:view", "sensors:update", "sensors:delete");
+        resource.addScope("sensors:view", "sensors:update");
 
         resourceIds.add(authorization.resources().create(resource).readEntity(ResourceRepresentation.class).getId());
 
@@ -695,6 +695,81 @@ public class EntitlementAPITest extends AbstractAuthzTest {
         }
     }
 
+    @Test
+    public void testObtainAllEntitlementsForResource() throws Exception {
+        ClientResource client = getClient(getRealm(), RESOURCE_SERVER_TEST);
+        AuthorizationResource authorization = client.authorization();
+
+        JSPolicyRepresentation policy = new JSPolicyRepresentation();
+
+        policy.setName(KeycloakModelUtils.generateId());
+        policy.setCode("$evaluation.grant();");
+
+        authorization.policies().js().create(policy).close();
+
+        ResourceRepresentation resource = new ResourceRepresentation();
+
+        resource.setName(KeycloakModelUtils.generateId());
+        resource.addScope("scope:view", "scope:update", "scope:delete");
+
+        resource = authorization.resources().create(resource).readEntity(ResourceRepresentation.class);
+
+        ResourcePermissionRepresentation permission = new ResourcePermissionRepresentation();
+
+        permission.setName(KeycloakModelUtils.generateId());
+        permission.addResource(resource.getId());
+        permission.addPolicy(policy.getName());
+
+        authorization.permissions().resource().create(permission);
+
+        String accessToken = new OAuthClient().realm("authz-test").clientId(RESOURCE_SERVER_TEST).doGrantAccessTokenRequest("secret", "kolo", "password").getAccessToken();
+        AuthzClient authzClient = getAuthzClient(AUTHZ_CLIENT_CONFIG);
+        AuthorizationRequest request = new AuthorizationRequest();
+
+        request.addPermission(null, "scope:view", "scope:update", "scope:delete");
+
+        AuthorizationResponse response = authzClient.authorization(accessToken).authorize(request);
+        assertNotNull(response.getToken());
+        Collection<Permission> permissions = toAccessToken(response.getToken()).getAuthorization().getPermissions();
+        assertEquals(1, permissions.size());
+
+        for (Permission grantedPermission : permissions) {
+            assertEquals(resource.getId(), grantedPermission.getResourceId());
+            assertEquals(3, grantedPermission.getScopes().size());
+            assertTrue(grantedPermission.getScopes().containsAll(Arrays.asList("scope:view")));
+        }
+
+        resource.setScopes(new HashSet<>());
+        resource.addScope("scope:view", "scope:update");
+
+        authorization.resources().resource(resource.getId()).update(resource);
+
+        request = new AuthorizationRequest();
+
+        request.addPermission(null, "scope:view", "scope:update", "scope:delete");
+
+        response = authzClient.authorization(accessToken).authorize(request);
+        assertNotNull(response.getToken());
+        permissions = toAccessToken(response.getToken()).getAuthorization().getPermissions();
+        assertEquals(1, permissions.size());
+
+        for (Permission grantedPermission : permissions) {
+            assertEquals(resource.getId(), grantedPermission.getResourceId());
+            assertEquals(2, grantedPermission.getScopes().size());
+            assertTrue(grantedPermission.getScopes().containsAll(Arrays.asList("scope:view", "scope:update")));
+        }
+
+        request = new AuthorizationRequest();
+
+        request.addPermission(resource.getId(), "scope:view", "scope:update", "scope:delete");
+
+        for (Permission grantedPermission : permissions) {
+            assertEquals(resource.getId(), grantedPermission.getResourceId());
+            assertEquals(2, grantedPermission.getScopes().size());
+            assertTrue(grantedPermission.getScopes().containsAll(Arrays.asList("scope:view", "scope:update")));
+        }
+    }
+
     private void testRptRequestWithResourceName(String configFile) {
         Metadata metadata = new Metadata();