ssh: extend connectfun with connnection info

alexandrejbr · Dec 6, 2024 · 056acac · 056acac
1 parent 50c0886
commit 056acac
Show file tree

Hide file tree

Showing 5 changed files with 65 additions and 4 deletions.
diff --git a/lib/ssh/src/ssh.hrl b/lib/ssh/src/ssh.hrl
@@ -1147,7 +1147,8 @@ in the User's Guide chapter.
 -doc(#{title => <<"Daemon Options">>}).
 -type callbacks_daemon_options() ::
         {failfun, fun((User::string(), PeerAddress::inet:ip_address(), Reason::term()) -> _)}
-      | {connectfun, fun((User::string(), PeerAddress::inet:ip_address(), Method::string()) ->_)} .
+      | {connectfun, fun((User::string(), PeerAddress::inet:ip_address(), Method::string()) ->_)
+        | fun((User::string(), PeerAddress::inet:ip_address(), Method::string(), Info::proplists:proplist()) ->_)} .
 
 -doc(#{title => <<"Other data types">>}).
 -type opaque_daemon_options()  ::

diff --git a/lib/ssh/src/ssh_connection_handler.erl b/lib/ssh/src/ssh_connection_handler.erl
@@ -68,6 +68,7 @@
          retrieve/2,
 	 info/1, info/2,
 	 connection_info/2,
+	 connection_info_server/1,
 	 channel_info/3,
 	 adjust_window/3, close/2,
 	 disconnect/4,
@@ -308,6 +309,17 @@ connection_info(ConnectionHandler, Key) when is_atom(Key) ->
 connection_info(ConnectionHandler, Options) ->
     call(ConnectionHandler, {connection_info, Options}).
 
+%%--------------------------------------------------------------------
+connection_info_server(D) when is_tuple(D) ->
+    Keys = [client_version,
+            server_version,
+            peer,
+            sockname,
+            options,
+            algorithms
+           ],
+    fold_keys(Keys, fun conn_info/2, D).
+
 %%--------------------------------------------------------------------
 -spec channel_info(connection_ref(),
 		   channel_id(),

diff --git a/lib/ssh/src/ssh_fsm_userauth_server.erl b/lib/ssh/src/ssh_fsm_userauth_server.erl
@@ -184,8 +184,14 @@ set_max_initial_idle_timeout(#data{ssh_params = #ssh{opts=Opts}}) ->
     {{timeout,max_initial_idle_time}, ?GET_OPT(max_initial_idle_time,Opts), none}.
 
 connected_fun(User, Method, #data{ssh_params = #ssh{peer = {_,Peer}}} = D) ->
-    ?CALL_FUN(connectfun,D)(User, Peer, Method).
-
+    Fun = ?GET_OPT(connectfun, (D#data.ssh_params)#ssh.opts),
+    ConnInfo = ssh_connection_handler:connection_info_server(D),
+    case erlang:fun_info(Fun, arity) of
+        {arity, 3} ->
+            Fun(User, Peer, Method);
+        {arity, 4} ->
+            Fun(User, Peer, Method, ConnInfo)
+    end.
 
 retry_fun(_, undefined, _) ->
     ok;

diff --git a/lib/ssh/src/ssh_options.erl b/lib/ssh/src/ssh_options.erl
@@ -584,7 +584,9 @@ default(server) ->
 
       connectfun =>
           #{default => fun(_,_,_) -> void end,
-            chk => fun(V) -> check_function3(V) end,
+            chk => fun(V) -> check_function3(V) orelse
+                                 check_function4(V) %% Adds ssh connection info
+                   end,
             class => user_option
            },
 

diff --git a/lib/ssh/test/ssh_options_SUITE.erl b/lib/ssh/test/ssh_options_SUITE.erl
@@ -74,6 +74,7 @@
 	 user_dir_option/1,
 	 user_dir_fun_option/1,
 	 connectfun_disconnectfun_server/1,
+	 connectfun4_server/1,
 	 hostkey_fingerprint_check/1,
 	 hostkey_fingerprint_check_md5/1,
 	 hostkey_fingerprint_check_sha/1,
@@ -114,6 +115,7 @@ suite() ->
 
 all() -> 
     [connectfun_disconnectfun_server,
+     connectfun4_server,
      connectfun_disconnectfun_client,
      server_password_option,
      server_userpassword_option,
@@ -778,6 +780,44 @@ connectfun_disconnectfun_server(Config) ->
 	    {fail, "No connectfun action"}
     end.
 
+
+%%--------------------------------------------------------------------
+connectfun4_server(Config) ->
+    UserDir = proplists:get_value(user_dir, Config),
+    SysDir = proplists:get_value(data_dir, Config),
+
+    Parent = self(),
+    Ref = make_ref(),
+    ConnFun = fun(User,_,Method,ConnInfo) -> Parent ! {connect,Ref,User,Method,ConnInfo} end,
+    DiscFun = fun(R) -> Parent ! {disconnect,Ref,R} end,
+
+    {Pid, Host, Port} = ssh_test_lib:daemon([{system_dir, SysDir},
+					     {user_dir, UserDir},
+					     {password, "morot"},
+					     {failfun, fun ssh_test_lib:failfun/2},
+					     {connectfun, ConnFun}]),
+    ConnectionRef =
+	ssh_test_lib:connect(Host, Port, [{silently_accept_hosts, true},
+					  {user, "foo"},
+					  {password, "morot"},
+					  {user_dir, UserDir},
+					  {user_interaction, false}]),
+    receive
+	{connect,Ref,User,Method,ConnInfo} ->
+            "foo" = User,
+            "keyboard-interactive" = Method,
+            Keys = [client_version, server_version, peer, sockname, options, algorithms],
+            true = lists:all(fun({K, _}) -> lists:member(K, Keys) end, ConnInfo),
+	    ssh:close(ConnectionRef),
+            ssh:stop_daemon(Pid)
+    after 10000 ->
+	    receive
+		X -> ct:log("received ~p",[X])
+	    after 0 -> ok
+	    end,
+	    {fail, "No connectfun action"}
+    end.
+
 %%--------------------------------------------------------------------
 connectfun_disconnectfun_client(Config) ->
     UserDir = proplists:get_value(user_dir, Config),