Don't lose carbons on presence change or session resumption

3 files changed, 77 insertions, 28 deletions

diff --git a/src/ejabberd_c2s.erl b/src/ejabberd_c2s.erl
index c65e71bdf..ba5b04af8 100644
--- a/src/ejabberd_c2s.erl
+++ b/src/ejabberd_c2s.erl
@@ -708,13 +708,11 @@ process_presence_out(#{lserver := LServer, jid := JID,
     end.
 
 -spec process_self_presence(state(), presence()) -> state().
-process_self_presence(#{ip := IP, conn := Conn, lserver := LServer,
-			auth_module := AuthMod, sid := SID,
+process_self_presence(#{lserver := LServer, sid := SID,
 			user := U, server := S,	resource := R} = State,
 		      #presence{type = unavailable} = Pres) ->
     Status = xmpp:get_text(Pres#presence.status),
-    Info = [{ip, IP}, {conn, Conn}, {auth_module, AuthMod}],
-    ejabberd_sm:unset_presence(SID, U, S, R, Status, Info),
+    ejabberd_sm:unset_presence(SID, U, S, R, Status),
     {Pres1, State1} = ejabberd_hooks:run_fold(
 			c2s_self_presence, LServer, {Pres, State}, []),
     State2 = broadcast_presence_unavailable(State1, Pres1),
@@ -732,13 +730,11 @@ process_self_presence(#{lserver := LServer} = State,
 process_self_presence(State, _Pres) ->
     State.
 
--spec update_priority(state(), presence()) -> ok.
-update_priority(#{ip := IP, conn := Conn, auth_module := AuthMod,
-		  sid := SID, user := U, server := S, resource := R},
+-spec update_priority(state(), presence()) -> ok | {error, notfound}.
+update_priority(#{sid := SID, user := U, server := S, resource := R},
 		Pres) ->
     Priority = get_priority_from_presence(Pres),
-    Info = [{ip, IP}, {conn, Conn}, {auth_module, AuthMod}],
-    ejabberd_sm:set_presence(SID, U, S, R, Priority, Pres, Info).
+    ejabberd_sm:set_presence(SID, U, S, R, Priority, Pres).
 
 -spec broadcast_presence_unavailable(state(), presence()) -> state().
 broadcast_presence_unavailable(#{jid := JID, pres_a := PresA} = State, Pres) ->
diff --git a/src/ejabberd_sm.erl b/src/ejabberd_sm.erl
index ab1251715..821a7883e 100644
--- a/src/ejabberd_sm.erl
+++ b/src/ejabberd_sm.erl
@@ -47,8 +47,8 @@
 	 disconnect_removed_user/2,
 	 get_user_resources/2,
 	 get_user_present_resources/2,
-	 set_presence/7,
-	 unset_presence/6,
+	 set_presence/6,
+	 unset_presence/5,
 	 close_session_unset_presence/5,
 	 dirty_get_sessions_list/0,
 	 dirty_get_my_sessions_list/0,
@@ -307,26 +307,48 @@ del_user_info(User, Server, Resource, Key) ->
     end.
 
 -spec set_presence(sid(), binary(), binary(), binary(),
-                   prio(), presence(), info()) -> ok.
+                   prio(), presence()) -> ok | {error, notfound}.
 
-set_presence(SID, User, Server, Resource, Priority,
-	     Presence, Info) ->
-    set_session(SID, User, Server, Resource, Priority,
-		Info),
-    ejabberd_hooks:run(set_presence_hook,
-		       jid:nameprep(Server),
-		       [User, Server, Resource, Presence]).
+set_presence(SID, User, Server, Resource, Priority, Presence) ->
+    LUser = jid:nodeprep(User),
+    LServer = jid:nameprep(Server),
+    LResource = jid:resourceprep(Resource),
+    Mod = get_sm_backend(LServer),
+    case get_sessions(Mod, LUser, LServer, LResource) of
+	[] -> {error, notfound};
+	Ss ->
+	    case lists:keyfind(SID, 1, Ss) of
+		#session{info = Info} ->
+		    set_session(SID, User, Server, Resource, Priority, Info),
+		    ejabberd_hooks:run(set_presence_hook,
+				       LServer,
+				       [User, Server, Resource, Presence]);
+		false ->
+		    {error, notfound}
+	    end
+    end.
 
 -spec unset_presence(sid(), binary(), binary(),
-                     binary(), binary(), info()) -> ok.
+                     binary(), binary()) -> ok | {error, notfound}.
 
-unset_presence(SID, User, Server, Resource, Status,
-	       Info) ->
-    set_session(SID, User, Server, Resource, undefined,
-		Info),
-    ejabberd_hooks:run(unset_presence_hook,
-		       jid:nameprep(Server),
-		       [User, Server, Resource, Status]).
+unset_presence(SID, User, Server, Resource, Status) ->
+    LUser = jid:nodeprep(User),
+    LServer = jid:nameprep(Server),
+    LResource = jid:resourceprep(Resource),
+    Mod = get_sm_backend(LServer),
+    case get_sessions(Mod, LUser, LServer, LResource) of
+	[] -> {error, notfound};
+	Ss ->
+	    case lists:keyfind(SID, 1, Ss) of
+		#session{info = Info} ->
+		    set_session(SID, User, Server, Resource, undefined,	Info),
+		    ejabberd_hooks:run(unset_presence_hook,
+				       LServer,
+				       [User, Server, Resource, Status]);
+		false ->
+		    {error, notfound}
+	    end
+    end.
 
 -spec close_session_unset_presence(sid(), binary(), binary(),
                                    binary(), binary()) -> ok.
diff --git a/src/mod_carboncopy.erl b/src/mod_carboncopy.erl
index 4c05f84c3..e1f82e872 100644
--- a/src/mod_carboncopy.erl
+++ b/src/mod_carboncopy.erl
@@ -38,6 +38,7 @@
 	 iq_handler/1, disco_features/5,
 	 is_carbon_copy/1, mod_opt_type/1, depends/2,
 	 mod_options/1]).
+-export([c2s_copy_session/2, c2s_session_opened/1, c2s_session_resumed/1]).
 %% For debugging purposes
 -export([list/2]).
 
@@ -45,6 +46,7 @@
 -include("xmpp.hrl").
 
 -type direction() :: sent | received.
+-type c2s_state() :: ejabberd_c2s:state().
 
 -spec is_carbon_copy(stanza()) -> boolean().
 is_carbon_copy(#message{meta = #{carbon_copy := true}}) ->
@@ -57,6 +59,9 @@ start(Host, _Opts) ->
     %% why priority 89: to define clearly that we must run BEFORE mod_logdb hook (90)
     ejabberd_hooks:add(user_send_packet,Host, ?MODULE, user_send_packet, 89),
     ejabberd_hooks:add(user_receive_packet,Host, ?MODULE, user_receive_packet, 89),
+    ejabberd_hooks:add(c2s_copy_session, Host, ?MODULE, c2s_copy_session, 50),
+    ejabberd_hooks:add(c2s_session_resumed, Host, ?MODULE, c2s_session_resumed, 50),
+    ejabberd_hooks:add(c2s_session_opened, Host, ?MODULE, c2s_session_opened, 50),
     gen_iq_handler:add_iq_handler(ejabberd_sm, Host, ?NS_CARBONS_2, ?MODULE, iq_handler).
 
 stop(Host) ->
@@ -64,7 +69,10 @@ stop(Host) ->
     ejabberd_hooks:delete(disco_local_features, Host, ?MODULE, disco_features, 50),
     %% why priority 89: to define clearly that we must run BEFORE mod_logdb hook (90)
     ejabberd_hooks:delete(user_send_packet,Host, ?MODULE, user_send_packet, 89),
-    ejabberd_hooks:delete(user_receive_packet,Host, ?MODULE, user_receive_packet, 89).
+    ejabberd_hooks:delete(user_receive_packet,Host, ?MODULE, user_receive_packet, 89),
+    ejabberd_hooks:delete(c2s_copy_session, Host, ?MODULE, c2s_copy_session, 50),
+    ejabberd_hooks:delete(c2s_session_resumed, Host, ?MODULE, c2s_session_resumed, 50),
+    ejabberd_hooks:delete(c2s_session_opened, Host, ?MODULE, c2s_session_opened, 50).
 
 reload(_Host, _NewOpts, _OldOpts) ->
     ok.
@@ -123,6 +131,29 @@ user_receive_packet({Packet, #{jid := JID} = C2SState}) ->
 	Pkt -> {Pkt, C2SState}
     end.
 
+-spec c2s_copy_session(c2s_state(), c2s_state()) -> c2s_state().
+c2s_copy_session(State, #{user := U, server := S, resource := R}) ->
+    case ejabberd_sm:get_user_info(U, S, R) of
+	offline -> State;
+	Info ->
+	    case lists:keyfind(carboncopy, 1, Info) of
+		{_, CC} -> State#{carboncopy => CC};
+		false -> State
+	    end
+    end.
+
+-spec c2s_session_resumed(c2s_state()) -> c2s_state().
+c2s_session_resumed(#{user := U, server := S, resource := R,
+		      carboncopy := CC} = State) ->
+    ejabberd_sm:set_user_info(U, S, R, carboncopy, CC),
+    maps:remove(carboncopy, State);
+c2s_session_resumed(State) ->
+    State.
+
+-spec c2s_session_opened(c2s_state()) -> c2s_state().
+c2s_session_opened(State) ->
+    maps:remove(carboncopy, State).
+
 % Modified from original version:
 %    - registered to the user_send_packet hook, to be called only once even for multicast
 %    - do not support "private" message mode, and do not modify the original packet in any way