From 67e06b317183df7c296b18b91501f5ec428a2736 Mon Sep 17 00:00:00 2001
From: Ilya Averyanov <av@rubybox.dev>
Date: Mon, 7 Aug 2023 19:29:50 +0300
Subject: [PATCH] chore(auth): make schema injection be more universal

---
 apps/emqx/src/emqx_schema.erl             |  4 +-
 apps/emqx/src/emqx_schema_hooks.erl       | 74 +++++++++++++----------
 apps/emqx_authn/src/emqx_authn_schema.erl |  4 +-
 apps/emqx_conf/src/emqx_conf_schema.erl   |  8 +--
 4 files changed, 48 insertions(+), 42 deletions(-)

diff --git a/apps/emqx/src/emqx_schema.erl b/apps/emqx/src/emqx_schema.erl
index 40654a84b..3f4507a3c 100644
--- a/apps/emqx/src/emqx_schema.erl
+++ b/apps/emqx/src/emqx_schema.erl
@@ -222,7 +222,7 @@ roots(high) ->
                 ref(?EMQX_AUTHORIZATION_CONFIG_ROOT_NAME),
                 #{importance => ?IMPORTANCE_HIDDEN}
             )}
-    ] ++ emqx_schema_hooks:injection_point('authentication');
+    ] ++ emqx_schema_hooks:injection_point('roots.high');
 roots(medium) ->
     [
         {"broker",
@@ -1749,7 +1749,7 @@ mqtt_listener(Bind) ->
                         default => <<"3s">>
                     }
                 )}
-        ] ++ emqx_schema_hooks:injection_point('listeners.authentication').
+        ] ++ emqx_schema_hooks:injection_point('mqtt.listener').
 
 base_listener(Bind) ->
     [
diff --git a/apps/emqx/src/emqx_schema_hooks.erl b/apps/emqx/src/emqx_schema_hooks.erl
index fd80db635..bd50a069f 100644
--- a/apps/emqx/src/emqx_schema_hooks.erl
+++ b/apps/emqx/src/emqx_schema_hooks.erl
@@ -25,11 +25,12 @@
 -optional_callbacks([injected_fields/0]).
 
 -define(HOOKPOINT_PT_KEY(POINT_NAME), {?MODULE, fields, POINT_NAME}).
--define(MODULE_PT_KEY(MOD_NAME), {?MODULE, mod, MOD_NAME}).
 
 -export([
     injection_point/1,
-    inject_fields_from_mod/1
+    any_injections/1,
+    inject_fields/2,
+    inject_from_modules/1
 ]).
 
 %% for tests
@@ -45,22 +46,15 @@
 injection_point(PointName) ->
     persistent_term:get(?HOOKPOINT_PT_KEY(PointName), []).
 
-inject_fields_from_mod(Module) ->
-    case persistent_term:get(?MODULE_PT_KEY(Module), false) of
-        false ->
-            persistent_term:put(?MODULE_PT_KEY(Module), true),
-            do_inject_fields_from_mod(Module);
-        true ->
-            ok
-    end.
+inject_fields(PointName, Fields) ->
+    Key = ?HOOKPOINT_PT_KEY(PointName),
+    persistent_term:put(Key, Fields).
 
 erase_injections() ->
     lists:foreach(
         fun
             ({?HOOKPOINT_PT_KEY(_) = Key, _}) ->
                 persistent_term:erase(Key);
-            ({?MODULE_PT_KEY(_) = Key, _}) ->
-                persistent_term:erase(Key);
             (_) ->
                 ok
         end,
@@ -72,35 +66,53 @@ any_injections() ->
         fun
             ({?HOOKPOINT_PT_KEY(_), _}) ->
                 true;
-            ({?MODULE_PT_KEY(_), _}) ->
-                true;
             (_) ->
                 false
         end,
         persistent_term:get()
     ).
 
+any_injections(PointName) ->
+    persistent_term:get(?HOOKPOINT_PT_KEY(PointName), undefined) =/= undefined.
+
+inject_from_modules(Modules) ->
+    Injections =
+        lists:foldl(
+            fun append_module_injections/2,
+            #{},
+            Modules
+        ),
+    ok = inject_fields(maps:to_list(Injections)).
+
 %%--------------------------------------------------------------------
 %% Internal functions
 %%--------------------------------------------------------------------
 
-do_inject_fields_from_mod(Module) ->
-    _ = Module:module_info(),
-    case erlang:function_exported(Module, injected_fields, 0) of
-        true ->
-            do_inject_fields_from_mod(Module, Module:injected_fields());
-        false ->
-            ok
-    end.
-
-do_inject_fields_from_mod(_Module, HookFields) ->
-    maps:foreach(
-        fun(PointName, Fields) ->
-            inject_fields(PointName, Fields)
+append_module_injections(Module, AllInjections) when is_atom(Module) ->
+    append_module_injections(Module:injected_fields(), AllInjections);
+append_module_injections(ModuleInjections, AllInjections) when is_map(ModuleInjections) ->
+    maps:fold(
+        fun(PointName, Fields, Acc) ->
+            maps:update_with(
+                PointName,
+                fun(Fields0) ->
+                    Fields0 ++ Fields
+                end,
+                Fields,
+                Acc
+            )
         end,
-        HookFields
+        AllInjections,
+        ModuleInjections
     ).
 
-inject_fields(PointName, Fields) ->
-    Key = ?HOOKPOINT_PT_KEY(PointName),
-    persistent_term:put(Key, Fields).
+inject_fields([]) ->
+    ok;
+inject_fields([{PointName, Fields} | Rest]) ->
+    case emqx_schema_hooks:any_injections(PointName) of
+        true ->
+            inject_fields(Rest);
+        false ->
+            ok = emqx_schema_hooks:inject_fields(PointName, Fields),
+            inject_fields(Rest)
+    end.
diff --git a/apps/emqx_authn/src/emqx_authn_schema.erl b/apps/emqx_authn/src/emqx_authn_schema.erl
index f77158f0c..b0a68e702 100644
--- a/apps/emqx_authn/src/emqx_authn_schema.erl
+++ b/apps/emqx_authn/src/emqx_authn_schema.erl
@@ -42,8 +42,8 @@ roots() -> [].
 
 injected_fields() ->
     #{
-        'authentication' => global_auth_fields(),
-        'listeners.authentication' => mqtt_listener_auth_fields()
+        'mqtt.listener' => global_auth_fields(),
+        'roots.high' => mqtt_listener_auth_fields()
     }.
 
 tags() ->
diff --git a/apps/emqx_conf/src/emqx_conf_schema.erl b/apps/emqx_conf/src/emqx_conf_schema.erl
index 14cd4f17c..5b73b9e03 100644
--- a/apps/emqx_conf/src/emqx_conf_schema.erl
+++ b/apps/emqx_conf/src/emqx_conf_schema.erl
@@ -80,7 +80,7 @@ tags() ->
     [<<"EMQX">>].
 
 roots() ->
-    ok = ensure_fields_injected(),
+    ok = emqx_schema_hooks:inject_from_modules(?INJECTING_CONFIGS),
     emqx_schema_high_prio_roots() ++
         [
             {"node",
@@ -1434,9 +1434,3 @@ ensure_unicode_path(Path, _) when is_list(Path) ->
     Path;
 ensure_unicode_path(Path, _) ->
     throw({"not_string", Path}).
-
-ensure_fields_injected() ->
-    lists:foreach(
-        fun(Module) -> emqx_schema_hooks:inject_fields_from_mod(Module) end,
-        ?INJECTING_CONFIGS
-    ).