esl · arcusfelis · Aug 12, 2022 · Aug 12, 2022 · Aug 12, 2022
diff --git a/big_tests/tests/mam_SUITE.erl b/big_tests/tests/mam_SUITE.erl
@@ -37,6 +37,7 @@
          easy_archive_request/1,
          easy_archive_request_for_the_receiver/1,
          text_search_query_fails_if_disabled/1,
+         pagination_simple_enforced/1,
          text_search_is_not_available/1,
          easy_text_search_request/1,
          long_text_search_request/1,
@@ -282,6 +283,7 @@ basic_group_names() ->
      prefs_cases,
      impl_specific,
      disabled_text_search,
+     disabled_complex_queries,
      disabled_retraction
     ].
 
@@ -349,6 +351,8 @@ basic_groups() ->
      {impl_specific,    [], impl_specific()},
      {disabled_text_search, [],
          [{mam04, [], disabled_text_search_cases()}]},
+     {disabled_complex_queries, [],
+         [{mam04, [], disabled_complex_queries_cases()}]},
      {chat_markers, [parallel],
          [{mam04, [parallel], chat_markers_cases()}]},
      {disabled_retraction, [],
@@ -392,6 +396,11 @@ disabled_text_search_cases() ->
      text_search_query_fails_if_disabled
     ].
 
+disabled_complex_queries_cases() ->
+    [
+     pagination_simple_enforced
+    ].
+
 muc_text_search_cases() ->
     [
      muc_text_search_request
@@ -707,6 +716,8 @@ muc_domain(Config) ->
 
 mam_opts_for_base_group(disabled_text_search) ->
     #{full_text_search => false};
+mam_opts_for_base_group(disabled_complex_queries) ->
+    #{enforce_simple_queries => true};
 mam_opts_for_base_group(BG) when BG =:= disabled_retraction;
                                  BG =:= muc_disabled_retraction ->
     #{message_retraction => false};
@@ -844,6 +855,9 @@ init_per_testcase(C=muc_show_x_user_to_moderators_in_anon_rooms, Config) ->
 init_per_testcase(C=muc_show_x_user_for_your_own_messages_in_anon_rooms, Config) ->
     Config1 = escalus_fresh:create_users(Config, [{alice, 1}, {bob, 1}]),
     escalus:init_per_testcase(C, start_alice_anonymous_room(Config1));
+init_per_testcase(C=pagination_simple_enforced, Config) ->
+    Config1 = escalus_fresh:create_users(Config, [{alice, 1}, {bob, 1}, {carol, 1}]),
+    escalus:init_per_testcase(C, bootstrap_archive(Config1));
 init_per_testcase(C=range_archive_request_not_empty, Config) ->
     Config1 = escalus_fresh:create_users(Config, [{alice, 1}, {bob, 1}, {carol, 1}]),
     escalus:init_per_testcase(C, bootstrap_archive(Config1));
@@ -1225,6 +1239,42 @@ text_search_query_fails_if_disabled(Config) ->
         end,
     escalus_fresh:story(Config, [{alice, 1}, {bob, 1}], F).
 
+pagination_simple_enforced(Config) ->
+    P = ?config(props, Config),
+    F = fun(Alice) ->
+        Msgs = ?config(pre_generated_msgs, Config),
+        [_, _, StartMsg, StopMsg | _] = Msgs,
+        {{StartMsgId, _}, _, _, _, _StartMsgPacket} = StartMsg,
+        {{StopMsgId, _}, _, _, _, _StopMsgPacket} = StopMsg,
+        {StartMicro, _} = rpc_apply(mod_mam_utils, decode_compact_uuid, [StartMsgId]),
+        {StopMicro, _} = rpc_apply(mod_mam_utils, decode_compact_uuid, [StopMsgId]),
+        StartTime = make_iso_time(StartMicro),
+        StopTime = make_iso_time(StopMicro),
+        %% Send
+        %% <iq type='get'>
+        %%   <query xmlns='urn:xmpp:mam:tmp'>
+        %%     <start>StartTime</start>
+        %%     <end>StopTime</end>
+        %%   </query>
+        %% </iq>
+        escalus:send(Alice, stanza_date_range_archive_request_not_empty(P, StartTime, StopTime)),
+        %% Receive two messages and IQ
+        Result = wait_archive_respond(Alice),
+        IQ = respond_iq(Result),
+        [M1, M2|_] = respond_messages(Result),
+        escalus:assert(is_iq_result, IQ),
+        SetEl = exml_query:path(IQ, [{element, <<"fin">>}, {element, <<"set">>}]),
+        ?assert_equal(true, undefined =/= SetEl),
+        ?assert_equal(undefined, exml_query:path(SetEl, [{element, <<"count">>}])),
+        ?assert_equal(undefined, exml_query:path(SetEl, [{element, <<"first">>}, {attr, <<"index">>}])),
+        #forwarded_message{delay_stamp = Stamp1} = parse_forwarded_message(M1),
+        #forwarded_message{delay_stamp = Stamp2} = parse_forwarded_message(M2),
+        ?assert_equal(list_to_binary(StartTime), Stamp1),
+        ?assert_equal(list_to_binary(StopTime), Stamp2)
+        end,
+    %% Made fresh in init_per_testcase
+    escalus:story(Config, [{alice, 1}], F).
+
 text_search_is_available(Config) ->
     P = ?config(props, Config),
     F = fun(Alice) ->

diff --git a/doc/modules/mod_mam.md b/doc/modules/mod_mam.md
@@ -352,6 +352,14 @@ This sets the default page size of returned results.
 
 This sets the maximum page size of returned results.
 
+#### `modules.mod_mam.enforce_simple_queries`
+* **Syntax:** boolean
+* **Default:** `false`
+* **Example:** `modules.mod_mam.enforce_simple_queries = true`
+
+This enforces all mam lookups to be "simple", i.e., they skip the RSM count.
+See [Message Archive Management extensions](../open-extensions/mam.md).
+
 #### `modules.mod_mam.db_jid_format`
 
 * **Syntax:** string, one of `"mam_jid_rfc"`, `"mam_jid_rfc_trust"`, `"mam_jid_mini"` or a module implementing `mam_jid` behaviour

diff --git a/src/mam/mam_iq.erl b/src/mam/mam_iq.erl
@@ -2,7 +2,7 @@
 
 -export([action/1]).
 -export([action_type/1]).
--export([form_to_lookup_params/4]).
+-export([form_to_lookup_params/5]).
 -export([lookup_params_with_archive_details/4]).
 
 -import(mod_mam_utils,
@@ -121,9 +121,9 @@ maybe_jid(<<>>) ->
 maybe_jid(JID) when is_binary(JID) ->
     jid:from_binary(JID).
 
--spec form_to_lookup_params(jlib:iq(), integer(), integer(), undefined | module()) ->
+-spec form_to_lookup_params(jlib:iq(), integer(), integer(), undefined | module(), boolean()) ->
     lookup_params().
-form_to_lookup_params(#iq{sub_el = QueryEl} = IQ, MaxResultLimit, DefaultResultLimit, Module) ->
+form_to_lookup_params(#iq{sub_el = QueryEl} = IQ, MaxResultLimit, DefaultResultLimit, Module, EnforceSimple) ->
     Params0 = common_lookup_params(QueryEl, MaxResultLimit, DefaultResultLimit),
     Params = Params0#{
                %% Filtering by date.
@@ -144,7 +144,7 @@ form_to_lookup_params(#iq{sub_el = QueryEl} = IQ, MaxResultLimit, DefaultResultL
                %% - true - do not count records (useful during pagination, when we already
                %%          know how many messages we have from a previous query);
                %% - false - count messages (slow, according XEP-0313);
-               is_simple => mod_mam_utils:form_decode_optimizations(QueryEl)},
+               is_simple => maybe_enforce_simple(QueryEl, EnforceSimple)},
     maybe_add_extra_lookup_params(Module, Params, IQ).
 
 -spec common_lookup_params(exml:element(), non_neg_integer(), non_neg_integer()) ->
@@ -175,3 +175,7 @@ maybe_add_extra_lookup_params(undefined, Params, _) ->
 maybe_add_extra_lookup_params(Module, Params, IQ) ->
     Module:extra_lookup_params(IQ, Params).
 
+maybe_enforce_simple(_, true) ->
+    true;
+maybe_enforce_simple(QueryEl, _) ->
+    mod_mam_utils:form_decode_optimizations(QueryEl).
diff --git a/src/mam/mod_mam.erl b/src/mam/mod_mam.erl
@@ -115,7 +115,8 @@ config_spec() ->
                     <<"full_text_search">> => true,
                     <<"cache_users">> => true,
                     <<"default_result_limit">> => 50,
-                    <<"max_result_limit">> => 50},
+                    <<"max_result_limit">> => 50,
+                    <<"enforce_simple_queries">> => false},
        process = fun ?MODULE:remove_unused_backend_opts/1
       }.
 
@@ -160,6 +161,7 @@ common_config_items() ->
       %% Low-level options
       <<"default_result_limit">> => #option{type = integer,
                                             validate = non_negative},
+      <<"enforce_simple_queries">> => #option{type = boolean},
       <<"max_result_limit">> => #option{type = integer,
                                         validate = non_negative},
       <<"db_jid_format">> => #option{type = atom,
@@ -262,6 +264,7 @@ common_opts() ->
      message_retraction,
      default_result_limit,
      max_result_limit,
+     enforce_simple_queries,
      no_stanzaid_element].
 
 -spec add_prefs_store_module(mam_backend(), mam_type(), module_opts(), module_map()) -> module_map().

diff --git a/src/mam/mod_mam_muc.erl b/src/mam/mod_mam_muc.erl
@@ -387,7 +387,8 @@ handle_set_message_form(HostType, #jid{} = From, #jid{} = ArcJID, IQ) ->
     ResLimit = mod_mam_params:max_result_limit(?MODULE, HostType),
     DefLimit = mod_mam_params:default_result_limit(?MODULE, HostType),
     ExtMod = mod_mam_params:extra_params_module(?MODULE, HostType),
-    try mam_iq:form_to_lookup_params(IQ, ResLimit, DefLimit, ExtMod) of
+    Sim = mod_mam_params:enforce_simple_queries(?MODULE, HostType),
+    try mam_iq:form_to_lookup_params(IQ, ResLimit, DefLimit, ExtMod, Sim) of
         Params0 ->
             do_handle_set_message_form(HostType, From, ArcID, ArcJID, IQ, Params0)
     catch _C:R:S ->

diff --git a/src/mam/mod_mam_params.erl b/src/mam/mod_mam_params.erl
@@ -20,7 +20,8 @@
 
 -export([extra_params_module/2, max_result_limit/2, default_result_limit/2,
          has_full_text_search/2, is_archivable_message_module/2, send_message_mod/2,
-         archive_chat_markers/2, add_stanzaid_element/2, extra_fin_element_module/2]).
+         archive_chat_markers/2, add_stanzaid_element/2, extra_fin_element_module/2,
+         enforce_simple_queries/2]).
 
 %%--------------------------------------------------------------------
 %% API
@@ -38,6 +39,9 @@ max_result_limit(Module, HostType) ->
 default_result_limit(Module, HostType) ->
     gen_mod:get_module_opt(HostType, Module, default_result_limit).
 
+-spec enforce_simple_queries(mam_module(), mongooseim:host_type()) -> module() | undefined.
+enforce_simple_queries(Module, HostType) ->
+    gen_mod:get_module_opt(HostType, Module, enforce_simple_queries, false).
 
 -spec has_full_text_search(Module :: mod_mam_pm | mod_mam_muc, mongooseim:host_type()) -> boolean().
 has_full_text_search(Module, HostType) ->

diff --git a/src/mam/mod_mam_pm.erl b/src/mam/mod_mam_pm.erl
@@ -419,7 +419,8 @@ iq_to_lookup_params(HostType, IQ) ->
     Max = mod_mam_params:max_result_limit(?MODULE, HostType),
     Def = mod_mam_params:default_result_limit(?MODULE, HostType),
     Ext = mod_mam_params:extra_params_module(?MODULE, HostType),
-    mam_iq:form_to_lookup_params(IQ, Max, Def, Ext).
+    Sim = mod_mam_params:enforce_simple_queries(?MODULE, HostType),
+    mam_iq:form_to_lookup_params(IQ, Max, Def, Ext, Sim).
 
 forward_messages(HostType, From, ArcJID, MamNs, QueryID, MessageRows, SetClientNs) ->
     %% Forward messages

diff --git a/test/common/config_parser_helper.erl b/test/common/config_parser_helper.erl
@@ -1295,6 +1295,7 @@ common_mam_config() ->
       full_text_search => true,
       default_result_limit => 50,
       max_result_limit => 50,
+      enforce_simple_queries => false,
       async_writer => default_config([modules, mod_mam, async_writer])}.
 
 mod_event_pusher_http_handler() ->

diff --git a/test/config_parser_SUITE.erl b/test/config_parser_SUITE.erl
@@ -2049,6 +2049,8 @@ test_mod_mam(P, T) ->
           T(#{<<"default_result_limit">> => 100})),
     ?cfgh(P ++ [max_result_limit], 1000,
           T(#{<<"max_result_limit">> => 1000})),
+    ?cfgh(P ++ [enforce_simple_queries], true,
+          T(#{<<"enforce_simple_queries">> => true})),
     ?cfgh(P ++ [db_jid_format], mam_jid_rfc,
           T(#{<<"db_jid_format">> => <<"mam_jid_rfc">>})),
     ?cfgh(P ++ [db_message_format], mam_message_xml,
@@ -2067,6 +2069,7 @@ test_mod_mam(P, T) ->
     ?errh(T(#{<<"cache_users">> => []})),
     ?errh(T(#{<<"default_result_limit">> => -1})),
     ?errh(T(#{<<"max_result_limit">> => -2})),
+    ?errh(T(#{<<"enforce_simple_queries">> => -2})),
     ?errh(T(#{<<"db_jid_format">> => <<"not_a_module">>})),
     ?errh(T(#{<<"db_message_format">> => <<"not_a_module">>})),
     ?errh(T(#{<<"extra_fin_element">> => <<"bad_module">>})),