Improve robustness of external authentication backends

Now all external ports are attached to supervising processes
and requests are balanced in round-robin manner until the pool
is exhausted.

The commit also deprecates `extauth_instances` option and introduces
`extauth_pool_size` option instead, with the default value of a number
of logical processors (i.e. CPU cores).

Fixes #2403

4 files changed, 343 insertions, 141 deletions

diff --git a/src/ejabberd_auth.erl b/src/ejabberd_auth.erl
index 585b4a9bc..847549c74 100644
--- a/src/ejabberd_auth.erl
+++ b/src/ejabberd_auth.erl
@@ -69,6 +69,7 @@
 
 -callback start(binary()) -> any().
 -callback stop(binary()) -> any().
+-callback reload(binary()) -> any().
 -callback plain_password_required(binary()) -> boolean().
 -callback store_type(binary()) -> plain | external | scram.
 -callback set_password(binary(), binary(), binary()) -> ok | {error, atom()}.
@@ -82,7 +83,8 @@
 -callback use_cache(binary()) -> boolean().
 -callback cache_nodes(binary()) -> boolean().
 
--optional_callbacks([set_password/3,
+-optional_callbacks([reload/1,
+		     set_password/3,
 		     remove_user/2,
 		     user_exists/2,
 		     check_password/4,
@@ -130,14 +132,16 @@ handle_cast({host_down, Host}, #state{host_modules = HostModules} = State) ->
     init_cache(NewHostModules),
     {noreply, State#state{host_modules = NewHostModules}};
 handle_cast(config_reloaded, #state{host_modules = HostModules} = State) ->
-    NewHostModules = lists:foldl(
-		       fun(Host, Acc) ->
-			       OldModules = maps:get(Host, HostModules, []),
-			       NewModules = auth_modules(Host),
-			       start(Host, NewModules -- OldModules),
-			       stop(Host, OldModules -- NewModules),
-			       maps:put(Host, NewModules, Acc)
-		       end, HostModules, ?MYHOSTS),
+    NewHostModules =
+	lists:foldl(
+	  fun(Host, Acc) ->
+		  OldModules = maps:get(Host, HostModules, []),
+		  NewModules = auth_modules(Host),
+		  start(Host, NewModules -- OldModules),
+		  stop(Host, OldModules -- NewModules),
+		  reload(Host, lists_intersection(OldModules, NewModules)),
+		  maps:put(Host, NewModules, Acc)
+	  end, HostModules, ?MYHOSTS),
     init_cache(NewHostModules),
     {noreply, State#state{host_modules = NewHostModules}};
 handle_cast(Msg, State) ->
@@ -165,6 +169,15 @@ start(Host, Modules) ->
 stop(Host, Modules) ->
     lists:foreach(fun(M) -> M:stop(Host) end, Modules).
 
+reload(Host, Modules) ->
+    lists:foreach(
+      fun(M) ->
+	      case erlang:function_exported(M, reload, 1) of
+		  true -> M:reload(Host);
+		  false -> ok
+	      end
+      end, Modules).
+
 host_up(Host) ->
     gen_server:cast(?MODULE, {host_up, Host}).
 
@@ -559,7 +572,9 @@ db_user_exists(User, Server, Mod) ->
 			{ok, _} ->
 			    true;
 			error ->
-			    false
+			    false;
+			{error, _} = Err ->
+			    Err
 		    end;
 		{external, false} ->
 		    Mod:user_exists(User, Server);
@@ -816,6 +831,12 @@ validate_credentials(User, Server, Password) ->
 	    end
     end.
 
+lists_intersection(L1, L2) ->
+    lists:filter(
+      fun(E) ->
+              lists:member(E, L2)
+      end, L1).
+
 import_info() ->
     [{<<"users">>, 3}].
 
diff --git a/src/ejabberd_auth_external.erl b/src/ejabberd_auth_external.erl
index 9404ff0b1..b36c519d6 100644
--- a/src/ejabberd_auth_external.erl
+++ b/src/ejabberd_auth_external.erl
@@ -31,7 +31,7 @@
 
 -behaviour(ejabberd_auth).
 
--export([start/1, stop/1, set_password/3, check_password/4,
+-export([start/1, stop/1, reload/1, set_password/3, check_password/4,
 	 try_register/3, user_exists/2, remove_user/2,
 	 store_type/1, plain_password_required/1, opt_type/1]).
 
@@ -42,12 +42,14 @@
 %%% API
 %%%----------------------------------------------------------------------
 start(Host) ->
-    Cmd = ejabberd_config:get_option({extauth_program, Host}, "extauth"),
-    extauth:start(Host, Cmd).
+    extauth:start(Host).
 
 stop(Host) ->
     extauth:stop(Host).
 
+reload(Host) ->
+    extauth:reload(Host).
+
 plain_password_required(_) -> true.
 
 store_type(_) -> external.
@@ -61,37 +63,47 @@ check_password(User, AuthzId, Server, Password) ->
 
 set_password(User, Server, Password) ->
     case extauth:set_password(User, Server, Password) of
-	true -> ok;
-	_ -> {error, db_failure}
+	Res when is_boolean(Res) -> ok;
+	{error, Reason} -> failure(Reason)
     end.
 
 try_register(User, Server, Password) ->
-    extauth:try_register(User, Server, Password).
+    case extauth:try_register(User, Server, Password) of
+	true -> ok;
+	false -> {error, not_allowed};
+	{error, Reason} -> failure(Reason)
+    end.
 
 user_exists(User, Server) ->
-    try extauth:user_exists(User, Server) of
-	Res -> Res
-    catch
-	_:Error ->
-	    ?ERROR_MSG("external authentication program failure: ~p",
-		       [Error]),
-	    {error, db_failure}
+    case extauth:user_exists(User, Server) of
+	Res when is_boolean(Res) -> Res;
+	{error, Reason} -> failure(Reason)
     end.
 
 remove_user(User, Server) ->
     case extauth:remove_user(User, Server) of
 	false -> {error, not_allowed};
-	true -> ok
+	true -> ok;
+	{error, Reason} -> failure(Reason)
     end.
 
 check_password_extauth(User, _AuthzId, Server, Password) ->
-    extauth:check_password(User, Server, Password) andalso
-      Password /= <<"">>.
+    if Password /= <<"">> ->
+	    case extauth:check_password(User, Server, Password) of
+		Res when is_boolean(Res) -> Res;
+		{error, Reason} ->
+		    failure(Reason),
+		    false
+	    end;
+       true ->
+	    false
+    end.
+
+-spec failure(any()) -> {error, db_failure}.
+failure(Reason) ->
+    ?ERROR_MSG("External authentication program failure: ~p", [Reason]),
+    {error, db_failure}.
 
--spec opt_type(extauth_cache) -> fun((false | non_neg_integer()) ->
-				      false | non_neg_integer());
-	      (extauth_program) -> fun((binary()) -> string());
-	      (atom()) -> [atom()].
 opt_type(extauth_cache) ->
     ?WARNING_MSG("option 'extauth_cache' is deprecated and has no effect, "
 		 "use authentication or global cache configuration "
@@ -100,6 +112,15 @@ opt_type(extauth_cache) ->
     fun (false) -> false;
 	(I) when is_integer(I), I >= 0 -> I
     end;
+opt_type(extauth_instances) ->
+    ?WARNING_MSG("option 'extauth_instances' is deprecated and has no effect, "
+		 "use 'extauth_pool_size'", []),
+    fun (V) when is_integer(V), V > 0 -> V end;
 opt_type(extauth_program) ->
     fun (V) -> binary_to_list(iolist_to_binary(V)) end;
-opt_type(_) -> [extauth_cache, extauth_program].
+opt_type(extauth_pool_size) ->
+    fun(I) when is_integer(I), I>0 -> I end;
+opt_type(_) ->
+    [extauth_program, extauth_pool_size,
+     %% Deprecated:
+     extauth_cache, extauth_instances].
diff --git a/src/extauth.erl b/src/extauth.erl
index ee2806dde..9c360bff2 100644
--- a/src/extauth.erl
+++ b/src/extauth.erl
@@ -1,8 +1,5 @@
-%%%----------------------------------------------------------------------
-%%% File    : extauth.erl
-%%% Author  : Leif Johansson <leifj@it.su.se>
-%%% Purpose : External authentication using a simple port-driver
-%%% Created : 30 Jul 2004 by Leif Johansson <leifj@it.su.se>
+%%%-------------------------------------------------------------------
+%%% Created : 7 May 2018 by Evgeny Khramtsov <ekhramtsov@process-one.net>
 %%%
 %%%
 %%% ejabberd, Copyright (C) 2002-2018   ProcessOne
@@ -21,54 +18,45 @@
 %%% with this program; if not, write to the Free Software Foundation, Inc.,
 %%% 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 %%%
-%%%----------------------------------------------------------------------
-
+%%%-------------------------------------------------------------------
 -module(extauth).
 
--behaviour(ejabberd_config).
+-ifndef(GEN_SERVER).
+-define(GEN_SERVER, gen_server).
+-endif.
+-behaviour(?GEN_SERVER).
 
--author('leifj@it.su.se').
+-define(CALL_TIMEOUT, timer:seconds(30)).
 
--export([start/2, stop/1, init/2, check_password/3,
-	 set_password/3, try_register/3, remove_user/2,
-	 remove_user/3, user_exists/2, opt_type/1]).
+%% API
+-export([start/1, stop/1, reload/1, start_link/2]).
+-export([check_password/3, set_password/3, try_register/3, remove_user/2,
+	 remove_user/3, user_exists/2]).
+-export([prog_name/1, pool_name/1, worker_name/2, pool_size/1]).
+-export([init/1, handle_call/3, handle_cast/2, handle_info/2,
+	 terminate/2, code_change/3]).
 
--include("ejabberd.hrl").
 -include("logger.hrl").
 
--define(INIT_TIMEOUT, 60000).
-
--define(CALL_TIMEOUT, 10000).
-
-start(Host, ExtPrg) ->
-    lists:foreach(fun (This) ->
-			  start_instance(get_process_name(Host, This), ExtPrg)
-		  end,
-		  lists:seq(0, get_instances(Host) - 1)).
-
-start_instance(ProcessName, ExtPrg) ->
-    spawn(?MODULE, init, [ProcessName, ExtPrg]).
+-record(state, {port :: port(),
+		prog :: string(),
+		start_time :: integer(),
+		os_pid :: integer() | undefined}).
 
-restart_instance(ProcessName, ExtPrg) ->
-    unregister(ProcessName),
-    start_instance(ProcessName, ExtPrg).
-
-init(ProcessName, ExtPrg) ->
-    register(ProcessName, self()),
-    process_flag(trap_exit, true),
-    Port = open_port({spawn, ExtPrg}, [{packet, 2}]),
-    loop(Port, ?INIT_TIMEOUT, ProcessName, ExtPrg).
+%%%===================================================================
+%%% API
+%%%===================================================================
+start(Host) ->
+    extauth_sup:start(Host).
 
 stop(Host) ->
-    lists:foreach(fun (This) ->
-			  get_process_name(Host, This) ! stop
-		  end,
-		  lists:seq(0, get_instances(Host) - 1)).
+    extauth_sup:stop(Host).
+
+reload(Host) ->
+    extauth_sup:reload(Host).
 
-get_process_name(Host, Integer) ->
-    gen_mod:get_module_proc(iolist_to_binary([Host,
-                                              integer_to_list(Integer)]),
-			    eauth).
+start_link(Name, Prog) ->
+    ?GEN_SERVER:start_link({local, Name}, ?MODULE, [Prog], []).
 
 check_password(User, Server, Password) ->
     call_port(Server, [<<"auth">>, User, Server, Password]).
@@ -80,82 +68,144 @@ set_password(User, Server, Password) ->
     call_port(Server, [<<"setpass">>, User, Server, Password]).
 
 try_register(User, Server, Password) ->
-    case call_port(Server,
-		   [<<"tryregister">>, User, Server, Password])
-	of
-      true -> ok;
-      false -> {error, not_allowed}
-    end.
+    call_port(Server, [<<"tryregister">>, User, Server, Password]).
 
 remove_user(User, Server) ->
     call_port(Server, [<<"removeuser">>, User, Server]).
 
 remove_user(User, Server, Password) ->
-    call_port(Server,
-	      [<<"removeuser3">>, User, Server, Password]).
-
-call_port(Server, Msg) ->
-    LServer = jid:nameprep(Server),
-    ProcessName = get_process_name(LServer,
-				   random_instance(get_instances(LServer))),
-    ProcessName ! {call, self(), Msg},
-    receive {eauth, Result} -> Result end.
-
-random_instance(MaxNum) ->
-    randoms:uniform(MaxNum) - 1.
-
-get_instances(Server) ->
-    ejabberd_config:get_option({extauth_instances, Server}, 1).
-
-loop(Port, Timeout, ProcessName, ExtPrg) ->
-    receive
-      {call, Caller, Msg} ->
-	  port_command(Port, encode(Msg)),
-	  receive
-	    {Port, {data, Data}} ->
-		?DEBUG("extauth call '~p' received data response:~n~p",
-		       [Msg, Data]),
-		Caller ! {eauth, decode(Data)},
-		loop(Port, ?CALL_TIMEOUT, ProcessName, ExtPrg);
-	    {Port, Other} ->
-		?ERROR_MSG("extauth call '~p' received strange response:~n~p",
-			   [Msg, Other]),
-		Caller ! {eauth, false},
-		loop(Port, ?CALL_TIMEOUT, ProcessName, ExtPrg)
-	    after Timeout ->
-		      ?ERROR_MSG("extauth call '~p' didn't receive response",
-				 [Msg]),
-		      Caller ! {eauth, false},
-		      Pid = restart_instance(ProcessName, ExtPrg),
-		      flush_buffer_and_forward_messages(Pid),
-		      exit(port_terminated)
-	  end;
-      stop ->
-	  Port ! {self(), close},
-	  receive {Port, closed} -> exit(normal) end;
-      {'EXIT', Port, Reason} ->
-	  ?CRITICAL_MSG("extauth script has exitted abruptly "
-			"with reason '~p'",
-			[Reason]),
-	  Pid = restart_instance(ProcessName, ExtPrg),
-	  flush_buffer_and_forward_messages(Pid),
-	  exit(port_terminated)
+    call_port(Server, [<<"removeuser3">>, User, Server, Password]).
+
+-spec prog_name(binary()) -> string() | undefined.
+prog_name(Host) ->
+    ejabberd_config:get_option({extauth_program, Host}).
+
+-spec pool_name(binary()) -> atom().
+pool_name(Host) ->
+    list_to_atom("extauth_pool_" ++ binary_to_list(Host)).
+
+-spec worker_name(atom(), integer()) -> atom().
+worker_name(Pool, N) ->
+    list_to_atom(atom_to_list(Pool) ++ "_" ++ integer_to_list(N)).
+
+-spec pool_size(binary()) -> pos_integer().
+pool_size(Host) ->
+    case ejabberd_config:get_option({extauth_pool_size, Host}) of
+	undefined ->
+	    try erlang:system_info(logical_processors)
+	    catch _:_ -> 1
+	    end;
+	Size ->
+	    Size
     end.
 
-flush_buffer_and_forward_messages(Pid) ->
-    receive
-      Message ->
-	  Pid ! Message, flush_buffer_and_forward_messages(Pid)
-      after 0 -> true
+%%%===================================================================
+%%% gen_server callbacks
+%%%===================================================================
+init([Prog]) ->
+    process_flag(trap_exit, true),
+    {Port, OSPid} = start_port(Prog),
+    Time = curr_time(),
+    {ok, #state{port = Port, start_time = Time,
+		prog = Prog, os_pid = OSPid}}.
+
+handle_call({cmd, Cmd, EndTime}, _From, State) ->
+    Timeout = EndTime - curr_time(),
+    if Timeout > 0 ->
+	    Port = State#state.port,
+	    port_command(Port, Cmd),
+	    receive
+		{Port, {data, [0, N] = Data}} when N == 0; N == 1 ->
+		    ?DEBUG("Received response from external authentication "
+			   "program: ~p", [Data]),
+		    {reply, decode_bool(N), State};
+		{Port, Data} ->
+		    ?ERROR_MSG("Received unexpected response from external "
+			       "authentication program '~s': ~p "
+			       "(port = ~p, pid = ~w)",
+			       [State#state.prog, Data, Port, State#state.os_pid]),
+		    {reply, {error, unexpected_response}, State};
+		{'EXIT', Port, Reason} ->
+		    handle_info({'EXIT', Port, Reason}, State)
+	    after Timeout ->
+		    {stop, extauth_program_timeout, State}
+	    end;
+       true ->
+	    {stop, extauth_program_timeout, State}
     end.
 
-encode(L) -> str:join(L, <<":">>).
+handle_cast(_Msg, State) ->
+    {noreply, State}.
+
+handle_info({'EXIT', Port, _Reason}, #state{port = Port,
+					    start_time = Time} = State) ->
+    case curr_time() - Time of
+	Diff when Diff < 1000 ->
+	    ?ERROR_MSG("Failed to start external authentication program '~s'",
+		       [State#state.prog]),
+	    {stop, normal, State};
+	_ ->
+	    ?ERROR_MSG("External authentication program '~s' has terminated "
+		       "unexpectedly (pid=~w), restarting via supervisor...",
+		       [State#state.prog, State#state.os_pid]),
+	    {stop, extauth_program_failure, State}
+    end;
+handle_info(Info, State) ->
+    ?WARNING_MSG("Unexpected info: ~p", [Info]),
+    {noreply, State}.
+
+terminate(_Reason, State) ->
+    catch port_close(State#state.port),
+    ok.
+
+code_change(_OldVsn, State, _Extra) ->
+    {ok, State}.
+
+%%%===================================================================
+%%% Internal functions
+%%%===================================================================
+-spec curr_time() -> non_neg_integer().
+curr_time() ->
+    p1_time_compat:monotonic_time(milli_seconds).
+
+-spec start_port(string()) -> {port(), integer() | undefined}.
+start_port(Path) ->
+    Port = open_port({spawn, Path}, [{packet, 2}]),
+    link(Port),
+    case erlang:port_info(Port, os_pid) of
+	{os_pid, OSPid} ->
+	    {Port, OSPid};
+	undefined ->
+	    {Port, undefined}
+    end.
 
-decode([0, 0]) -> false;
-decode([0, 1]) -> true.
+call_port(Server, Args) ->
+    call_port(Server, Args, ?CALL_TIMEOUT).
+
+call_port(Server, Args, Timeout) ->
+    StartTime = p1_time_compat:monotonic_time(milli_seconds),
+    Pool = pool_name(Server),
+    PoolSize = pool_size(Server),
+    I = randoms:round_robin(PoolSize),
+    Cmd = str:join(Args, <<":">>),
+    do_call(Cmd, I, I + PoolSize, Pool, PoolSize,
+	    StartTime + Timeout, StartTime).
+
+do_call(_, Max, Max, _, _, _, _) ->
+    {error, disconnected};
+do_call(Cmd, I, Max, Pool, PoolSize, EndTime, CurrTime) ->
+    Timeout = EndTime - CurrTime,
+    if Timeout > 0 ->
+	    Proc = worker_name(Pool, (I rem PoolSize) + 1),
+	    try ?GEN_SERVER:call(Proc, {cmd, Cmd, EndTime}, Timeout)
+	    catch exit:{timeout, {?GEN_SERVER, call, _}} ->
+		    {error, timeout};
+		  exit:{_, {?GEN_SERVER, call, _}} ->
+		    do_call(Cmd, I+1, Max, Pool, PoolSize, EndTime, curr_time())
+	    end;
+       true ->
+	    {error, timeout}
+    end.
 
--spec opt_type(extauth_instances) -> fun((pos_integer()) -> pos_integer());
-	      (atom()) -> [atom()].
-opt_type(extauth_instances) ->
-    fun (V) when is_integer(V), V > 0 -> V end;
-opt_type(_) -> [extauth_instances].
+decode_bool(0) -> false;
+decode_bool(1) -> true.
diff --git a/src/extauth_sup.erl b/src/extauth_sup.erl
new file mode 100644
index 000000000..697c1dec6
--- /dev/null
+++ b/src/extauth_sup.erl
@@ -0,0 +1,110 @@
+%%%-------------------------------------------------------------------
+%%% Created : 7 May 2018 by Evgeny Khramtsov <ekhramtsov@process-one.net>
+%%%
+%%%
+%%% ejabberd, Copyright (C) 2002-2018   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(extauth_sup).
+-behaviour(supervisor).
+
+%% API
+-export([start/1, stop/1, reload/1, start_link/3]).
+%% Supervisor callbacks
+-export([init/1]).
+
+-include("logger.hrl").
+
+%%%===================================================================
+%%% API functions
+%%%===================================================================
+start(Host) ->
+    case extauth:prog_name(Host) of
+	undefined ->
+	    ?ERROR_MSG("Option 'extauth_program' is not set for '~s'",
+		       [Host]),
+	    ignore;
+	Prog ->
+	    Pool = extauth:pool_name(Host),
+	    ChildSpec = {Pool, {?MODULE, start_link, [Host, Prog, Pool]},
+			 transient, infinity, supervisor, [?MODULE]},
+	    supervisor:start_child(ejabberd_backend_sup, ChildSpec)
+    end.
+
+stop(Host) ->
+    Pool = extauth:pool_name(Host),
+    supervisor:terminate_child(ejabberd_backend_sup, Pool),
+    supervisor:delete_child(ejabberd_backend_sup, Pool).
+
+reload(Host) ->
+    Pool = extauth:pool_name(Host),
+    Prog = extauth:prog_name(Host),
+    PoolSize = extauth:pool_size(Host),
+    try process_info(whereis(Pool), dictionary) of
+	{dictionary, Dict} ->
+	    case proplists:get_value(extauth_program, Dict) of
+		Prog ->
+		    OldPoolSize = try supervisor:which_children(Pool) of
+				      Children -> length(Children)
+				  catch _:_ -> PoolSize
+				  end,
+		    if OldPoolSize > PoolSize ->
+			    lists:foreach(
+			      fun(I) ->
+				      Worker = extauth:worker_name(Pool, I),
+				      supervisor:terminate_child(Pool, Worker),
+				      supervisor:delete_child(Pool, Worker)
+			      end, lists:seq(PoolSize+1, OldPoolSize));
+		       OldPoolSize < PoolSize ->
+			    lists:foreach(
+			      fun(I) ->
+				      Spec = worker_spec(Pool, Prog, I),
+				      supervisor:start_child(Pool, Spec)
+			      end, lists:seq(OldPoolSize+1, PoolSize));
+		       OldPoolSize == PoolSize ->
+			    ok
+		    end;
+		_ ->
+		    stop(Host),
+		    start(Host)
+	    end
+    catch _:badarg ->
+	    ok
+    end.
+
+start_link(Host, Prog, Pool) ->
+    supervisor:start_link({local, Pool}, ?MODULE, [Host, Prog, Pool]).
+
+%%%===================================================================
+%%% Supervisor callbacks
+%%%===================================================================
+init([Host, Prog, Pool]) ->
+    PoolSize = extauth:pool_size(Host),
+    Children = lists:map(
+		 fun(I) ->
+			 worker_spec(Pool, Prog, I)
+		 end, lists:seq(1, PoolSize)),
+    put(extauth_program, Prog),
+    {ok, {{one_for_one, PoolSize, 1}, Children}}.
+
+%%%===================================================================
+%%% Internal functions
+%%%===================================================================
+worker_spec(Pool, Prog, I) ->
+    Worker = extauth:worker_name(Pool, I),
+    {Worker, {extauth, start_link, [Worker, Prog]},
+     permanent, 5000, worker, [extauth]}.