Add support for s2s bidi

mix.exs

@@ -144,7 +144,7 @@ defmodule Ejabberd.MixProject do
      {:p1_utils, "~> 1.0"},
      {:pkix, "~> 1.0"},
      {:stringprep, ">= 1.0.26"},
-     {:xmpp, git: "https://github.com/processone/xmpp.git", ref: "c045d4d8555e251f2212743db8af90255da2ab57", override: true},
+     {:xmpp, git: "https://github.com/processone/xmpp.git", ref: "85d4cb6c080f6328a174bdc12103fd65834a400b", override: true},
      {:yconf, "~> 1.0"}]
     ++ cond_deps()
   end

mix.lock

@@ -32,6 +32,6 @@
   "stringprep": {:hex, :stringprep, "1.0.30", "46cf0ff631b3e7328f61f20b454d59428d87738f25d709798b5dcbb9b83c23f1", [:rebar3], [{:p1_utils, "1.0.26", [hex: :p1_utils, repo: "hexpm", optional: false]}], "hexpm", "f6fc9b3384a03877830f89b2f38580caf3f4a27448a4a333d6a8c3975c220b9a"},
   "stun": {:hex, :stun, "1.2.14", "6f538ac80c842131dbd149055570d116bfabc9b5ebff4bd6af2e7888958c660c", [:rebar3], [{:fast_tls, "1.1.21", [hex: :fast_tls, repo: "hexpm", optional: false]}, {:p1_utils, "1.0.26", [hex: :p1_utils, repo: "hexpm", optional: false]}], "hexpm", "e134807b1b7a8dffd94e64eefee00e65c7b4042f3d14e16f8f43566d20371583"},
   "unicode_util_compat": {:hex, :unicode_util_compat, "0.7.0", "bc84380c9ab48177092f43ac89e4dfa2c6d62b40b8bd132b1059ecc7232f9a78", [:rebar3], [], "hexpm", "25eee6d67df61960cf6a794239566599b09e17e668d3700247bc498638152521"},
-  "xmpp": {:git, "https://github.com/processone/xmpp.git", "c045d4d8555e251f2212743db8af90255da2ab57", [ref: "c045d4d8555e251f2212743db8af90255da2ab57"]},
+  "xmpp": {:git, "https://github.com/processone/xmpp.git", "85d4cb6c080f6328a174bdc12103fd65834a400b", [ref: "85d4cb6c080f6328a174bdc12103fd65834a400b"]},
   "yconf": {:hex, :yconf, "1.0.16", "d59521d66ff89f219411b6e9277cd6feec7cc6fce11554e67de02a8d0a470479", [:rebar3], [{:fast_yaml, "1.0.37", [hex: :fast_yaml, repo: "hexpm", optional: false]}], "hexpm", "e947813273f38711c7b2e5a8e4acc9a51c7bbe854f744a345f60300b38586c89"},
 }

rebar.config

@@ -69,7 +69,7 @@
         {stringprep, "~> 1.0.29", {git, "https://github.com/processone/stringprep", {tag, "1.0.30"}}},
         {if_var_true, stun,
          {stun, "~> 1.2.12", {git, "https://github.com/processone/stun", {tag, "1.2.14"}}}},
-        {xmpp, "~> 1.8.3", {git, "https://github.com/processone/xmpp", "c045d4d8555e251f2212743db8af90255da2ab57"}},
+        {xmpp, "~> 1.8.3", {git, "https://github.com/processone/xmpp", "85d4cb6c080f6328a174bdc12103fd65834a400b"}},
         {yconf, "~> 1.0.15", {git, "https://github.com/processone/yconf", {tag, "1.0.16"}}}
        ]}.
 

rebar.lock

@@ -32,7 +32,7 @@
  {<<"unicode_util_compat">>,{pkg,<<"unicode_util_compat">>,<<"0.7.0">>},1},
  {<<"xmpp">>,
   {git,"https://github.com/processone/xmpp",
-       {ref,"c045d4d8555e251f2212743db8af90255da2ab57"}},
+       {ref,"85d4cb6c080f6328a174bdc12103fd65834a400b"}},
   0},
  {<<"yconf">>,{pkg,<<"yconf">>,<<"1.0.16">>},0}]}.
 [

src/ejabberd_s2s.erl

@@ -43,7 +43,7 @@
 	 external_host_overloaded/1, is_temporarly_blocked/1,
 	 get_commands_spec/0, zlib_enabled/1, get_idle_timeout/1,
 	 tls_required/1, tls_enabled/1, tls_options/3,
-	 host_up/1, host_down/1, queue_type/1]).
+	 host_up/1, host_down/1, queue_type/1, register_connection/1]).
 
 %% gen_server callbacks
 -export([init/1, handle_call/3, handle_cast/2,
@@ -130,6 +130,10 @@ get_connections_pids(FromTo) ->
 	    []
     end.
 
+-spec register_connection(FromTo :: {binary(), binary()}) -> ok.
+register_connection(FromTo) ->
+	gen_server:call(ejabberd_s2s, {register_connection, FromTo, self()}).
+
 -spec dirty_get_connections() -> [{binary(), binary()}].
 dirty_get_connections() ->
     mnesia:dirty_all_keys(s2s).
@@ -228,6 +232,8 @@ init([]) ->
 
 handle_call({new_connection, Args}, _From, State) ->
     {reply, erlang:apply(fun new_connection_int/7, Args), State};
+handle_call({register_connection, FromTo, Pid}, _From, State) ->
+    {reply, register_connection_int(FromTo, Pid), State};
 handle_call(Request, From, State) ->
     ?WARNING_MSG("Unexpected call from ~p: ~p", [From, Request]),
     {noreply, State}.
@@ -478,6 +484,20 @@ new_connection_int(MyServer, Server, From, FromTo,
 	    []
     end.
 
+-spec register_connection_int(FromTo :: {binary(), binary()}, Pid :: pid()) -> ok.
+register_connection_int(FromTo, Pid) ->
+    F = fun() ->
+			mnesia:write(#s2s{fromto = FromTo, pid = Pid})
+		end,
+    TRes = mnesia:transaction(F),
+    case TRes of
+	{atomic, _} ->
+		erlang:monitor(process, Pid),
+		ok;
+	_ ->
+		ok
+	end.
+
 -spec max_s2s_connections_number({binary(), binary()}) -> pos_integer().
 max_s2s_connections_number({From, To}) ->
     case ejabberd_shaper:match(From, max_s2s_connections, jid:make(To)) of

src/ejabberd_s2s_out.erl

@@ -34,7 +34,7 @@
 	 terminate/2, code_change/3]).
 %% Hooks
 -export([process_auth_result/2, process_closed/2, handle_unexpected_info/2,
-	 handle_unexpected_cast/2, process_downgraded/2]).
+	 handle_unexpected_cast/2, process_downgraded/2, handle_unauthenticated_features/2]).
 %% API
 -export([start/3, start_link/3, connect/1, close/1, close/2, stop_async/1, send/2,
 	 route/2, establish/1, update_state/2, host_up/1, host_down/1]).
@@ -216,6 +216,9 @@ dns_retries(#{server_host := ServerHost}) ->
 dns_timeout(#{server_host := ServerHost}) ->
     ejabberd_option:s2s_dns_timeout(ServerHost).
 
+handle_unauthenticated_features(Features, #{server_host := ServerHost} = State) ->
+    ejabberd_hooks:run_fold(s2s_out_unauthenticated_features, ServerHost, State, [Features]).
+
 handle_auth_success(Mech, #{socket := Socket, ip := IP,
 			    remote_server := RServer,
 			    server_host := ServerHost,

src/mod_s2s_bidi.erl

@@ -0,0 +1,142 @@
+%%%-------------------------------------------------------------------
+%%% Created : 20 Oct 2024 by Pawel Chmielowski <pawel@process-one.net>
+%%%
+%%%
+%%% ejabberd, Copyright (C) 2002-2024   ProcessOne
+%%%
+%%% This program is free software; you can redistribute it and/or
+%%% modify it under the terms of the GNU General Public License as
+%%% published by the Free Software Foundation; either version 2 of the
+%%% License, or (at your option) any later version.
+%%%
+%%% This program is distributed in the hope that it will be useful,
+%%% but WITHOUT ANY WARRANTY; without even the implied warranty of
+%%% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+%%% General Public License for more details.
+%%%
+%%% You should have received a copy of the GNU General Public License along
+%%% with this program; if not, write to the Free Software Foundation, Inc.,
+%%% 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+%%%
+%%%-------------------------------------------------------------------
+-module(mod_s2s_bidi).
+-behaviour(gen_mod).
+-protocol({xep, 288, '1.0.1'}).
+
+%% gen_mod API
+-export([start/2, stop/1, reload/3, depends/2, mod_options/1]).
+-export([mod_doc/0]).
+%% Hooks
+-export([s2s_in_packet/2, s2s_out_packet/2,
+	 s2s_in_features/2, s2s_in_auth_result/3, s2s_out_unauthenticated_features/2, s2s_in_handle_info/2]).
+
+-include_lib("xmpp/include/xmpp.hrl").
+-include("logger.hrl").
+-include("translate.hrl").
+
+%%%===================================================================
+%%% API
+%%%===================================================================
+start(_Host, _Opts) ->
+    {ok, [{hook, s2s_in_pre_auth_features, s2s_in_features, 50},
+	  {hook, s2s_in_post_auth_features, s2s_in_features, 50},
+	  {hook, s2s_in_unauthenticated_packet, s2s_in_packet, 50},
+	  {hook, s2s_in_authenticated_packet, s2s_in_packet, 50},
+	  {hook, s2s_in_handle_info, s2s_in_handle_info, 50},
+	  {hook, s2s_in_auth_result, s2s_in_auth_result, 50},
+	  {hook, s2s_out_unauthenticated_features, s2s_out_unauthenticated_features, 50},
+	  {hook, s2s_out_packet, s2s_out_packet, 50}]}.
+
+stop(_Host) ->
+    ok.
+
+reload(_Host, _NewOpts, _OldOpts) ->
+    ok.
+
+depends(_Host, _Opts) ->
+    [].
+
+mod_options(_Host) ->
+    [].
+
+mod_doc() ->
+    #{desc =>
+      [?T("The module adds support for "
+	  "https://xmpp.org/extensions/xep-0288.html"
+	  "[XEP-0288: Bidirectional Server-to-Server Connections] that allows using "
+	  "single s2s connection to communicate in both directions.")],
+      opts => [],
+      example =>
+      ["modules:",
+       "  mod_s2s_bidi: {}"]}.
+
+s2s_in_features(Acc, _) ->
+    [#s2s_bidi{}|Acc].
+
+s2s_in_packet(State, #s2s_bidi{}) ->
+    {stop, State#{bidi_enabled => true}};
+s2s_in_packet(State, _) ->
+    State.
+
+s2s_out_unauthenticated_features(#{db_verify := _} = State, _) ->
+    State;
+s2s_out_unauthenticated_features(State, #stream_features{} = Pkt) ->
+    try xmpp:try_subtag(Pkt, #s2s_bidi{}) of
+	#s2s_bidi{} ->
+	    ejabberd_s2s_out:send(State#{bidi_enabled => true}, #s2s_bidi{})
+    catch _:{xmpp_codec, _Why} ->
+	State
+    end.
+
+s2s_out_packet(#{bidi_enabled := true, ip := {IP, _}} = State, Pkt0)
+    when ?is_stanza(Pkt0) ->
+    To = xmpp:get_to(Pkt0),
+    case check_from_to(State, xmpp:get_from(Pkt0), To) of
+	ok ->
+	    Pkt = xmpp:put_meta(Pkt0, ip, IP),
+	    LServer = ejabberd_router:host_of_route(To#jid.lserver),
+	    State1 = ejabberd_hooks:run_fold(s2s_in_authenticated_packet,
+					     LServer, State, [Pkt]),
+	    {Pkt1, State2} = ejabberd_hooks:run_fold(s2s_receive_packet, LServer,
+						     {Pkt, State1}, []),
+	    case Pkt1 of
+		drop -> ok;
+		_ -> ejabberd_router:route(Pkt1)
+	    end,
+	    {stop, State2};
+	{error, Err} ->
+	    {stop, ejabberd_s2s_out:send(State, Err)}
+    end;
+s2s_out_packet(#{db_verify := _} = State, #stream_features{}) ->
+    State;
+s2s_out_packet(State, #stream_features{} = Pkt) ->
+    try xmpp:try_subtag(Pkt, #s2s_bidi{}) of
+	#s2s_bidi{} ->
+	    ejabberd_s2s_out:send(State#{bidi_enabled => true}, #s2s_bidi{})
+    catch _:{xmpp_codec, _Why} ->
+	State
+    end;
+s2s_out_packet(State, _Pkt) ->
+    State.
+
+s2s_in_handle_info(State, {route, Pkt}) when ?is_stanza(Pkt) ->
+    ejabberd_s2s_in:send(State, Pkt);
+s2s_in_handle_info(State, _Info) ->
+    State.
+
+check_from_to(#{remote_server := RServer}, #jid{lserver = FromServer},
+	      #jid{lserver = ToServer}) ->
+    if
+	RServer /= FromServer -> {error, xmpp:serr_invalid_from()};
+	true ->
+	    case ejabberd_router:is_my_route(ToServer) of
+		false -> {error, xmpp:serr_host_unknown()};
+		_ -> ok
+	    end
+    end.
+
+s2s_in_auth_result(#{server := LServer, bidi_enabled := true} = State, true, RServer) ->
+    ejabberd_s2s:register_connection({LServer, RServer}),
+    State;
+s2s_in_auth_result(State, _, _) ->
+    State.