new CAPS processing

2 files changed, 194 insertions, 37 deletions

diff --git a/src/ejabberd_c2s.erl b/src/ejabberd_c2s.erl
index d4b68096e..88075300c 100644
--- a/src/ejabberd_c2s.erl
+++ b/src/ejabberd_c2s.erl
@@ -40,6 +40,11 @@
 	 send_element/2,
 	 socket_type/0,
 	 get_presence/1,
+	 get_aux_field/2,
+	 set_aux_field/3,
+	 del_aux_field/2,
+	 get_subscription/2,
+	 broadcast/4,
 	 get_subscribed/1]).
 
 %% gen_fsm callbacks
@@ -96,6 +101,7 @@
 		conn = unknown,
 		auth_module = unknown,
 		ip,
+		aux_fields = [],
 		lang}).
 
 %-define(DBGFSM, true).
@@ -154,6 +160,39 @@ socket_type() ->
 get_presence(FsmRef) ->
     ?GEN_FSM:sync_send_all_state_event(FsmRef, {get_presence}, 1000).
 
+get_aux_field(Key, #state{aux_fields = Opts}) ->
+    case lists:keysearch(Key, 1, Opts) of
+	{value, {_, Val}} ->
+	    {ok, Val};
+	_ ->
+	    error
+    end.
+
+set_aux_field(Key, Val, #state{aux_fields = Opts} = State) ->
+    Opts1 = lists:keydelete(Key, 1, Opts),
+    State#state{aux_fields = [{Key, Val}|Opts1]}.
+
+del_aux_field(Key, #state{aux_fields = Opts} = State) ->
+    Opts1 = lists:keydelete(Key, 1, Opts),
+    State#state{aux_fields = Opts1}.
+
+get_subscription(From = #jid{}, StateData) ->
+    get_subscription(jlib:jid_tolower(From), StateData);
+get_subscription(LFrom, StateData) ->
+    LBFrom = setelement(3, LFrom, ""),
+    F = ?SETS:is_element(LFrom, StateData#state.pres_f) orelse
+	?SETS:is_element(LBFrom, StateData#state.pres_f),
+    T = ?SETS:is_element(LFrom, StateData#state.pres_t) orelse
+	?SETS:is_element(LBFrom, StateData#state.pres_t),
+    if F and T -> both;
+       F -> from;
+       T -> to;
+       true -> none
+    end.
+
+broadcast(FsmRef, Type, From, Packet) ->
+    FsmRef ! {broadcast, Type, From, Packet}.
+
 stop(FsmRef) ->
     ?GEN_FSM:send_event(FsmRef, closed).
 
@@ -1103,35 +1142,39 @@ handle_info({route, From, To, Packet}, StateName, StateData) ->
     {Pass, NewAttrs, NewState} =
 	case Name of
 	    "presence" ->
+		State = ejabberd_hooks:run_fold(
+			  c2s_presence_in, StateData#state.server,
+			  StateData,
+			  [{From, To, Packet}]),
 		case xml:get_attr_s("type", Attrs) of
 		    "probe" ->
 			LFrom = jlib:jid_tolower(From),
 			LBFrom = jlib:jid_remove_resource(LFrom),
 			NewStateData =
 			    case ?SETS:is_element(
-				    LFrom, StateData#state.pres_a) orelse
+				    LFrom, State#state.pres_a) orelse
 				?SETS:is_element(
-				   LBFrom, StateData#state.pres_a) of
+				   LBFrom, State#state.pres_a) of
 				true ->
-				    StateData;
+				    State;
 				false ->
 				    case ?SETS:is_element(
-					    LFrom, StateData#state.pres_f) of
+					    LFrom, State#state.pres_f) of
 					true ->
 					    A = ?SETS:add_element(
 						   LFrom,
-						   StateData#state.pres_a),
-					    StateData#state{pres_a = A};
+						   State#state.pres_a),
+					    State#state{pres_a = A};
 					false ->
 					    case ?SETS:is_element(
-						    LBFrom, StateData#state.pres_f) of
+						    LBFrom, State#state.pres_f) of
 						true ->
 						    A = ?SETS:add_element(
 							   LBFrom,
-							   StateData#state.pres_a),
-						    StateData#state{pres_a = A};
+							   State#state.pres_a),
+						    State#state{pres_a = A};
 						false ->
-						    StateData
+						    State
 					    end
 				    end
 			    end,
@@ -1139,66 +1182,66 @@ handle_info({route, From, To, Packet}, StateName, StateData) ->
 			{false, Attrs, NewStateData};
 		    "error" ->
 			NewA = remove_element(jlib:jid_tolower(From),
-					      StateData#state.pres_a),
-			{true, Attrs, StateData#state{pres_a = NewA}};
+					      State#state.pres_a),
+			{true, Attrs, State#state{pres_a = NewA}};
 		    "invisible" ->
 			Attrs1 = lists:keydelete("type", 1, Attrs),
-			{true, [{"type", "unavailable"} | Attrs1], StateData};
+			{true, [{"type", "unavailable"} | Attrs1], State};
 		    "subscribe" ->
-			SRes = is_privacy_allow(From, To, Packet, StateData#state.privacy_list),
-			{SRes, Attrs, StateData};
+			SRes = is_privacy_allow(From, To, Packet, State#state.privacy_list),
+			{SRes, Attrs, State};
 		    "subscribed" ->
-			SRes = is_privacy_allow(From, To, Packet, StateData#state.privacy_list),
-			{SRes, Attrs, StateData};
+			SRes = is_privacy_allow(From, To, Packet, State#state.privacy_list),
+			{SRes, Attrs, State};
 		    "unsubscribe" ->
-			SRes = is_privacy_allow(From, To, Packet, StateData#state.privacy_list),
-			{SRes, Attrs, StateData};
+			SRes = is_privacy_allow(From, To, Packet, State#state.privacy_list),
+			{SRes, Attrs, State};
 		    "unsubscribed" ->
-			SRes = is_privacy_allow(From, To, Packet, StateData#state.privacy_list),
-			{SRes, Attrs, StateData};
+			SRes = is_privacy_allow(From, To, Packet, State#state.privacy_list),
+			{SRes, Attrs, State};
 		    _ ->
 			case ejabberd_hooks:run_fold(
-			       privacy_check_packet, StateData#state.server,
+			       privacy_check_packet, State#state.server,
 			       allow,
-			       [StateData#state.user,
-				StateData#state.server,
-				StateData#state.privacy_list,
+			       [State#state.user,
+				State#state.server,
+				State#state.privacy_list,
 				{From, To, Packet},
 				in]) of
 			    allow ->
 				LFrom = jlib:jid_tolower(From),
 				LBFrom = jlib:jid_remove_resource(LFrom),
 				case ?SETS:is_element(
-					LFrom, StateData#state.pres_a) orelse
+					LFrom, State#state.pres_a) orelse
 				    ?SETS:is_element(
-				       LBFrom, StateData#state.pres_a) of
+				       LBFrom, State#state.pres_a) of
 				    true ->
-					{true, Attrs, StateData};
+					{true, Attrs, State};
 				    false ->
 					case ?SETS:is_element(
-						LFrom, StateData#state.pres_f) of
+						LFrom, State#state.pres_f) of
 					    true ->
 						A = ?SETS:add_element(
 						       LFrom,
-						       StateData#state.pres_a),
+						       State#state.pres_a),
 						{true, Attrs,
-						 StateData#state{pres_a = A}};
+						 State#state{pres_a = A}};
 					    false ->
 						case ?SETS:is_element(
-							LBFrom, StateData#state.pres_f) of
+							LBFrom, State#state.pres_f) of
 						    true ->
 							A = ?SETS:add_element(
 							       LBFrom,
-							       StateData#state.pres_a),
+							       State#state.pres_a),
 							{true, Attrs,
-							 StateData#state{pres_a = A}};
+							 State#state{pres_a = A}};
 						    false ->
-							{true, Attrs, StateData}
+							{true, Attrs, State}
 						end
 					end
 				end;
 			    deny ->
-				{false, Attrs, StateData}
+				{false, Attrs, State}
 			end
 		end;
 	    "broadcast" ->
@@ -1360,6 +1403,17 @@ handle_info({force_update_presence, LUser}, StateName,
 		StateData
 	end,
     {next_state, StateName, NewStateData};
+handle_info({broadcast, Type, From, Packet}, StateName, StateData) ->
+    Recipients = ejabberd_hooks:run_fold(
+		   c2s_broadcast_recipients, StateData#state.server,
+		   [],
+		   [StateData, Type, From, Packet]),
+    lists:foreach(
+      fun(USR) ->
+	      ejabberd_router:route(
+		From, jlib:make_jid(USR), Packet)
+      end, lists:usort(Recipients)),
+    fsm_next_state(StateName, StateData);
 handle_info(Info, StateName, StateData) ->
     ?ERROR_MSG("Unexpected info: ~p", [Info]),
     fsm_next_state(StateName, StateData).
diff --git a/src/mod_caps.erl b/src/mod_caps.erl
index 319e37e6c..10f2f3414 100644
--- a/src/mod_caps.erl
+++ b/src/mod_caps.erl
@@ -53,7 +53,10 @@
 	]).
 
 %% hook handlers
--export([user_send_packet/3]).
+-export([user_send_packet/3,
+	 user_receive_packet/4,
+	 c2s_presence_in/2,
+	 c2s_broadcast_recipients/5]).
 
 -include("ejabberd.hrl").
 -include("jlib.hrl").
@@ -159,6 +162,22 @@ user_send_packet(#jid{luser = User, lserver = Server} = From,
 user_send_packet(_From, _To, _Packet) ->
     ok.
 
+user_receive_packet(#jid{lserver = Server}, From, _To,
+		    {xmlelement, "presence", Attrs, Els}) ->
+    Type = xml:get_attr_s("type", Attrs),
+    if Type == ""; Type == "available" ->
+	    case read_caps(Els) of
+		nothing ->
+		    ok;
+		#caps{version = Version, exts = Exts} = Caps ->
+		    feature_request(Server, From, Caps, [Version | Exts])
+	    end;
+       true ->
+	    ok
+    end;
+user_receive_packet(_JID, _From, _To, _Packet) ->
+    ok.
+
 caps_stream_features(Acc, MyHost) ->
     case make_my_disco_hash(MyHost) of
 	"" ->
@@ -194,6 +213,65 @@ disco_info(_Acc, Host, Module, ?EJABBERD_URI ++ "#" ++ [_|_], Lang) ->
 disco_info(Acc, _Host, _Module, _Node, _Lang) ->
     Acc.
 
+c2s_presence_in(C2SState, {From, To, {_, _, Attrs, Els}}) ->
+    Type = xml:get_attr_s("type", Attrs),
+    Subscription = ejabberd_c2s:get_subscription(From, C2SState),
+    Insert = ((Type == "") or (Type == "available"))
+	and ((Subscription == both) or (Subscription == to)),
+    Delete = (Type == "unavailable") or (Type == "error") or (Type == "invisible"),
+    if Insert or Delete ->
+	    LFrom = jlib:jid_tolower(From),
+	    Rs = case ejabberd_c2s:get_aux_field(caps_resources, C2SState) of
+		     {ok, Rs1} ->
+			 Rs1;
+		     error ->
+			 gb_trees:empty()
+		 end,
+	    Caps = read_caps(Els),
+	    {FromUnavail, NewRs} =
+		case Caps of
+		    nothing when Insert == true ->
+			{false, Rs};
+		    _ when Insert == true ->
+			case gb_trees:lookup(LFrom, Rs) of
+			    none ->
+				{true, gb_trees:insert(LFrom, Caps, Rs)};
+			    _ ->
+				{false, gb_trees:update(LFrom, Caps, Rs)}
+			end;
+		    _ ->
+			{false, gb_trees:delete_any(LFrom, Rs)}
+		end,
+	    if FromUnavail ->
+		    ejabberd_hooks:run(caps_user_available, To#jid.lserver,
+				       [From, To, get_features(Caps)]);
+	       true ->
+		    ok
+	    end,
+	    ejabberd_c2s:set_aux_field(caps_resources, NewRs, C2SState);
+       true ->
+	    C2SState
+    end.
+
+c2s_broadcast_recipients(InAcc, C2SState, {pep_message, Feature},
+			 _From, _Packet) ->
+    case ejabberd_c2s:get_aux_field(caps_resources, C2SState) of
+	{ok, Rs} ->
+	    gb_trees_fold(
+	      fun(USR, Caps, Acc) ->
+		      case lists:member(Feature, get_features(Caps)) of
+			  true ->
+			      [USR|Acc];
+			  false ->
+			      Acc
+		      end
+	      end, InAcc, Rs);
+	_ ->
+	    InAcc
+    end;
+c2s_broadcast_recipients(Acc, _, _, _, _) ->
+    Acc.
+
 %%====================================================================
 %% gen_server callbacks
 %%====================================================================
@@ -214,8 +292,14 @@ init([Host, Opts]) ->
     MaxSize = gen_mod:get_opt(cache_size, Opts, 1000),
     LifeTime = gen_mod:get_opt(cache_life_time, Opts, timer:hours(24) div 1000),
     cache_tab:new(caps_features, [{max_size, MaxSize}, {life_time, LifeTime}]),
+    ejabberd_hooks:add(c2s_presence_in, Host,
+		       ?MODULE, c2s_presence_in, 75),
+    ejabberd_hooks:add(c2s_broadcast_recipients, Host,
+		       ?MODULE, c2s_broadcast_recipients, 75),
     ejabberd_hooks:add(user_send_packet, Host,
 		       ?MODULE, user_send_packet, 75),
+    ejabberd_hooks:add(user_receive_packet, Host,
+		       ?MODULE, user_receive_packet, 75),
     ejabberd_hooks:add(c2s_stream_features, Host,
 		       ?MODULE, caps_stream_features, 75),
     ejabberd_hooks:add(s2s_stream_features, Host,
@@ -241,8 +325,14 @@ handle_info(_Info, State) ->
 
 terminate(_Reason, State) ->
     Host = State#state.host,
+    ejabberd_hooks:delete(c2s_presence_in, Host,
+			  ?MODULE, c2s_presence_in, 75),
+    ejabberd_hooks:delete(c2s_broadcast_recipients, Host,
+			  ?MODULE, c2s_broadcast_recipients, 75),
     ejabberd_hooks:delete(user_send_packet, Host,
 			  ?MODULE, user_send_packet, 75),
+    ejabberd_hooks:delete(user_receive_packet, Host,
+			  ?MODULE, user_receive_packet, 75),
     ejabberd_hooks:delete(c2s_stream_features, Host,
 			  ?MODULE, caps_stream_features, 75),
     ejabberd_hooks:delete(s2s_stream_features, Host,
@@ -506,3 +596,16 @@ concat_xdata_fields(Fields) ->
 		  Acc
 	  end, ["", []], Fields),
     [Form, $<, lists:sort(Res)].
+
+gb_trees_fold(F, Acc, Tree) ->
+    Iter = gb_trees:iterator(Tree),
+    gb_trees_fold_iter(F, Acc, Iter).
+
+gb_trees_fold_iter(F, Acc, Iter) ->
+    case gb_trees:next(Iter) of
+	{Key, Val, NewIter} ->
+	    NewAcc = F(Key, Val, Acc),
+	    gb_trees_fold_iter(F, NewAcc, NewIter);
+	_ ->
+	    Acc
+    end.