Add a callback to modify mam iq-fin element

esl · Jul 21, 2021 · ddf25a9 · ddf25a9
1 parent 83a741c
commit ddf25a9
Show file tree

Hide file tree

Showing 6 changed files with 35 additions and 15 deletions.
diff --git a/src/mam/mod_mam.erl b/src/mam/mod_mam.erl
@@ -82,7 +82,7 @@
          wrap_message/6,
          result_set/4,
          result_prefs/4,
-         make_fin_element/4,
+         make_fin_element/5,
          parse_prefs/1,
          is_mam_result_message/1,
          features/2]).
@@ -440,7 +440,8 @@ handle_set_message_form(#jid{} = From, #jid{} = ArcJID,
             IsComplete = is_complete_result_page(TotalCount, Offset, MessageRows, Params),
             IsStable = true,
             ResultSetEl = result_set(FirstMessID, LastMessID, Offset, TotalCount),
-            FinElem = make_fin_element(IQ#iq.xmlns, IsComplete, IsStable, ResultSetEl),
+            ExtFinMod = mod_mam_params:extra_fin_element_module(?MODULE, HostType),
+            FinElem = make_fin_element(IQ#iq.xmlns, IsComplete, IsStable, ResultSetEl, ExtFinMod),
             IQ#iq{type = result, sub_el = [FinElem]}
     end.
 

diff --git a/src/mam/mod_mam_meta.erl b/src/mam/mod_mam_meta.erl
@@ -94,6 +94,8 @@ config_items() ->
       <<"db_message_format">> => #option{type = atom,
                                          validate = module},
       <<"simple">> => #option{type = boolean},
+      <<"extra_fin_element">> => #option{type = atom,
+                                         validate = module},
       <<"extra_lookup_params">> => #option{type = atom,
                                            validate = module}
      }.
@@ -213,6 +215,7 @@ common_opts() ->
     [is_archivable_message,
      send_message,
      archive_chat_markers,
+     extra_fin_element,
      extra_lookup_params,
      full_text_search,
      message_retraction,

diff --git a/src/mam/mod_mam_muc.erl b/src/mam/mod_mam_muc.erl
@@ -79,7 +79,7 @@
          result_query/2,
          result_prefs/4,
          make_fin_message/5,
-         make_fin_element/4,
+         make_fin_element/5,
          parse_prefs/1,
          borders_decode/1,
          features/2]).
@@ -411,7 +411,8 @@ send_messages_and_iq_result({TotalCount, Offset, MessageRows}, HostType, From,
     IsComplete = is_complete_result_page(TotalCount, Offset, MessageRows, Params),
     IsStable = true,
     ResultSetEl = result_set(FirstMessID, LastMessID, Offset, TotalCount),
-    FinElem = make_fin_element(IQ#iq.xmlns, IsComplete, IsStable, ResultSetEl),
+    ExtFinMod = mod_mam_params:extra_fin_element_module(?MODULE, HostType),
+    FinElem = make_fin_element(IQ#iq.xmlns, IsComplete, IsStable, ResultSetEl, ExtFinMod),
     IQ#iq{type = result, sub_el = [FinElem]}.
 
 forward_messages(HostType, From, ArcJID, MamNs, QueryID, MessageRows, SetClientNs) ->

diff --git a/src/mam/mod_mam_params.erl b/src/mam/mod_mam_params.erl
@@ -20,7 +20,7 @@
 
 -export([extra_params_module/2, max_result_limit/2, default_result_limit/2,
          has_full_text_search/2, is_archivable_message_fun/2, send_message_mod/2,
-         archive_chat_markers/2, add_stanzaid_element/2]).
+         archive_chat_markers/2, add_stanzaid_element/2, extra_fin_element_module/2]).
 
 %%--------------------------------------------------------------------
 %% API
@@ -70,6 +70,10 @@ archive_chat_markers(Module, HostType) ->
 add_stanzaid_element(Module, HostType) ->
     not param(Module, HostType, no_stanzaid_element, false).
 
+-spec extra_fin_element_module(mam_module(), mongooseim:host_type()) -> module() | undefined.
+extra_fin_element_module(Module, HostType) ->
+    param(Module, HostType, extra_fin_element, undefined).
+
 %%--------------------------------------------------------------------
 %% Internal functions
 %%--------------------------------------------------------------------

diff --git a/src/mam/mod_mam_utils.erl b/src/mam/mod_mam_utils.erl
@@ -38,7 +38,7 @@
          result_set/4,
          result_query/2,
          result_prefs/4,
-         make_fin_element/4,
+         make_fin_element/5,
          parse_prefs/1,
          borders_decode/1,
          decode_optimizations/1,
@@ -92,6 +92,9 @@
 %% Shared logic
 -export([check_result_for_policy_violation/2]).
 
+-callback extra_fin_element(exml:element()) -> exml:element().
+-optional_callbacks([extra_fin_element/1]).
+
 -ignore_xref([append_arcid_elem/4, delete_arcid_elem/3, form_field_value/2,
               get_one_of_path/3, is_arcid_elem_for/3, maybe_encode_compact_uuid/2,
               maybe_last/1, result_query/2, send_message/4, wrap_message/7, wrapper_id/0]).
@@ -514,15 +517,20 @@ encode_jids(JIDs) ->
 
 
 %% MAM v0.4.1 and above
--spec make_fin_element(binary(), boolean(), boolean(), exml:element()) -> exml:element().
-make_fin_element(MamNs, IsComplete, IsStable, ResultSetEl) ->
-    #xmlel{
-       name = <<"fin">>,
-       attrs = [{<<"xmlns">>, MamNs}]
-        ++ [{<<"complete">>, <<"true">>} || IsComplete]
-        ++ [{<<"stable">>, <<"false">>} || not IsStable],
-       children = [ResultSetEl]}.
-
+-spec make_fin_element(binary(), boolean(), boolean(), exml:element(), module()) -> exml:element().
+make_fin_element(MamNs, IsComplete, IsStable, ResultSetEl, ExtFinMod) ->
+    FinEl = #xmlel{
+               name = <<"fin">>,
+               attrs = [{<<"xmlns">>, MamNs}]
+               ++ [{<<"complete">>, <<"true">>} || IsComplete]
+               ++ [{<<"stable">>, <<"false">>} || not IsStable],
+               children = [ResultSetEl]},
+    maybe_transform_fin_elem(ExtFinMod, FinEl).
+
+maybe_transform_fin_elem(undefined, FinEl) ->
+    FinEl;
+maybe_transform_fin_elem(Module, FinEl) ->
+    Module:extra_fin_element(FinEl).
 
 -spec parse_prefs(PrefsEl :: exml:element()) -> mod_mam:preference().
 parse_prefs(El = #xmlel{ name = <<"prefs">> }) ->

diff --git a/test/config_parser_SUITE.erl b/test/config_parser_SUITE.erl
@@ -2125,6 +2125,8 @@ test_mod_mam_meta(T, M) ->
          T(#{<<"db_message_format">> => <<"mam_message_xml">>})),
     ?eqf(M([{simple, false}]),
          T(#{<<"simple">> => false})),
+    ?eqf(M([{extra_fin_element, mod_mam_utils}]),
+         T(#{<<"extra_fin_element">> => <<"mod_mam_utils">>})),
     ?eqf(M([{extra_lookup_params, mod_mam_utils}]),
          T(#{<<"extra_lookup_params">> => <<"mod_mam_utils">>})),
     ?errf(T(#{<<"backend">> => <<"notepad">>})),
@@ -2145,6 +2147,7 @@ test_mod_mam_meta(T, M) ->
     ?errf(T(#{<<"db_jid_format">> => <<"not_a_module">>})),
     ?errf(T(#{<<"db_message_format">> => <<"not_a_module">>})),
     ?errf(T(#{<<"simple">> => <<"yes">>})),
+    ?errf(T(#{<<"extra_fin_element">> => <<"bad_module">>})),
     ?errf(T(#{<<"extra_lookup_params">> => <<"bad_module">>})).
 
 mod_muc(_Config) ->