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

SVN Revision: 2365

diff --git a/src/ejabberd_auth_internal.erl b/src/ejabberd_auth_internal.erl
index a3851b532546d9dc50712b770085cc423437bba1..2a6e729655185a54ed4101d672cc0cf2b99ccb70 100644 (file)
--- 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;

diff --git a/src/ejabberd_sm.erl b/src/ejabberd_sm.erl
index 883a90103fd73dcc626e71faf9ee368c5e206eff..b42746ea80d0ef4fc508fccd660dfaa1e7aa642f 100644 (file)
--- a/src/ejabberd_sm.erl
+++ b/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).
+
 
 %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%