From 41484213fde0611f3f55a1fd1b1131cef7f0e961 Mon Sep 17 00:00:00 2001
From: Badlop <badlop@process-one.net>
Date: Wed, 19 May 2010 23:56:55 +0200
Subject: [PATCH] Use multicast in C2S presence broadcast (EJAB-267)

---
 src/ejabberd_c2s.erl | 124 ++++++++++++++++++-------------------------
 1 file changed, 53 insertions(+), 71 deletions(-)

diff --git a/src/ejabberd_c2s.erl b/src/ejabberd_c2s.erl
index 195492268..2991677e1 100644
--- a/src/ejabberd_c2s.erl
+++ b/src/ejabberd_c2s.erl
@@ -1777,87 +1777,69 @@ is_privacy_allow(From, To, Packet, PrivacyList) ->
 		{From, To, Packet},
 		in]).
 
+%% Send presence when disconnecting
 presence_broadcast(StateData, From, JIDSet, Packet) ->
-    lists:foreach(fun({U, S, R}) ->
-			  FJID = exmpp_jid:make(U, S, R),
-			  case ejabberd_hooks:run_fold(
-				 privacy_check_packet, StateData#state.server,
-				 allow,
-				 [StateData#state.user,
-				  StateData#state.server,
-				  StateData#state.privacy_list,
-				  {From, FJID, Packet},
-				  out]) of
-			      deny ->
-				  ok;
-			      allow ->
-				  ejabberd_router:route(From, FJID, Packet)
-			  end
-		  end, ?SETS:to_list(JIDSet)).
-
-presence_broadcast_to_trusted(StateData, From, T, A, Packet) ->
-    lists:foreach(
-      fun({U, S, R} = JID) ->
-	      case ?SETS:is_element(JID, T) of
-		  true ->
-		      FJID = exmpp_jid:make(U, S, R),
-		      case ejabberd_hooks:run_fold(
-			     privacy_check_packet, StateData#state.server,
-			     allow,
-			     [StateData#state.user,
-			      StateData#state.server,
-			      StateData#state.privacy_list,
-			      {From, FJID, Packet},
-			      out]) of
-			  deny ->
-			      ok;
-			  allow ->
-			      ejabberd_router:route(From, FJID, Packet)
-		      end;
-		  _ ->
-		      ok
-	      end
-      end, ?SETS:to_list(A)).
+    JIDs = ?SETS:to_list(JIDSet),
+    JIDs2 = format_and_check_privacy(From, StateData, Packet, JIDs),
+    Server = StateData#state.server,
+    send_multiple(From, Server, JIDs2, Packet).
 
+%% Send presence when updating presence
+presence_broadcast_to_trusted(StateData, From, Trusted, JIDSet, Packet) ->
+    JIDs = ?SETS:to_list(JIDSet),
+    JIDs_trusted = [JID || JID <- JIDs, ?SETS:is_element(JID, Trusted)],
+    JIDs2 = format_and_check_privacy(From, StateData, Packet, JIDs_trusted),
+    Server = StateData#state.server,
+    send_multiple(From, Server, JIDs2, Packet).
 
+%% Send presence when connecting
 presence_broadcast_first(From, StateData, Packet) ->
-    Probe = exmpp_presence:probe(),
-    ?SETS:fold(fun({U, S, R}, X) ->
-		       FJID = exmpp_jid:make(U, S, R),
-		       ejabberd_router:route(
-			 From,
-			 FJID,
-			 Probe),
-		       X
-	       end,
-	       [],
-	       StateData#state.pres_t),
+    JIDsProbe = 
+	?SETS:fold(
+	   fun(JID, L) -> [JID | L] end,
+	   [],
+	   StateData#state.pres_t),
+    PacketProbe = exmpp_presence:probe(),
+    JIDs2Probe = format_and_check_privacy(From, StateData, Packet, JIDsProbe),
+    Server = StateData#state.server,
+    send_multiple(From, Server, JIDs2Probe, PacketProbe),
     if
 	StateData#state.pres_invis ->
 	    StateData;
 	true ->
-	    As = ?SETS:fold(
-		    fun({U, S, R} = JID, A) ->
-			    FJID = exmpp_jid:make(U, S, R),
-			    case ejabberd_hooks:run_fold(
-				   privacy_check_packet, StateData#state.server,
-				   allow,
-				   [StateData#state.user,
-				    StateData#state.server,
-				    StateData#state.privacy_list,
-				    {From, FJID, Packet},
-				    out]) of
-				deny ->
-				    ok;
-				allow ->
-				    ejabberd_router:route(From, FJID, Packet)
-			    end,
-			    ?SETS:add_element(JID, A)
-		    end,
-		    StateData#state.pres_a,
-		    StateData#state.pres_f),
+	    {As, JIDs} = 
+		?SETS:fold(
+		   fun(JID, {A, JID_list}) ->
+			   {?SETS:add_element(JID, A), JID_list++[JID]}
+		   end,
+		   {StateData#state.pres_a, []},
+		   StateData#state.pres_f),
+	    JIDs2 = format_and_check_privacy(From, StateData, Packet, JIDs),
+	    Server = StateData#state.server,
+	    send_multiple(From, Server, JIDs2, Packet),
 	    StateData#state{pres_a = As}
     end.
+ 
+format_and_check_privacy(From, StateData, Packet, JIDs) ->
+    FJIDs = [exmpp_jid:make(JID) || JID <- JIDs],
+    lists:filter(
+      fun(FJID) ->
+	      case ejabberd_hooks:run_fold(
+		     privacy_check_packet, StateData#state.server,
+		     allow,
+		     [StateData#state.user,
+		      StateData#state.server,
+		      StateData#state.privacy_list,
+		      {From, FJID, Packet},
+		      out]) of
+		  deny -> false;
+		  allow -> true
+	      end
+      end,
+      FJIDs).
+
+send_multiple(From, Server, JIDs, Packet) ->
+    ejabberd_router_multicast:route_multicast(From, Server, JIDs, Packet).
 
 
 remove_element(E, Set) ->