diff --git a/src/ejabberd_auth_internal.erl b/src/ejabberd_auth_internal.erl index 74aa82c7c..61c57d7e7 100644 --- a/src/ejabberd_auth_internal.erl +++ b/src/ejabberd_auth_internal.erl @@ -49,6 +49,7 @@ -include("ejabberd.hrl"). -record(passwd, {us, password}). +-record(reg_users_counter, {vhost, count}). %%%---------------------------------------------------------------------- %%% API @@ -57,12 +58,22 @@ %% @spec (Host) -> ok %% Host = string() -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 = exmpp_jid:prep_domain(exmpp_jid:parse(Server)), + set_vh_registered_users_counter(LServer, Size). + %% @spec () -> bool() plain_password_required() -> @@ -151,6 +162,7 @@ try_register(User, Server, Password) -> [] -> mnesia:write(#passwd{us = US, password = Password}), + inc_vh_registered_users_counter(LServer), ok; [_E] -> exists @@ -249,8 +261,17 @@ get_vh_registered_users(Server, _) -> %% Users_Number = integer() get_vh_registered_users_number(Server) -> - Set = get_vh_registered_users(Server), - length(Set). + LServer = exmpp_jid:prep_domain(exmpp_jid:parse(Server)), + Query = mnesia:dirty_select( + reg_users_counter, + [{#reg_users_counter{vhost = LServer, count = '$1'}, + [], + ['$1']}]), + case Query of + [Count] -> + Count; + _ -> 0 + end. %% @spec (Server, [{prefix, Prefix}]) -> Users_Number %% Server = string() @@ -264,6 +285,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). + %% @spec (User, Server) -> Password | false %% User = string() %% Server = string() @@ -341,7 +396,8 @@ remove_user(User, Server) -> LServer = exmpp_stringprep: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 @@ -365,6 +421,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 511f7ec5e..009d32bd7 100644 --- 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, @@ -65,6 +66,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 @@ -111,6 +113,7 @@ route(From, To, Packet) -> open_session(SID, JID, Info) when ?IS_JID(JID) -> set_session(SID, JID, undefined, Info), + inc_session_counter(exmpp_jid:domain(JID)), check_for_sessions_to_replace(JID), ejabberd_hooks:run(sm_register_connection_hook, exmpp_jid:prep_domain(JID), [SID, JID, Info]). @@ -121,7 +124,8 @@ close_session(SID, JID ) when ?IS_JID(JID)-> [#session{info=I}] -> I end, F = fun() -> - mnesia:delete({session, SID}) + mnesia:delete({session, SID}), + dec_session_counter(exmpp_jid:domain(JID)) end, mnesia:sync_dirty(F), ejabberd_hooks:run(sm_remove_connection_hook, exmpp_jid:prep_domain(JID), @@ -249,6 +253,19 @@ get_vh_session_list(Server) when is_binary(Server) -> [{'==', {element, 2, '$1'}, LServer}], ['$1']}]). +get_vh_session_number(Server) -> + LServer = exmpp_jid:prep_domain(exmpp_jid:parse(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}. @@ -275,6 +292,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), @@ -406,6 +426,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, @@ -684,6 +706,33 @@ get_max_user_sessions(JID) -> _ -> ?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). + %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%