Implement database backend interface for MUC, BOSH and auth_anonyous

1 files changed, 206 insertions, 12 deletions

diff --git a/src/mod_muc_mnesia.erl b/src/mod_muc_mnesia.erl
index 6464d431..0552184b 100644
--- a/src/mod_muc_mnesia.erl
+++ b/src/mod_muc_mnesia.erl
@@ -14,28 +14,55 @@
 %% API
 -export([init/2, import/3, store_room/4, restore_room/3, forget_room/3,
 	 can_use_nick/4, get_rooms/2, get_nick/3, set_nick/4]).
+-export([register_online_room/3, unregister_online_room/3, find_online_room/2,
+	 get_online_rooms/2, count_online_rooms/1, rsm_supported/0,
+	 register_online_user/3, unregister_online_user/3,
+	 count_online_rooms_by_user/2, get_online_rooms_by_user/2,
+	 handle_event/1]).
 -export([set_affiliation/6, set_affiliations/4, get_affiliation/5,
 	 get_affiliations/3, search_affiliation/4]).
 
--include("jid.hrl").
 -include("mod_muc.hrl").
 -include("logger.hrl").
+-include("xmpp.hrl").
+-include_lib("stdlib/include/ms_transform.hrl").
 
 %%%===================================================================
 %%% API
 %%%===================================================================
-init(_Host, Opts) ->
+init(Host, Opts) ->
     MyHost = proplists:get_value(host, Opts),
-    ejabberd_mnesia:create(?MODULE, muc_room,
-			[{disc_copies, [node()]},
-			 {attributes,
-			  record_info(fields, muc_room)}]),
-    ejabberd_mnesia:create(?MODULE, muc_registered,
-			[{disc_copies, [node()]},
-			 {attributes,
-			  record_info(fields, muc_registered)}]),
-    update_tables(MyHost),
-    mnesia:add_table_index(muc_registered, nick).
+    case gen_mod:db_mod(Host, Opts, mod_muc) of
+	?MODULE ->
+	    ejabberd_mnesia:create(?MODULE, muc_room,
+				   [{disc_copies, [node()]},
+				    {attributes,
+				     record_info(fields, muc_room)}]),
+	    ejabberd_mnesia:create(?MODULE, muc_registered,
+				   [{disc_copies, [node()]},
+				    {attributes,
+				     record_info(fields, muc_registered)}]),
+	    update_tables(MyHost),
+	    mnesia:add_table_index(muc_registered, nick);
+	_ ->
+	    ok
+    end,
+    case gen_mod:ram_db_mod(Host, Opts, mod_muc) of
+	?MODULE ->
+	    update_muc_online_table(),
+	    ejabberd_mnesia:create(?MODULE, muc_online_room,
+				   [{ram_copies, [node()]},
+				    {type, ordered_set},
+				    {attributes,
+				     record_info(fields, muc_online_room)}]),
+	    mnesia:add_table_copy(muc_online_room, node(), ram_copies),
+	    catch ets:new(muc_online_users,
+			  [bag, named_table, public, {keypos, 2}]),
+	    clean_table_from_bad_node(node(), MyHost),
+	    mnesia:subscribe(system);
+	_ ->
+	    ok
+    end.
 
 store_room(_LServer, Host, Name, Opts) ->
     F = fun () ->
@@ -132,6 +159,128 @@ get_affiliations(_ServerHost, _Room, _Host) ->
 search_affiliation(_ServerHost, _Room, _Host, _Affiliation) ->
     {error, not_implemented}.
 
+register_online_room(Room, Host, Pid) ->
+    F = fun() ->
+		mnesia:write(
+		  #muc_online_room{name_host = {Room, Host}, pid = Pid})
+	end,
+    mnesia:transaction(F).
+
+unregister_online_room(Room, Host, Pid) ->
+    F = fun () ->
+		mnesia:delete_object(
+		  #muc_online_room{name_host = {Room, Host}, pid = Pid})
+	end,
+    mnesia:transaction(F).
+
+find_online_room(Room, Host) ->
+    case mnesia:dirty_read(muc_online_room, {Room, Host}) of
+	[] -> error;
+	[#muc_online_room{pid = Pid}] -> {ok, Pid}
+    end.
+
+count_online_rooms(Host) ->
+    ets:select_count(
+      muc_online_room,
+      ets:fun2ms(
+	fun(#muc_online_room{name_host = {_, H}}) ->
+		H == Host
+	end)).
+
+get_online_rooms(Host,
+		 #rsm_set{max = Max, 'after' = After, before = undefined})
+  when is_binary(After), After /= <<"">> ->
+    lists:reverse(get_online_rooms(next, {After, Host}, Host, 0, Max, []));
+get_online_rooms(Host,
+		 #rsm_set{max = Max, 'after' = undefined, before = Before})
+  when is_binary(Before), Before /= <<"">> ->
+    get_online_rooms(prev, {Before, Host}, Host, 0, Max, []);
+get_online_rooms(Host,
+		 #rsm_set{max = Max, 'after' = undefined, before = <<"">>}) ->
+    get_online_rooms(last, {<<"">>, Host}, Host, 0, Max, []);
+get_online_rooms(Host, #rsm_set{max = Max}) ->
+    lists:reverse(get_online_rooms(first, {<<"">>, Host}, Host, 0, Max, []));
+get_online_rooms(Host, undefined) ->
+    mnesia:dirty_select(
+      muc_online_room,
+      ets:fun2ms(
+	fun(#muc_online_room{name_host = {Name, H}, pid = Pid})
+	      when H == Host -> {Name, Host, Pid}
+	end)).
+
+-spec get_online_rooms(prev | next | last | first,
+		       {binary(), binary()}, binary(),
+		       non_neg_integer(), non_neg_integer() | undefined,
+		       [{binary(), binary(), pid()}]) ->
+			      [{binary(), binary(), pid()}].
+get_online_rooms(_Action, _Key, _Host, Count, Max, Items) when Count >= Max ->
+    Items;
+get_online_rooms(Action, Key, Host, Count, Max, Items) ->
+    Call = fun() ->
+		   case Action of
+		       prev -> mnesia:dirty_prev(muc_online_room, Key);
+		       next -> mnesia:dirty_next(muc_online_room, Key);
+		       last -> mnesia:dirty_last(muc_online_room);
+		       first -> mnesia:dirty_first(muc_online_room)
+		   end
+	   end,
+    NewAction = case Action of
+		    last -> prev;
+		    first -> next;
+		    _ -> Action
+		end,
+    try Call() of
+	'$end_of_table' ->
+	    Items;
+	{Room, Host} = NewKey ->
+	    case find_online_room(Room, Host) of
+		{ok, Pid} ->
+		    get_online_rooms(NewAction, NewKey, Host,
+				     Count + 1, Max, [{Room, Host, Pid}|Items]);
+		{error, _} ->
+		    get_online_rooms(NewAction, NewKey, Host,
+				     Count, Max, Items)
+	    end;
+	NewKey ->
+	    get_online_rooms(NewAction, NewKey, Host, Count, Max, Items)
+    catch _:{aborted, {badarg, _}} ->
+	    Items
+    end.
+
+handle_event({mnesia_system_event, {mnesia_down, Node}}) ->
+    clean_table_from_bad_node(Node);
+handle_event(_) ->
+    ok.
+
+rsm_supported() ->
+    true.
+
+register_online_user({U, S, R}, Room, Host) ->
+    ets:insert(muc_online_users,
+	       #muc_online_users{us = {U, S}, resource = R,
+				 room = Room, host = Host}).
+
+unregister_online_user({U, S, R}, Room, Host) ->
+    ets:delete_object(muc_online_users,
+		      #muc_online_users{us = {U, S}, resource = R,
+					room = Room, host = Host}).
+
+count_online_rooms_by_user(U, S) ->
+    ets:select_count(
+      muc_online_users,
+      ets:fun2ms(
+	fun(#muc_online_users{us = {U1, S1}}) ->
+		U == U1 andalso S == S1
+	end)).
+
+get_online_rooms_by_user(U, S) ->
+    ets:select(
+      muc_online_users,
+      ets:fun2ms(
+	fun(#muc_online_users{us = {U1, S1}, room = Room, host = Host})
+	      when U == U1 andalso S == S1 -> {Room, Host}
+	end)).
+
 import(_LServer, <<"muc_room">>,
        [Name, RoomHost, SOpts, _TimeStamp]) ->
     Opts = mod_muc:opts_to_binary(ejabberd_sql:decode_term(SOpts)),
@@ -148,6 +297,34 @@ import(_LServer, <<"muc_registered">>,
 %%%===================================================================
 %%% Internal functions
 %%%===================================================================
+clean_table_from_bad_node(Node) ->
+    F = fun() ->
+		Es = mnesia:select(
+		       muc_online_room,
+		       [{#muc_online_room{pid = '$1', _ = '_'},
+			 [{'==', {node, '$1'}, Node}],
+			 ['$_']}]),
+		lists:foreach(fun(E) ->
+				      mnesia:delete_object(E)
+			      end, Es)
+        end,
+    mnesia:async_dirty(F).
+
+clean_table_from_bad_node(Node, Host) ->
+    F = fun() ->
+		Es = mnesia:select(
+		       muc_online_room,
+		       [{#muc_online_room{pid = '$1',
+					  name_host = {'_', Host},
+					  _ = '_'},
+			 [{'==', {node, '$1'}, Node}],
+			 ['$_']}]),
+		lists:foreach(fun(E) ->
+				      mnesia:delete_object(E)
+			      end, Es)
+        end,
+    mnesia:async_dirty(F).
+
 update_tables(Host) ->
     update_muc_room_table(Host),
     update_muc_registered_table(Host).
@@ -188,3 +365,20 @@ update_muc_registered_table(_Host) ->
 	  ?INFO_MSG("Recreating muc_registered table", []),
 	  mnesia:transform_table(muc_registered, ignore, Fields)
     end.
+
+update_muc_online_table() ->
+    try
+	case mnesia:table_info(muc_online_room, type) of
+	    ordered_set -> ok;
+	    _ ->
+		case mnesia:delete_table(muc_online_room) of
+		    {atomic, ok} -> ok;
+		    Err -> erlang:error(Err)
+		end
+	end
+    catch _:{aborted, {no_exists, muc_online_room}} -> ok;
+	  _:{aborted, {no_exists, muc_online_room, type}} -> ok;
+	  E:R ->
+	    ?ERROR_MSG("failed to update mnesia table '~s': ~p",
+		       [muc_online_room, {E, R, erlang:get_stacktrace()}])
+    end.