ssh: Simplification of using fun:s as exec subsystems

3 files changed, 255 insertions, 127 deletions

diff --git a/lib/ssh/src/ssh_cli.erl b/lib/ssh/src/ssh_cli.erl
index 958c342f5f..783f2f80c0 100644
--- a/lib/ssh/src/ssh_cli.erl
+++ b/lib/ssh/src/ssh_cli.erl
@@ -118,42 +118,53 @@ handle_ssh_msg({ssh_cm, ConnectionHandler,
     write_chars(ConnectionHandler, ChannelId, Chars),
     {ok, State#state{pty = Pty, buf = NewBuf}};
 
-handle_ssh_msg({ssh_cm, ConnectionHandler,
-	    {shell, ChannelId, WantReply}}, State) ->
+handle_ssh_msg({ssh_cm, ConnectionHandler,  {shell, ChannelId, WantReply}}, State) ->
     NewState = start_shell(ConnectionHandler, State),
-    ssh_connection:reply_request(ConnectionHandler, WantReply,
-				 success, ChannelId),
-    {ok, NewState#state{channel = ChannelId,
-			cm = ConnectionHandler}};
-
-handle_ssh_msg({ssh_cm, ConnectionHandler,
-		{exec, ChannelId, WantReply, Cmd}}, #state{exec=undefined,
-                                                           shell=?DEFAULT_SHELL} = State) ->
-    {Reply, Status} = exec(Cmd),
-    write_chars(ConnectionHandler,
-		ChannelId, io_lib:format("~p\n", [Reply])),
-    ssh_connection:reply_request(ConnectionHandler, WantReply,
-				 success, ChannelId),
-    ssh_connection:exit_status(ConnectionHandler, ChannelId, Status),
-    ssh_connection:send_eof(ConnectionHandler, ChannelId),
-    {stop, ChannelId, State#state{channel = ChannelId, cm = ConnectionHandler}};
-
-handle_ssh_msg({ssh_cm, ConnectionHandler,
-		{exec, ChannelId, WantReply, _Cmd}}, #state{exec = undefined} = State) ->
-    write_chars(ConnectionHandler, ChannelId, 1, "Prohibited.\n"),
     ssh_connection:reply_request(ConnectionHandler, WantReply, success, ChannelId),
-    ssh_connection:exit_status(ConnectionHandler, ChannelId, 255),
-    ssh_connection:send_eof(ConnectionHandler, ChannelId),
-    {stop, ChannelId, State#state{channel = ChannelId, cm = ConnectionHandler}};
-
-handle_ssh_msg({ssh_cm, ConnectionHandler,
-		{exec, ChannelId, WantReply, Cmd}}, State) ->
-    NewState = start_shell(ConnectionHandler, Cmd, State),
-    ssh_connection:reply_request(ConnectionHandler, WantReply,
-				 success, ChannelId),
     {ok, NewState#state{channel = ChannelId,
 			cm = ConnectionHandler}};
 
+handle_ssh_msg({ssh_cm, ConnectionHandler,  {exec, ChannelId, WantReply, Cmd}}, S0) ->
+    case
+        case S0#state.exec of
+            {direct,F} ->
+                %% Exec called and a Fun or MFA is defined to use.  The F returns the
+                %% value to return.
+                exec_direct(ConnectionHandler, F, Cmd);
+
+            undefined when S0#state.shell == ?DEFAULT_SHELL ->
+                %% Exec called and the shell is the default shell (= Erlang shell).
+                %% To be exact, eval the term as an Erlang term (but not using the
+                %% ?DEFAULT_SHELL directly). This disables banner, prompts and such.
+                exec_in_erlang_default_shell(Cmd);
+
+            undefined ->
+                %% Exec called, but the a shell other than the default shell is defined.
+                %% No new exec shell is defined, so don't execute!
+                %% We don't know if it is intended to use the new shell or not.
+                {"Prohibited.", 255, 1};
+
+            _ ->
+                %% Exec called and a Fun or MFA is defined to use.  The F communicates via
+                %% standard io:write/read.
+                %% Kept for compatibility.
+                S1 = start_exec_shell(ConnectionHandler, Cmd, S0),
+                ssh_connection:reply_request(ConnectionHandler, WantReply, success, ChannelId),
+                {ok, S1}
+        end
+    of
+        {Reply, Status, Type} ->
+            write_chars(ConnectionHandler, ChannelId, Type, Reply),
+            ssh_connection:reply_request(ConnectionHandler, WantReply, success, ChannelId),
+            ssh_connection:exit_status(ConnectionHandler, ChannelId, Status),
+            ssh_connection:send_eof(ConnectionHandler, ChannelId),
+            {stop, ChannelId, S0#state{channel = ChannelId, cm = ConnectionHandler}};
+            
+        {ok, S} ->
+            {ok, S#state{channel = ChannelId,
+                         cm = ConnectionHandler}}
+    end;
+
 handle_ssh_msg({ssh_cm, _ConnectionHandler, {eof, _ChannelId}}, State) ->
     {ok, State};
 
@@ -259,35 +270,7 @@ to_group(Data, Group) ->
     end,
     to_group(Tail, Group).
 
-exec(Cmd) ->
-    case eval(parse(scan(Cmd))) of
-	{error, _} ->
-	    {Cmd, 0}; %% This should be an external call
-	Term ->
-	    Term
-    end.
-
-scan(Cmd) ->
-    erl_scan:string(Cmd). 
-
-parse({ok, Tokens, _}) ->
-    erl_parse:parse_exprs(Tokens);
-parse(Error) ->
-    Error.
-
-eval({ok, Expr_list}) ->
-    case (catch erl_eval:exprs(Expr_list,
- 			       erl_eval:new_bindings())) of
- 	{value, Value, _NewBindings} ->
- 	    {Value, 0};
- 	{'EXIT', {Error, _}} -> 
- 	    {Error, -1};
- 	Error -> 
- 	    {Error, -1}
-    end;
-eval(Error) ->
-    {Error, -1}.
-
+%%--------------------------------------------------------------------
 %%% io_request, handle io requests from the user process,
 %%% Note, this is not the real I/O-protocol, but the mockup version
 %%% used between edlin and a user_driver. The protocol tags are
@@ -506,53 +489,130 @@ bin_to_list(L) when is_list(L) ->
 bin_to_list(I) when is_integer(I) ->
     I.
 
+
+%%--------------------------------------------------------------------
 start_shell(ConnectionHandler, State) ->
-    Shell = State#state.shell,
-    ConnectionInfo = ssh_connection_handler:connection_info(ConnectionHandler,
-						  [peer, user]),
-    ShellFun = case is_function(Shell) of
-		   true ->
-		       User = proplists:get_value(user, ConnectionInfo),
-		       case erlang:fun_info(Shell, arity) of
-			   {arity, 1} ->
-			       fun() -> Shell(User) end;
-			   {arity, 2} ->
-			       {_, PeerAddr} = proplists:get_value(peer, ConnectionInfo),
-			       fun() -> Shell(User, PeerAddr) end;
-			   _ ->
-			       Shell
-		       end;
-		   _ ->
-		       Shell
-	       end,
-    Echo = get_echo(State#state.pty),
-    Group = group:start(self(), ShellFun, [{echo, Echo}]),
-    State#state{group = Group, buf = empty_buf()}.
-
-start_shell(_ConnectionHandler, Cmd, #state{exec={M, F, A}} = State) ->
-    Group = group:start(self(), {M, F, A++[Cmd]}, [{echo, false}]),
-    State#state{group = Group, buf = empty_buf()};
-start_shell(ConnectionHandler, Cmd, #state{exec=Shell} = State) when is_function(Shell) ->
-
-    ConnectionInfo = ssh_connection_handler:connection_info(ConnectionHandler,
-						 [peer, user]),
-    User = proplists:get_value(user, ConnectionInfo),
-    ShellFun = 
-	case erlang:fun_info(Shell, arity) of
-	    {arity, 1} ->
-		fun() -> Shell(Cmd) end;
-	    {arity, 2} ->
-		fun() -> Shell(Cmd, User) end;
-	    {arity, 3} ->
-		{_, PeerAddr} = proplists:get_value(peer, ConnectionInfo),
-		fun() -> Shell(Cmd, User, PeerAddr) end;
-	    _ ->
-		Shell
-	end,
-    Echo = get_echo(State#state.pty),
-    Group = group:start(self(), ShellFun, [{echo,Echo}]),
-    State#state{group = Group, buf = empty_buf()}.
+    ShellSpawner =
+        case State#state.shell of
+            Shell when is_function(Shell, 1) ->
+                [{user,User}] = ssh_connection_handler:connection_info(ConnectionHandler, [user]),
+                fun() -> Shell(User) end;
+            Shell when is_function(Shell, 2) ->
+                ConnectionInfo =
+                    ssh_connection_handler:connection_info(ConnectionHandler, [peer, user]),
+                User = proplists:get_value(user, ConnectionInfo),
+                {_, PeerAddr} = proplists:get_value(peer, ConnectionInfo),
+                fun() -> Shell(User, PeerAddr) end;
+            {_,_,_} = Shell ->
+                Shell
+        end,
+    State#state{group = group:start(self(), ShellSpawner, [{echo, get_echo(State#state.pty)}]),
+                buf = empty_buf()}.
+
+%%--------------------------------------------------------------------
+start_exec_shell(ConnectionHandler, Cmd, State) ->
+    ExecShellSpawner =
+        case State#state.exec of
+            ExecShell when is_function(ExecShell, 1) ->
+                fun() -> ExecShell(Cmd) end;
+            ExecShell when is_function(ExecShell, 2) ->
+                [{user,User}] = ssh_connection_handler:connection_info(ConnectionHandler, [user]),
+                fun() -> ExecShell(Cmd, User) end;
+            ExecShell when is_function(ExecShell, 3) ->
+                ConnectionInfo =
+                    ssh_connection_handler:connection_info(ConnectionHandler, [peer, user]),
+                User = proplists:get_value(user, ConnectionInfo),
+                {_, PeerAddr} = proplists:get_value(peer, ConnectionInfo),
+                fun() -> ExecShell(Cmd, User, PeerAddr) end;
+            {M,F,A} ->
+                {M, F, A++[Cmd]}
+        end,
+    State#state{group = group:start(self(), ExecShellSpawner, [{echo,false}]),
+                buf = empty_buf()}.
+
+%%--------------------------------------------------------------------
+exec_in_erlang_default_shell(Cmd) ->
+    case eval(parse(scan(Cmd))) of
+	{ok, Term} ->
+            {io_lib:format("~p\n", [Term]), 0, 0};
+        {error, Error} when is_atom(Error) ->
+            {io_lib:format("Error in ~p: ~p\n", [Cmd,Error]), -1, 1};
+        _ ->
+            {io_lib:format("Error: ~p\n", [Cmd]), -1, 1}
+    end.
+
+scan(Cmd) ->
+    erl_scan:string(Cmd). 
+
+parse({ok, Tokens, _}) ->
+    erl_parse:parse_exprs(Tokens);
+parse(Error) ->
+    Error.
 
+eval({ok, Expr_list}) ->
+    case (catch erl_eval:exprs(Expr_list,
+                              erl_eval:new_bindings())) of
+        {value, Value, _NewBindings} ->
+            {ok, Value};
+        {'EXIT', {Error, _}} -> 
+            {error, Error};
+        {error, Error} -> 
+            {error, Error};
+        Error -> 
+            {error, Error}
+    end;
+eval({error,Error}) ->
+    {error, Error};
+eval(Error) ->
+    {error, Error}.
+
+%%--------------------------------------------------------------------
+exec_direct(ConnectionHandler, ExecSpec, Cmd) ->
+    try
+        case ExecSpec of
+            _ when is_function(ExecSpec, 1) ->
+                ExecSpec(Cmd);
+            _ when is_function(ExecSpec, 2) ->
+                [{user,User}] = ssh_connection_handler:connection_info(ConnectionHandler, [user]),
+                ExecSpec(Cmd, User);
+            _ when is_function(ExecSpec, 3) ->
+                ConnectionInfo =
+                    ssh_connection_handler:connection_info(ConnectionHandler, [peer, user]),
+                User = proplists:get_value(user, ConnectionInfo),
+                {_, PeerAddr} = proplists:get_value(peer, ConnectionInfo),
+                ExecSpec(Cmd, User, PeerAddr)
+        end
+    of
+        Reply ->
+            return_direct_exec_reply(Reply, Cmd)
+    catch
+        C:Error ->
+            {io_lib:format("Error in \"~s\": ~p ~p~n", [Cmd,C,Error]), -1, 1}
+    end.
+
+
+
+return_direct_exec_reply(Reply, Cmd) ->
+    case fmt_exec_repl(Reply) of
+        {ok,S} ->
+            {S, 0, 0};
+        {error,S} ->
+            {io_lib:format("Error in \"~s\": ~s~n", [Cmd,S]), -1, 1}
+    end.
+
+fmt_exec_repl({T,A}) when T==ok ; T==error ->
+    try
+        {T, io_lib:format("~s",[A])}
+    catch
+        error:badarg ->
+            {T, io_lib:format("~p", [A])};
+        C:Err ->
+            {error, io_lib:format("~p:~p~n",[C,Err])}
+    end;
+fmt_exec_repl(Other) ->
+    {error, io_lib:format("Bad exec-plugin return: ~p",[Other])}.
+
+%%--------------------------------------------------------------------
 % Pty can be undefined if the client never sets any pty options before
 % starting the shell.
 get_echo(undefined) ->
diff --git a/lib/ssh/src/ssh_options.erl b/lib/ssh/src/ssh_options.erl
index 1e10f72956..c05293d1ae 100644
--- a/lib/ssh/src/ssh_options.erl
+++ b/lib/ssh/src/ssh_options.erl
@@ -275,10 +275,12 @@ default(server) ->
             class => user_options
            },
 
-      {exec, def} =>                 % FIXME: need some archeology....
+      {exec, def} =>
           #{default => undefined,
-            chk => fun({M,F,_}) -> is_atom(M) andalso is_atom(F);
-                      (V) -> is_function(V)
+            chk => fun({direct, V}) ->  check_function1(V) orelse check_function2(V) orelse check_function3(V);
+                      %% Compatibility (undocumented):
+                      ({M,F,A}) -> is_atom(M) andalso is_atom(F) andalso is_list(A);
+                      (V) -> check_function1(V) orelse check_function2(V) orelse check_function3(V)
                    end,
             class => user_options
            },
diff --git a/lib/ssh/test/ssh_connection_SUITE.erl b/lib/ssh/test/ssh_connection_SUITE.erl
index 9587c0c251..257f2f70d7 100644
--- a/lib/ssh/test/ssh_connection_SUITE.erl
+++ b/lib/ssh/test/ssh_connection_SUITE.erl
@@ -50,6 +50,13 @@ all() ->
      start_shell,
      start_shell_exec,
      start_shell_exec_fun,
+     start_shell_exec_fun2,
+     start_shell_exec_fun3,
+     start_shell_exec_direct_fun,
+     start_shell_exec_direct_fun2,
+     start_shell_exec_direct_fun3,
+     start_shell_exec_direct_fun1_error,
+     start_shell_exec_direct_fun1_error_type,
      start_shell_sock_exec_fun,
      start_shell_sock_daemon_exec,
      connect_sock_not_tcp,
@@ -522,7 +529,7 @@ start_shell_exec(Config) when is_list(Config) ->
     {Pid, Host, Port} = ssh_test_lib:daemon([{system_dir, SysDir},
 					     {user_dir, UserDir},
 					     {password, "morot"},
-					     {exec, {?MODULE,ssh_exec,[]}} ]),
+					     {exec, {?MODULE,ssh_exec_echo,[]}} ]),
 
     ConnectionRef = ssh_test_lib:connect(Host, Port, [{silently_accept_hosts, true},
 						      {user, "foo"},
@@ -535,7 +542,7 @@ start_shell_exec(Config) when is_list(Config) ->
     success = ssh_connection:exec(ConnectionRef, ChannelId0,
 				  "testing", infinity),
     receive
-	{ssh_cm, ConnectionRef, {data, _ChannelId, 0, <<"testing\r\n">>}} ->
+	{ssh_cm, ConnectionRef, {data, _ChannelId, 0, <<"echo testing\r\n">>}} ->
 	    ok
     after 5000 ->
 	    ct:fail("Exec Timeout")
@@ -618,10 +625,49 @@ exec_erlang_term_non_default_shell(Config) when is_list(Config) ->
     TestResult.
 
 %%--------------------------------------------------------------------
-start_shell_exec_fun() ->
-    [{doc, "start shell to exec command"}].
+start_shell_exec_fun(Config) ->
+    do_start_shell_exec_fun(fun ssh_exec_echo/1,
+                            "testing", <<"echo testing\r\n">>, 0,
+                            Config).
+
+start_shell_exec_fun2(Config) ->
+    do_start_shell_exec_fun(fun ssh_exec_echo/2,
+                            "testing", <<"echo foo testing\r\n">>, 0,
+                            Config).
+
+start_shell_exec_fun3(Config) ->
+    do_start_shell_exec_fun(fun ssh_exec_echo/3,
+                            "testing", <<"echo foo testing\r\n">>, 0,
+                            Config).
+
+start_shell_exec_direct_fun(Config) ->
+    do_start_shell_exec_fun({direct, fun ssh_exec_direct_echo/1},
+                            "testing", <<"echo testing\n">>, 0,
+                            Config).
+
+start_shell_exec_direct_fun2(Config) ->
+    do_start_shell_exec_fun({direct, fun ssh_exec_direct_echo/2},
+                            "testing", <<"echo foo testing">>, 0,
+                            Config).
+
+start_shell_exec_direct_fun3(Config) ->
+    do_start_shell_exec_fun({direct, fun ssh_exec_direct_echo/3},
+                            "testing", <<"echo foo testing">>, 0,
+                            Config).
+
+start_shell_exec_direct_fun1_error(Config) ->
+    do_start_shell_exec_fun({direct, fun ssh_exec_direct_echo_error_return/1},
+                            "testing", <<"Error in \"testing\": {bad}\n">>, 1,
+                            Config).
+
+start_shell_exec_direct_fun1_error_type(Config) ->
+    do_start_shell_exec_fun({direct, fun ssh_exec_direct_echo_error_return_type/1},
+                            "testing", <<"Error in \"testing\": Bad exec-plugin return: very_bad\n">>, 1,
+                            Config).
+
+
 
-start_shell_exec_fun(Config) when is_list(Config) ->
+do_start_shell_exec_fun(Fun, Command, Expect, ExpectType, Config) ->
     PrivDir = proplists:get_value(priv_dir, Config),
     UserDir = filename:join(PrivDir, nopubkey), % to make sure we don't use public-key-auth
     file:make_dir(UserDir),
@@ -629,7 +675,7 @@ start_shell_exec_fun(Config) when is_list(Config) ->
     {Pid, Host, Port} = ssh_test_lib:daemon([{system_dir, SysDir},
 					     {user_dir, UserDir},
 					     {password, "morot"},
-					     {exec, fun ssh_exec/1}]),
+					     {exec, Fun}]),
 
     ConnectionRef = ssh_test_lib:connect(Host, Port, [{silently_accept_hosts, true},
 						      {user, "foo"},
@@ -639,14 +685,19 @@ start_shell_exec_fun(Config) when is_list(Config) ->
 
     {ok, ChannelId0} = ssh_connection:session_channel(ConnectionRef, infinity),
 
-    success = ssh_connection:exec(ConnectionRef, ChannelId0,
-				  "testing", infinity),
+    success = ssh_connection:exec(ConnectionRef, ChannelId0, Command, infinity),
 
     receive
-	{ssh_cm, ConnectionRef, {data, _ChannelId, 0, <<"testing\r\n">>}} ->
+	{ssh_cm, ConnectionRef, {data, _ChannelId, ExpectType, Expect}} ->
 	    ok
     after 5000 ->
-	    ct:fail("Exec Timeout")
+            receive
+                Other ->
+                    ct:pal("Received other:~n~p",[Other]),
+                    ct:fail("Unexpected response")
+            after 0 ->
+                    ct:fail("Exec Timeout")
+            end
     end,
 
     ssh:close(ConnectionRef),
@@ -664,7 +715,7 @@ start_shell_sock_exec_fun(Config) when is_list(Config) ->
     {Pid, HostD, Port} = ssh_test_lib:daemon([{system_dir, SysDir},
                                               {user_dir, UserDir},
                                               {password, "morot"},
-                                              {exec, fun ssh_exec/1}]),
+                                              {exec, fun ssh_exec_echo/1}]),
     Host = ssh_test_lib:ntoa(ssh_test_lib:mangle_connect_address(HostD)),
 
     {ok, Sock} = ssh_test_lib:gen_tcp_connect(Host, Port, [{active,false}]),
@@ -680,7 +731,7 @@ start_shell_sock_exec_fun(Config) when is_list(Config) ->
 				  "testing", infinity),
 
     receive
-	{ssh_cm, ConnectionRef, {data, _ChannelId, 0, <<"testing\r\n">>}} ->
+	{ssh_cm, ConnectionRef, {data, _ChannelId, 0, <<"echo testing\r\n">>}} ->
 	    ok
     after 5000 ->
 	    ct:fail("Exec Timeout")
@@ -704,7 +755,7 @@ start_shell_sock_daemon_exec(Config) ->
 		       {ok, _Pid} = ssh:daemon(Ss, [{system_dir, SysDir},
 						    {user_dir, UserDir},
 						    {password, "morot"},
-						    {exec, fun ssh_exec/1}])
+						    {exec, fun ssh_exec_echo/1}])
 	       end),
     {ok,Sc} = gen_tcp:accept(Sl),
     {ok,ConnectionRef} = ssh:connect(Sc, [{silently_accept_hosts, true},
@@ -719,7 +770,7 @@ start_shell_sock_daemon_exec(Config) ->
 				  "testing", infinity),
 
     receive
-	{ssh_cm, ConnectionRef, {data, _ChannelId, 0, <<"testing\r\n">>}} ->
+	{ssh_cm, ConnectionRef, {data, _ChannelId, 0, <<"echo testing\r\n">>}} ->
 	    ok
     after 5000 ->
 	    ct:fail("Exec Timeout")
@@ -830,7 +881,7 @@ stop_listener(Config) when is_list(Config) ->
     {Pid0, Host, Port} = ssh_test_lib:daemon([{system_dir, SysDir},
 					      {user_dir, UserDir},
 					      {password, "morot"},
-					      {exec, fun ssh_exec/1}]),
+					      {exec, fun ssh_exec_echo/1}]),
 
     ConnectionRef0 = ssh_test_lib:connect(Host, Port, [{silently_accept_hosts, true},
 						       {user, "foo"},
@@ -850,7 +901,7 @@ stop_listener(Config) when is_list(Config) ->
     success = ssh_connection:exec(ConnectionRef0, ChannelId0,
 				  "testing", infinity),
     receive
-	{ssh_cm, ConnectionRef0, {data, ChannelId0, 0, <<"testing\r\n">>}} ->
+	{ssh_cm, ConnectionRef0, {data, ChannelId0, 0, <<"echo testing\r\n">>}} ->
 	    ok
     after 5000 ->
 	    ct:fail("Exec Timeout")
@@ -859,7 +910,7 @@ stop_listener(Config) when is_list(Config) ->
     case ssh_test_lib:daemon(Port, [{system_dir, SysDir},
                                     {user_dir, UserDir},
                                     {password, "potatis"},
-                                    {exec, fun ssh_exec/1}]) of
+                                    {exec, fun ssh_exec_echo/1}]) of
 	{Pid1, Host, Port} ->
 	    ConnectionRef1 = ssh_test_lib:connect(Host, Port, [{silently_accept_hosts, true},
 							       {user, "foo"},
@@ -1070,7 +1121,22 @@ start_our_shell(_User, _Peer) ->
 		  %% Don't actually loop, just exit
           end).
 
-ssh_exec(Cmd) ->
+
+ssh_exec_echo(Cmd) ->
     spawn(fun() ->
-		  io:format(Cmd ++ "\n")
+                  io:format("echo "++Cmd ++ "\n")
           end).
+
+ssh_exec_echo(Cmd, User) ->
+    spawn(fun() ->
+                  io:format(io_lib:format("echo ~s ~s\n",[User,Cmd]))
+          end).
+ssh_exec_echo(Cmd, User, _PeerAddr) ->
+    ssh_exec_echo(Cmd,User).
+
+ssh_exec_direct_echo(Cmd) -> {ok, io_lib:format("echo ~s~n",[Cmd])}.
+ssh_exec_direct_echo(Cmd, User) -> {ok, io_lib:format("echo ~s ~s",[User,Cmd])}.
+ssh_exec_direct_echo(Cmd, User, _PeerAddr) -> ssh_exec_direct_echo(Cmd,User).
+
+ssh_exec_direct_echo_error_return(_Cmd) -> {error, {bad}}.
+ssh_exec_direct_echo_error_return_type(_Cmd) -> very_bad.