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

SVN Revision: 2365

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),
+    LServer = jlib:nameprep(Server),
-    length(Set).
+    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;

src/ejabberd_sm.erl

@@ -43,6 +43,7 @@
 	 dirty_get_sessions_list/0,
 	 dirty_get_my_sessions_list/0,
 	 get_vh_session_list/1,
+	 get_vh_session_number/1,
 	 register_iq_handler/4,
 	 register_iq_handler/5,
 	 unregister_iq_handler/2,
@@ -64,6 +65,7 @@
 -include("mod_privacy.hrl").
 
 -record(session, {sid, usr, us, priority, info}).
+-record(session_counter, {vhost, count}).
 -record(state, {}).
 
 %% default value for the maximum number of user connections
@@ -90,6 +92,7 @@ route(From, To, Packet) ->
 
 open_session(SID, User, Server, Resource, Info) ->
     set_session(SID, User, Server, Resource, undefined, Info),
+    inc_session_counter(jlib:nameprep(Server)),
     check_for_sessions_to_replace(User, Server, Resource),
     JID = jlib:make_jid(User, Server, Resource),
     ejabberd_hooks:run(sm_register_connection_hook, JID#jid.lserver,
@@ -101,7 +104,8 @@ close_session(SID, User, Server, Resource) ->
 	[#session{info=I}] -> I
     end,
     F = fun() ->
-		mnesia:delete({session, SID})
+		mnesia:delete({session, SID}),
+		dec_session_counter(jlib:nameprep(Server))
 	end,
     mnesia:sync_dirty(F),
     JID = jlib:make_jid(User, Server, Resource),
@@ -214,6 +218,19 @@ get_vh_session_list(Server) ->
 	[{'==', {element, 2, '$1'}, LServer}],
 	['$1']}]).
 
+get_vh_session_number(Server) ->
+    LServer = jlib:nameprep(Server),
+    Query = mnesia:dirty_select(
+		session_counter,
+		[{#session_counter{vhost = LServer, count = '$1'},
+		  [],
+		  ['$1']}]),
+    case Query of
+	[Count] ->
+	    Count;
+	_ -> 0
+    end.
+    
 register_iq_handler(Host, XMLNS, Module, Fun) ->
     ejabberd_sm ! {register_iq_handler, Host, XMLNS, Module, Fun}.
 
@@ -240,6 +257,9 @@ init([]) ->
     mnesia:create_table(session,
 			[{ram_copies, [node()]},
 			 {attributes, record_info(fields, session)}]),
+    mnesia:create_table(session_counter,
+			[{ram_copies, [node()]},
+			 {attributes, record_info(fields, session_counter)}]),
     mnesia:add_table_index(session, usr),
     mnesia:add_table_index(session, us),
     mnesia:add_table_copy(session, node(), ram_copies),
@@ -361,6 +381,8 @@ clean_table_from_bad_node(Node) ->
 			 [{'==', {node, '$1'}, Node}],
 			 ['$_']}]),
 		lists:foreach(fun(E) ->
+				      {_, LServer} = E#session.us,
+				      dec_session_counter(LServer),
 				      mnesia:delete({session, E#session.sid})
 			      end, Es)
 	end,
@@ -641,6 +663,33 @@ get_max_user_sessions(LUser, Host) ->
 	_ -> ?MAX_USER_SESSIONS
     end.
 
+inc_session_counter(LServer) ->
+    F = fun() ->
+		case mnesia:wread({session_counter, LServer}) of
+		    [C] ->
+			Count = C#session_counter.count + 1,
+			C2 = C#session_counter{count = Count},
+			mnesia:write(C2);
+		    _ ->
+			mnesia:write(#session_counter{vhost = LServer,
+						      count = 1})
+		end
+	end,
+    mnesia:sync_dirty(F).
+
+dec_session_counter(LServer) ->
+    F = fun() ->
+		case mnesia:wread({session_counter, LServer}) of
+		    [C] ->
+			Count = C#session_counter.count - 1,
+			C2 = C#session_counter{count = Count},
+			mnesia:write(C2);
+		    _ ->
+			error
+		end
+	end,
+    mnesia:sync_dirty(F).
+
 
 %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%