Extend gen_mod API to simplify hooks and IQ handlers registration

src/gen_mod.erl

@@ -49,6 +49,7 @@
 -record(ejabberd_module,
         {module_host = {undefined, <<"">>} :: {atom(), binary()},
          opts = [] :: opts() | '_' | '$2',
+         registrations = [] :: [registration()],
 	 order = 0 :: integer()}).
 
 -type opts() :: #{atom() => term()}.
@@ -57,7 +58,14 @@
                       value => string() | binary()}.
 -type opt_doc() :: {atom(), opt_desc()} | {atom(), opt_desc(), [opt_doc()]}.
 
--callback start(binary(), opts()) -> ok | {ok, pid()} | {error, term()}.
+-type component() :: ejabberd_sm | ejabberd_local.
+-type registration() ::
+        {hook, atom(), module(), atom(), integer()} |
+        {iq_handler, component(), binary(), module(), atom()}.
+
+-callback start(binary(), opts()) ->
+    ok | {ok, pid()} |
+    {ok, [registration()]} | {error, term()}.
 -callback stop(binary()) -> any().
 -callback reload(binary(), opts(), opts()) -> ok | {ok, pid()} | {error, term()}.
 -callback mod_opt_type(atom()) -> econf:validator().
@@ -155,6 +163,10 @@ start_module(Host, Module, Opts, Order) ->
     try case Module:start(Host, Opts) of
 	    ok -> ok;
 	    {ok, Pid} when is_pid(Pid) -> {ok, Pid};
+	    {ok, Registrations} when is_list(Registrations) ->
+                store_options(Host, Module, Opts, Registrations, Order),
+                add_registrations(Host, Registrations),
+                ok;
 	    Err ->
 		ets:delete(ejabberd_modules, {Module, Host}),
 		erlang:error({bad_return, Module, Err})
@@ -246,9 +258,21 @@ update_module(Host, Module, Opts) ->
 
 -spec store_options(binary(), module(), opts(), integer()) -> true.
 store_options(Host, Module, Opts, Order) ->
+    case ets:lookup(ejabberd_modules, {Module, Host}) of
+	[M] ->
+            store_options(
+              Host, Module, Opts, M#ejabberd_module.registrations, Order);
+	[] ->
+            store_options(Host, Module, Opts, [], Order)
+    end.
+
+-spec store_options(binary(), module(), opts(), [registration()], integer()) -> true.
+store_options(Host, Module, Opts, Registrations, Order) ->
     ets:insert(ejabberd_modules,
 	       #ejabberd_module{module_host = {Module, Host},
-				opts = Opts, order = Order}).
+				opts = Opts,
+                                registrations = Registrations,
+                                order = Order}).
 
 maybe_halt_ejabberd() ->
     case is_app_running(ejabberd) of
@@ -281,16 +305,21 @@ stop_modules(Host) ->
 		stop_module_keep_config(Host, Module)
 	end, Modules).
 
--spec stop_module(binary(), atom()) -> error | {aborted, any()} | {atomic, any()}.
+-spec stop_module(binary(), atom()) -> error | ok.
 stop_module(Host, Module) ->
-    case stop_module_keep_config(Host, Module) of
-      error -> error;
-      ok -> ok
-    end.
+    stop_module_keep_config(Host, Module).
 
 -spec stop_module_keep_config(binary(), atom()) -> error | ok.
 stop_module_keep_config(Host, Module) ->
     ?DEBUG("Stopping ~ts at ~ts", [Module, Host]),
+    Registrations =
+        case ets:lookup(ejabberd_modules, {Module, Host}) of
+            [M] ->
+                M#ejabberd_module.registrations;
+            [] ->
+                []
+        end,
+    del_registrations(Host, Registrations),
     try Module:stop(Host) of
 	_ ->
 	    ets:delete(ejabberd_modules, {Module, Host}),
@@ -303,6 +332,25 @@ stop_module_keep_config(Host, Module) ->
 	    error
     end.
 
+-spec add_registrations(binary(), [registration()]) -> ok.
+add_registrations(Host, Registrations) ->
+    lists:foreach(
+      fun({hook, Hook, Module, Function, Seq}) ->
+              ejabberd_hooks:add(Hook, Host, Module, Function, Seq);
+         ({iq_handler, Component, NS, Module, Function}) ->
+              gen_iq_handler:add_iq_handler(
+                Component, Host, NS, Module, Function)
+      end, Registrations).
+
+-spec del_registrations(binary(), [registration()]) -> ok.
+del_registrations(Host, Registrations) ->
+    lists:foreach(
+      fun({hook, Hook, Module, Function, Seq}) ->
+              ejabberd_hooks:delete(Hook, Host, Module, Function, Seq);
+         ({iq_handler, Component, NS, _Module, _Function}) ->
+              gen_iq_handler:remove_iq_handler(Component, Host, NS)
+      end, Registrations).
+
 -spec get_opt(atom(), opts()) -> any().
 get_opt(Opt, Opts) ->
     maps:get(Opt, Opts).