From 08a72c1693dd1f3faeae01ee5cf35a2d651b63b8 Mon Sep 17 00:00:00 2001 From: Badlop Date: Fri, 17 Jul 2009 19:05:55 +0000 Subject: [PATCH] Store account number in internal auth of (EJAB-981)(thanks to Juan Pablo Carlino) SVN Revision: 2365 --- src/ejabberd_auth_internal.erl | 65 +++++++++++++++++++++++++++++++--- src/ejabberd_sm.erl | 51 +++++++++++++++++++++++++- 2 files changed, 111 insertions(+), 5 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; diff --git a/src/ejabberd_sm.erl b/src/ejabberd_sm.erl index 883a90103..b42746ea8 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, @@ -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). + %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%