Add auth:which_user_exist to bulk checking existence of list of users

2 files changed, 70 insertions, 2 deletions

diff --git a/src/ejabberd_auth.erl b/src/ejabberd_auth.erl
index 5659ee38..bc021154 100644
--- a/src/ejabberd_auth.erl
+++ b/src/ejabberd_auth.erl
@@ -41,7 +41,8 @@
 	 get_password_s/2, get_password_with_authmodule/2,
 	 user_exists/2, user_exists_in_other_modules/3,
 	 remove_user/2, remove_user/3, plain_password_required/1,
-	 store_type/1, entropy/1, backend_type/1, password_format/1]).
+	 store_type/1, entropy/1, backend_type/1, password_format/1,
+	 which_users_exists/1]).
 %% gen_server callbacks
 -export([init/1, handle_call/3, handle_cast/2, handle_info/2,
 	 terminate/2, code_change/3]).
@@ -411,6 +412,47 @@ user_exists_in_other_modules_loop([AuthModule | AuthModules], User, Server) ->
 	    maybe
     end.
 
+-spec which_users_exists(list({binary(), binary()})) -> list({binary(), binary()}).
+which_users_exists(USPairs) ->
+    ByServer = lists:foldl(
+	fun({User, Server}, Dict) ->
+	    LServer = jid:nameprep(Server),
+	    LUser =  jid:nodeprep(User),
+	    case gb_trees:lookup(LServer, Dict) of
+		none ->
+		    gb_trees:insert(LServer, gb_sets:singleton(LUser), Dict);
+		{value, Set} ->
+		    gb_trees:update(LServer, gb_sets:add(LUser, Set), Dict)
+	    end
+	end, gb_trees:empty(), USPairs),
+    Set = lists:foldl(
+	fun({LServer, UsersSet}, Results) ->
+	    UsersList = gb_sets:to_list(UsersSet),
+	    lists:foldl(
+		fun(M, Results2) ->
+		    try M:which_users_exists(LServer, UsersList) of
+			{error, _} ->
+			    Results2;
+			Res ->
+			    gb_sets:union(
+				gb_sets:from_list([{U, LServer} || U <- Res]),
+				Results2)
+		    catch
+			_:undef ->
+			    lists:foldl(
+				fun(U, R2) ->
+				    case user_exists(U, LServer) of
+					true ->
+					    gb_sets:add({U, LServer}, R2);
+					_ ->
+					    R2
+				    end
+				end, Results2, UsersList)
+		    end
+		end, Results, auth_modules(LServer))
+	end, gb_sets:empty(), gb_trees:to_list(ByServer)),
+    gb_sets:to_list(Set).
+
 -spec remove_user(binary(), binary()) -> ok.
 remove_user(User, Server) ->
     case validate_credentials(User, Server) of
diff --git a/src/ejabberd_auth_sql.erl b/src/ejabberd_auth_sql.erl
index 4b774642..cd9c02b9 100644
--- a/src/ejabberd_auth_sql.erl
+++ b/src/ejabberd_auth_sql.erl
@@ -35,7 +35,7 @@
 -export([start/1, stop/1, set_password/3, try_register/3,
 	 get_users/2, count_users/2, get_password/2,
 	 remove_user/2, store_type/1, plain_password_required/1,
-	 convert_to_scram/1, opt_type/1, export/1]).
+	 convert_to_scram/1, opt_type/1, export/1, which_users_exists/2]).
 
 -include("scram.hrl").
 -include("logger.hrl").
@@ -247,6 +247,32 @@ users_number(LServer, [{prefix, Prefix}])
 users_number(LServer, []) ->
     users_number(LServer).
 
+which_users_exists(LServer, LUsers) when length(LUsers) =< 100 ->
+    try ejabberd_sql:sql_query(
+        LServer,
+        ?SQL("select @(username)s from users where username in %(LUsers)ls")) of
+        {selected, Matching} ->
+            [U || {U} <- Matching];
+        {error, _} = E ->
+            E
+    catch _:B ->
+        {error, B}
+    end;
+which_users_exists(LServer, LUsers) ->
+    {First, Rest} = lists:split(100, LUsers),
+    case which_users_exists(LServer, First) of
+        {error, _} = E ->
+            E;
+        V ->
+            case which_users_exists(LServer, Rest) of
+                {error, _} = E2 ->
+                    E2;
+                V2 ->
+                    V ++ V2
+            end
+    end.
+
+
 convert_to_scram(Server) ->
     LServer = jid:nameprep(Server),
     if