ssh: make disconnectfun work for both server&client

3 files changed, 84 insertions, 12 deletions

diff --git a/lib/ssh/doc/src/ssh.xml b/lib/ssh/doc/src/ssh.xml
index cf58806aa8..0516945c0e 100644
--- a/lib/ssh/doc/src/ssh.xml
+++ b/lib/ssh/doc/src/ssh.xml
@@ -175,6 +175,11 @@
 	    to use from a security point of view.</p>
 	  </item>
 
+	  <tag><c><![CDATA[{disconnectfun, fun(Reason:term()) -> _}]]></c></tag>
+	  <item>
+	    <p>Provides a fun to implement your own logging when a server disconnects the client.</p>
+	  </item>
+
 	  <tag><c><![CDATA[{public_key_alg, 'ssh-rsa' | 'ssh-dss'}]]></c></tag>
 	  <item>
 	    <note>
diff --git a/lib/ssh/src/ssh_connection_handler.erl b/lib/ssh/src/ssh_connection_handler.erl
index 3bdca4ba94..ab1fc93a1b 100644
--- a/lib/ssh/src/ssh_connection_handler.erl
+++ b/lib/ssh/src/ssh_connection_handler.erl
@@ -750,15 +750,12 @@ handle_sync_event({info, ChannelPid}, _From, StateName,
     {reply, {ok, Result}, StateName, State};
 
 handle_sync_event(stop, _, _StateName, #state{connection_state = Connection0,
-					     role = Role,
-					     opts = Opts} = State0) ->
-    {disconnect, Reason, {{replies, Replies}, Connection}} =
+					     role = Role} = State0) ->
+    {disconnect, _Reason, {{replies, Replies}, Connection}} =
 	ssh_connection:handle_msg(#ssh_msg_disconnect{code = ?SSH_DISCONNECT_BY_APPLICATION,
 						      description = "User closed down connection",
 						      language = "en"}, Connection0, Role),
     State = send_replies(Replies, State0),
-    SSHOpts = proplists:get_value(ssh_opts, Opts),
-    disconnect_fun(Reason, SSHOpts),
     {stop, normal, ok, State#state{connection_state = Connection}};
 
 
@@ -1275,7 +1272,6 @@ generate_event(<<?BYTE(Byte), _/binary>> = Msg, StateName,
 	       #state{
 		  role = Role,
 		  starter = User,
-		  opts = Opts,
 		  renegotiate = Renegotiation,
 		  connection_state = Connection0} = State0, EncData)
   when  Byte == ?SSH_MSG_GLOBAL_REQUEST;
@@ -1315,21 +1311,17 @@ generate_event(<<?BYTE(Byte), _/binary>> = Msg, StateName,
 	    User ! {self(), not_connected, Reason},
 	    {stop, {shutdown, normal},
 	     next_packet(State#state{connection_state = Connection})};
-	{disconnect, Reason, {{replies, Replies}, Connection}} ->
+	{disconnect, _Reason, {{replies, Replies}, Connection}} ->
 	    State = send_replies(Replies,  State1#state{connection_state = Connection}),
-	    SSHOpts = proplists:get_value(ssh_opts, Opts),
-	    disconnect_fun(Reason, SSHOpts),
 	    {stop, {shutdown, normal}, State#state{connection_state = Connection}}
     catch
 	_:Error ->
-	    {disconnect, Reason, {{replies, Replies}, Connection}} =
+	    {disconnect, _Reason, {{replies, Replies}, Connection}} =
 		ssh_connection:handle_msg(
 		  #ssh_msg_disconnect{code = ?SSH_DISCONNECT_BY_APPLICATION,
 					  description = "Internal error",
 				      language = "en"}, Connection0, Role),
 	    State = send_replies(Replies,  State1#state{connection_state = Connection}),
-	    SSHOpts = proplists:get_value(ssh_opts, Opts),
-	    disconnect_fun(Reason, SSHOpts),
 	    {stop, {shutdown, Error}, State#state{connection_state = Connection}}
     end;
 
@@ -1576,12 +1568,14 @@ handle_disconnect(#ssh_msg_disconnect{} = DisconnectMsg, State, Error) ->
 handle_disconnect(Type,  #ssh_msg_disconnect{description = Desc} = Msg, #state{connection_state = Connection0,									role = Role} = State0) ->
     {disconnect, _, {{replies, Replies}, Connection}} = ssh_connection:handle_msg(Msg, Connection0, Role),
     State = send_replies(disconnect_replies(Type, Msg, Replies), State0),
+    disconnect_fun(Desc, State#state.opts),
     {stop, {shutdown, Desc}, State#state{connection_state = Connection}}.
 
 handle_disconnect(Type, #ssh_msg_disconnect{description = Desc} = Msg, #state{connection_state = Connection0,
 									      role = Role} = State0, ErrorMsg) ->
     {disconnect, _, {{replies, Replies}, Connection}} = ssh_connection:handle_msg(Msg, Connection0, Role),
     State = send_replies(disconnect_replies(Type, Msg, Replies), State0),
+    disconnect_fun(Desc, State#state.opts),
     {stop, {shutdown, {Desc, ErrorMsg}}, State#state{connection_state = Connection}}.
 
 disconnect_replies(own, Msg, Replies) ->
@@ -1700,6 +1694,8 @@ send_reply({flow_control, Cache, Channel, From, Msg}) ->
 send_reply({flow_control, From, Msg}) ->
     gen_fsm:reply(From, Msg).
 
+disconnect_fun({disconnect,Msg}, Opts) ->
+    disconnect_fun(Msg, Opts);
 disconnect_fun(_, undefined) ->
     ok;
 disconnect_fun(Reason, Opts) ->
diff --git a/lib/ssh/test/ssh_basic_SUITE.erl b/lib/ssh/test/ssh_basic_SUITE.erl
index cff695681e..4b72c57cca 100644
--- a/lib/ssh/test/ssh_basic_SUITE.erl
+++ b/lib/ssh/test/ssh_basic_SUITE.erl
@@ -56,6 +56,8 @@ all() ->
      ssh_daemon_minimal_remote_max_packet_size_option,
      ssh_msg_debug_fun_option_client,
      ssh_msg_debug_fun_option_server,
+     disconnectfun_option_server,
+     disconnectfun_option_client,
      preferred_algorithms,
      id_string_no_opt_client,
      id_string_own_string_client,
@@ -738,6 +740,75 @@ ssh_msg_debug_fun_option_server(Config) ->
     end.
 
 %%--------------------------------------------------------------------
+disconnectfun_option_server(Config) ->
+    PrivDir = ?config(priv_dir, Config),
+    UserDir = filename:join(PrivDir, nopubkey), % to make sure we don't use public-key-auth
+    file:make_dir(UserDir),
+    SysDir = ?config(data_dir, Config),
+
+    Parent = self(),
+    DisConnFun = fun(Reason) -> Parent ! {disconnect,Reason} end,
+
+    {Pid, Host, Port} = ssh_test_lib:daemon([{system_dir, SysDir},
+					     {user_dir, UserDir},
+					     {password, "morot"},
+					     {failfun, fun ssh_test_lib:failfun/2},
+					     {disconnectfun, DisConnFun}]),
+    ConnectionRef =
+	ssh_test_lib:connect(Host, Port, [{silently_accept_hosts, true},
+					  {user, "foo"},
+					  {password, "morot"},
+					  {user_dir, UserDir},
+					  {user_interaction, false}]),
+    ssh:close(ConnectionRef),
+    receive
+	{disconnect,Reason} ->
+	    ct:log("Server detected disconnect: ~p",[Reason]),
+	    ssh:stop_daemon(Pid),
+	    ok
+    after 3000 ->
+	    receive
+		X -> ct:log("received ~p",[X])
+	    after 0 -> ok
+	    end,
+	    {fail,"Timeout waiting for disconnect"}
+    end.
+
+%%--------------------------------------------------------------------
+disconnectfun_option_client(Config) ->
+    PrivDir = ?config(priv_dir, Config),
+    UserDir = filename:join(PrivDir, nopubkey), % to make sure we don't use public-key-auth
+    file:make_dir(UserDir),
+    SysDir = ?config(data_dir, Config),
+
+    Parent = self(),
+    DisConnFun = fun(Reason) -> Parent ! {disconnect,Reason} end,
+
+    {Pid, Host, Port} = ssh_test_lib:daemon([{system_dir, SysDir},
+					     {user_dir, UserDir},
+					     {password, "morot"},
+					     {failfun, fun ssh_test_lib:failfun/2}]),
+    _ConnectionRef =
+	ssh_test_lib:connect(Host, Port, [{silently_accept_hosts, true},
+					  {user, "foo"},
+					  {password, "morot"},
+					  {user_dir, UserDir},
+					  {user_interaction, false},
+					  {disconnectfun, DisConnFun}]),
+    ssh:stop_daemon(Pid),
+    receive
+	{disconnect,Reason} ->
+	    ct:log("Client detected disconnect: ~p",[Reason]),
+	    ok
+    after 3000 ->
+	    receive
+		X -> ct:log("received ~p",[X])
+	    after 0 -> ok
+	    end,
+	    {fail,"Timeout waiting for disconnect"}
+    end.
+
+%%--------------------------------------------------------------------
 known_hosts() ->
     [{doc, "check that known_hosts is updated correctly"}].
 known_hosts(Config) when is_list(Config) ->