Store account number in internal auth of (EJAB-981)(thanks to Juan Pablo Carlino)

SVN Revision: 2365

1 files changed, 61 insertions, 4 deletions

diff --git a/src/ejabberd_auth_internal.erl b/src/ejabberd_auth_internal.erl
index a3851b532..2a6e72965 100644
--- a/src/ejabberd_auth_internal.erl
+++ b/src/ejabberd_auth_internal.erl
@@ -49,16 +49,27 @@
 -include("ejabberd.hrl").
 
 -record(passwd, {us, password}).
+-record(reg_users_counter, {vhost, count}).
 
 %%%----------------------------------------------------------------------
 %%% API
 %%%----------------------------------------------------------------------
-start(_Host) ->
+start(Host) ->
     mnesia:create_table(passwd, [{disc_copies, [node()]},
 				 {attributes, record_info(fields, passwd)}]),
+    mnesia:create_table(reg_users_counter,
+			[{ram_copies, [node()]},
+			 {attributes, record_info(fields, reg_users_counter)}]),
     update_table(),
+    update_reg_users_counter_table(Host),
     ok.
 
+update_reg_users_counter_table(Server) ->
+    Set = get_vh_registered_users(Server),
+    Size = length(Set),
+    LServer = jlib:nameprep(Server),
+    set_vh_registered_users_counter(LServer, Size).
+
 plain_password_required() ->
     false.
 
@@ -126,6 +137,7 @@ try_register(User, Server, Password) ->
 			    [] ->
 				mnesia:write(#passwd{us = US,
 						     password = Password}),
+				inc_vh_registered_users_counter(LServer),
 				ok;
 			    [_E] ->
 				exists
@@ -193,8 +205,17 @@ get_vh_registered_users(Server, _) ->
     get_vh_registered_users(Server).
 
 get_vh_registered_users_number(Server) ->
-    Set = get_vh_registered_users(Server),
-    length(Set).
+    LServer = jlib:nameprep(Server),
+    Query = mnesia:dirty_select(
+		reg_users_counter,
+		[{#reg_users_counter{vhost = LServer, count = '$1'},
+		  [],
+		  ['$1']}]),
+    case Query of
+	[Count] ->
+	    Count;
+	_ -> 0
+    end.
 
 get_vh_registered_users_number(Server, [{prefix, Prefix}]) when is_list(Prefix) ->
     Set = [{U, S} || {U, S} <- get_vh_registered_users(Server), lists:prefix(Prefix, U)],
@@ -203,6 +224,40 @@ get_vh_registered_users_number(Server, [{prefix, Prefix}]) when is_list(Prefix)
 get_vh_registered_users_number(Server, _) ->
     get_vh_registered_users_number(Server).
 
+inc_vh_registered_users_counter(LServer) ->
+    F = fun() ->
+		case mnesia:wread({reg_users_counter, LServer}) of
+		    [C] ->
+			Count = C#reg_users_counter.count + 1,
+			C2 = C#reg_users_counter{count = Count},
+			mnesia:write(C2);
+		    _ ->
+			mnesia:write(#reg_users_counter{vhost = LServer,
+						      count = 1})
+		end
+	end,
+    mnesia:sync_dirty(F).
+
+dec_vh_registered_users_counter(LServer) ->
+    F = fun() ->
+		case mnesia:wread({reg_users_counter, LServer}) of
+		    [C] ->
+			Count = C#reg_users_counter.count - 1,
+			C2 = C#reg_users_counter{count = Count},
+			mnesia:write(C2);
+		    _ ->
+			error
+		end
+	end,
+    mnesia:sync_dirty(F).
+
+set_vh_registered_users_counter(LServer, Count) ->
+    F = fun() ->
+		mnesia:write(#reg_users_counter{vhost = LServer,
+						count = Count})
+	end,
+    mnesia:sync_dirty(F).
+
 get_password(User, Server) ->
     LUser = jlib:nodeprep(User),
     LServer = jlib:nameprep(Server),
@@ -247,7 +302,8 @@ remove_user(User, Server) ->
     LServer = jlib:nameprep(Server),
     US = {LUser, LServer},
     F = fun() ->
-		mnesia:delete({passwd, US})
+		mnesia:delete({passwd, US}),
+		dec_vh_registered_users_counter(LServer)
         end,
     mnesia:transaction(F),
 	ok.
@@ -262,6 +318,7 @@ remove_user(User, Server, Password) ->
 		case mnesia:read({passwd, US}) of
 		    [#passwd{password = Password}] ->
 			mnesia:delete({passwd, US}),
+			dec_vh_registered_users_counter(LServer),
 			ok;
 		    [_] ->
 			not_allowed;