From bf19ec11cfd7189141a0ed67f57b464b55fd54d1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt?= <49686195+newwdles@users.noreply.github.com>
Date: Tue, 24 Sep 2024 09:51:35 +0200
Subject: [PATCH] Fix UserStorageManager.getGroupMembersStream potentially
 fetching all user (#33145)

Closes #32761
Signed-off-by: Benoit Messager <benoit.messager@liksi.fr>
Co-authored-by: Benoit Messager <benoit.messager@liksi.fr>
---
 .../keycloak/storage/UserStorageManager.java  | 23 ++++++++++++-------
 .../federation/UserPropertyFileStorage.java   |  6 +++++
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/model/storage-private/src/main/java/org/keycloak/storage/UserStorageManager.java b/model/storage-private/src/main/java/org/keycloak/storage/UserStorageManager.java
index 8c546afdee..b86a44b131 100755
--- a/model/storage-private/src/main/java/org/keycloak/storage/UserStorageManager.java
+++ b/model/storage-private/src/main/java/org/keycloak/storage/UserStorageManager.java
@@ -431,15 +431,22 @@ public class UserStorageManager extends AbstractStorageManager<UserStorageProvid
     @Override
     public Stream<UserModel> getGroupMembersStream(final RealmModel realm, final GroupModel group, Integer firstResult, Integer maxResults) {
         Stream<UserModel> results = query((provider, firstResultInQuery, maxResultsInQuery) -> {
-            if (provider instanceof UserQueryMethodsProvider) {
-                return ((UserQueryMethodsProvider)provider).getGroupMembersStream(realm, group, firstResultInQuery, maxResultsInQuery);
+                    if (provider instanceof UserQueryMethodsProvider) {
+                        return ((UserQueryMethodsProvider) provider).getGroupMembersStream(realm, group, firstResultInQuery, maxResultsInQuery);
 
-            } else if (provider instanceof UserFederatedStorageProvider) {
-                return ((UserFederatedStorageProvider)provider).getMembershipStream(realm, group, firstResultInQuery, maxResultsInQuery).
-                        map(id -> getUserById(realm, id));
-           }
-            return Stream.empty();
-        }, realm, firstResult, maxResults);
+                    } else if (provider instanceof UserFederatedStorageProvider) {
+                        return ((UserFederatedStorageProvider) provider).getMembershipStream(realm, group, firstResultInQuery, maxResultsInQuery).
+                                map(id -> getUserById(realm, id));
+                    }
+                    return Stream.empty();
+                },
+                (provider, firstResultInQuery, maxResultsInQuery) -> {
+                    if (provider instanceof UserCountMethodsProvider) {
+                        return ((UserCountMethodsProvider) provider).getUsersCount(realm, Set.of(group.getId()));
+                    }
+                    return 0;
+                },
+                realm, firstResult, maxResults);
 
         return importValidation(realm, results);
     }
diff --git a/testsuite/integration-arquillian/servers/auth-server/services/testsuite-providers/src/main/java/org/keycloak/testsuite/federation/UserPropertyFileStorage.java b/testsuite/integration-arquillian/servers/auth-server/services/testsuite-providers/src/main/java/org/keycloak/testsuite/federation/UserPropertyFileStorage.java
index 59a86989b0..604a78190b 100644
--- a/testsuite/integration-arquillian/servers/auth-server/services/testsuite-providers/src/main/java/org/keycloak/testsuite/federation/UserPropertyFileStorage.java
+++ b/testsuite/integration-arquillian/servers/auth-server/services/testsuite-providers/src/main/java/org/keycloak/testsuite/federation/UserPropertyFileStorage.java
@@ -43,6 +43,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Properties;
+import java.util.Set;
 import java.util.function.Predicate;
 import java.util.stream.Stream;
 
@@ -216,6 +217,11 @@ public class UserPropertyFileStorage implements UserLookupProvider, UserStorageP
         return userPasswords.size();
     }
 
+    @Override
+    public int getUsersCount(RealmModel realm, Set<String> groupIds) {
+        return 0;
+    }
+
 //    @Override
 //    public Stream<UserModel> searchForUserStream(RealmModel realm, String search, Integer firstResult, Integer maxResults) {
 //        addCall(SEARCH_METHOD, firstResult, maxResults);