From c290b4284f86487958d760417261c2ec473f2989 Mon Sep 17 00:00:00 2001
From: Evgeniy Khramtsov <ekhramtsov@process-one.net>
Date: Sat, 15 Apr 2017 11:15:50 +0300
Subject: [PATCH] Fix closing of outbound S2S connections

---
 src/mod_stream_mgmt.erl |  3 ++-
 src/xmpp_stream_in.erl  | 31 +++++++++++++++----------------
 src/xmpp_stream_out.erl | 31 +++++++++++++++----------------
 3 files changed, 32 insertions(+), 33 deletions(-)

diff --git a/src/mod_stream_mgmt.erl b/src/mod_stream_mgmt.erl
index 97875fa81..0f64a47cc 100644
--- a/src/mod_stream_mgmt.erl
+++ b/src/mod_stream_mgmt.erl
@@ -233,7 +233,8 @@ c2s_handle_info(#{mgmt_ack_timer := TRef, jid := JID, mod := Mod} = State,
 		{timeout, TRef, ack_timeout}) ->
     ?DEBUG("Timed out waiting for stream management acknowledgement of ~s",
 	   [jid:encode(JID)]),
-    {stop, Mod:close(State, ack_timeout)};
+    State1 = Mod:close(State),
+    {stop, transition_to_pending(State1)};
 c2s_handle_info(#{mgmt_state := pending, jid := JID, mod := Mod} = State,
 		{timeout, _, pending_timeout}) ->
     ?DEBUG("Timed out waiting for resumption of stream for ~s",
diff --git a/src/xmpp_stream_in.erl b/src/xmpp_stream_in.erl
index 073ae3d1f..253adbf95 100644
--- a/src/xmpp_stream_in.erl
+++ b/src/xmpp_stream_in.erl
@@ -151,23 +151,17 @@ send(_, _) ->
 
 -spec close(pid()) -> ok;
 	   (state()) -> state().
-close(Ref) ->
-    close(Ref, closed).
-
--spec close(pid(), atom()) -> ok;
-	   (state(), atom()) -> state().
-close(Pid, Reason) when is_pid(Pid) ->
-    cast(Pid, {close, Reason});
-close(#{owner := Owner} = State, Reason) when Owner == self() ->
-    case is_disconnected(State) of
-	true -> State;
-	false ->
-	    _IgnoreState = close_socket(State),
-	    process_stream_end({socket, Reason}, State)
-    end;
-close(_, _) ->
+close(Pid) when is_pid(Pid) ->
+    close(Pid, closed);
+close(#{owner := Owner} = State) when Owner == self() ->
+    close_socket(State);
+close(_) ->
     erlang:error(badarg).
 
+-spec close(pid(), atom()) -> ok.
+close(Pid, Reason) ->
+    cast(Pid, {close, Reason}).
+
 -spec establish(state()) -> state().
 establish(State) ->
     process_stream_established(State).
@@ -275,7 +269,12 @@ handle_cast({send, Pkt}, State) ->
 handle_cast(stop, State) ->
     {stop, normal, State};
 handle_cast({close, Reason}, State) ->
-    noreply(close(State, Reason));
+    State1 = close_socket(State),
+    noreply(
+      case is_disconnected(State) of
+	  true -> State1;
+	  false -> process_stream_end({socket, Reason}, State)
+      end);
 handle_cast(Cast, #{mod := Mod} = State) ->
     noreply(try Mod:handle_cast(Cast, State)
 	      catch _:undef -> State
diff --git a/src/xmpp_stream_out.erl b/src/xmpp_stream_out.erl
index c8a2fba1e..af5c67c66 100644
--- a/src/xmpp_stream_out.erl
+++ b/src/xmpp_stream_out.erl
@@ -161,23 +161,17 @@ send(_, _) ->
 
 -spec close(pid()) -> ok;
 	   (state()) -> state().
-close(Ref) ->
-    close(Ref, closed).
-
--spec close(pid(), atom()) -> ok;
-	   (state(), atom()) -> state().
-close(Pid, Reason) when is_pid(Pid) ->
-    cast(Pid, {close, Reason});
-close(#{owner := Owner} = State, Reason) when Owner == self() ->
-    case is_disconnected(State) of
-	true -> State;
-	false ->
-	    _IgnoreState = close_socket(State),
-	    process_stream_end({socket, Reason}, State)
-    end;
-close(_, _) ->
+close(Pid) when is_pid(Pid) ->
+    close(Pid, closed);
+close(#{owner := Owner} = State) when Owner == self() ->
+    close_socket(State);
+close(_) ->
     erlang:error(badarg).
 
+-spec close(pid(), atom()) -> ok.
+close(Pid, Reason) ->
+    cast(Pid, {close, Reason}).
+
 -spec establish(state()) -> state().
 establish(State) ->
     process_stream_established(State).
@@ -306,7 +300,12 @@ handle_cast({send, Pkt}, State) ->
 handle_cast(stop, State) ->
     {stop, normal, State};
 handle_cast({close, Reason}, State) ->
-    noreply(close(State, Reason));
+    State1 = close_socket(State),
+    noreply(
+      case is_disconnected(State) of
+	  true -> State1;
+	  false -> process_stream_end({socket, Reason}, State)
+      end);
 handle_cast(Cast, #{mod := Mod} = State) ->
     noreply(try Mod:handle_cast(Cast, State)
 	    catch _:undef -> State