Add remove_domain handler for MAM

big_tests/default.spec

@@ -84,6 +84,7 @@
 {suites, "tests", xep_0352_csi_SUITE}.
 {suites, "tests", service_domain_db_SUITE}.
 {suites, "tests", domain_isolation_SUITE}.
+{suites, "tests", domain_removal_SUITE}.
 
 {config, ["test.config"]}.
 {logdir, "ct_report"}.

big_tests/tests/domain_removal_SUITE.erl

@@ -0,0 +1,114 @@
+-module(domain_removal_SUITE).
+
+%% API
+-export([all/0,
+         groups/0,
+         init_per_suite/1,
+         end_per_suite/1,
+         init_per_group/2,
+         end_per_group/2,
+         init_per_testcase/2,
+         end_per_testcase/2]).
+
+-export([mam_pm_removal/1,
+         mam_muc_removal/1]).
+
+-import(distributed_helper, [mim/0, rpc/4]).
+
+-include("mam_helper.hrl").
+-include_lib("escalus/include/escalus.hrl").
+-include_lib("escalus/include/escalus_xmlns.hrl").
+-include_lib("common_test/include/ct.hrl").
+-include_lib("exml/include/exml_stream.hrl").
+
+all() ->
+    [{group, mam_removal}].
+
+groups() ->
+    [
+     {mam_removal, [], [mam_pm_removal,
+                        mam_muc_removal]}
+    ].
+
+domain() ->
+    ct:get_config({hosts, mim, domain}).
+
+%%%===================================================================
+%%% Overall setup/teardown
+%%%===================================================================
+init_per_suite(Config) ->
+    escalus:init_per_suite(Config).
+
+end_per_suite(Config) ->
+    escalus:end_per_suite(Config).
+
+%%%===================================================================
+%%% Group specific setup/teardown
+%%%===================================================================
+init_per_group(Group, Config) ->
+    case mongoose_helper:is_rdbms_enabled(domain()) of
+        true ->
+            Config2 = dynamic_modules:save_modules(domain(), Config),
+            rpc(mim(), gen_mod_deps, start_modules, [domain(), group_to_modules(Group)]),
+            Config2;
+        false ->
+            {skip, require_rdbms}
+    end.
+
+end_per_group(_Groupname, Config) ->
+    case mongoose_helper:is_rdbms_enabled(domain()) of
+        true ->
+            dynamic_modules:restore_modules(domain(), Config);
+        false ->
+            ok
+    end,
+    ok.
+
+group_to_modules(mam_removal) ->
+    MH = muc_light_helper:muc_host(),
+    [{mod_mam_meta, [{backend, rdbms}, {pm, []}, {muc, [{host, MH}]}]},
+     {mod_muc_light, []}].
+
+%%%===================================================================
+%%% Testcase specific setup/teardown
+%%%===================================================================
+
+init_per_testcase(TestCase, Config) ->
+    escalus:init_per_testcase(TestCase, Config).
+
+end_per_testcase(TestCase, Config) ->
+    escalus:end_per_testcase(TestCase, Config).
+
+%%%===================================================================
+%%% Test Cases
+%%%===================================================================
+
+mam_pm_removal(Config) ->
+    F = fun(Alice, Bob) ->
+        escalus:send(Alice, escalus_stanza:chat_to(Bob, <<"OH, HAI!">>)),
+        escalus:wait_for_stanza(Bob),
+        mam_helper:wait_for_archive_size(Alice, 1),
+        mam_helper:wait_for_archive_size(Bob, 1),
+        run_remove_domain(),
+        mam_helper:wait_for_archive_size(Alice, 0),
+        mam_helper:wait_for_archive_size(Bob, 0)
+        end,
+    escalus_fresh:story(Config, [{alice, 1}, {bob, 1}], F).
+
+mam_muc_removal(Config0) ->
+    F = fun(Config, Alice) ->
+        Room = muc_helper:fresh_room_name(),
+        MucHost = muc_light_helper:muc_host(),
+        muc_light_helper:create_room(Room, MucHost, alice,
+                                     [alice], Config, muc_light_helper:ver(1)),
+        RoomAddr = <<Room/binary, "@", MucHost/binary>>,
+        escalus:send(Alice, escalus_stanza:groupchat_to(RoomAddr, <<"text">>)),
+        escalus:wait_for_stanza(Alice),
+        mam_helper:wait_for_room_archive_size(MucHost, Room, 1),
+        run_remove_domain(),
+        mam_helper:wait_for_room_archive_size(MucHost, Room, 0)
+        end,
+    escalus_fresh:story_with_config(Config0, [{alice, 1}], F).
+
+run_remove_domain() ->
+    rpc(mim(), mongoose_hooks, remove_domain, [domain(), domain()]).

priv/mysql.sql

@@ -265,6 +265,7 @@ CREATE TABLE mam_config(
 ) CHARACTER SET utf8mb4
   ROW_FORMAT=DYNAMIC;
 
+-- The server field is a MUC host for MUC rooms
 CREATE TABLE mam_server_user(
   id INT UNSIGNED NOT NULL AUTO_INCREMENT,
   server    varchar(250) CHARACTER SET binary NOT NULL,

src/auth/ejabberd_auth.erl

@@ -81,15 +81,15 @@ start() ->
 -spec start(HostType :: binary()) -> 'ok'.
 start(HostType) ->
     ensure_metrics(HostType),
-    ejabberd_hooks:add(remove_domain, global, fun remove_domain/3, 50),
+    ejabberd_hooks:add(remove_domain, HostType, ?MODULE, remove_domain, 50),
     lists:foreach(
       fun(M) ->
               M:start(HostType)
       end, auth_modules_for_host_type(HostType)).
 
 -spec stop(HostType :: binary()) -> 'ok'.
 stop(HostType) ->
-    ejabberd_hooks:delete(remove_domain, global, fun remove_domain/3, 50),
+    ejabberd_hooks:delete(remove_domain, HostType, ?MODULE, remove_domain, 50),
     lists:foreach(
       fun(M) ->
               M:stop(HostType)

src/mam/mod_mam_muc_rdbms_arch.erl

@@ -23,7 +23,8 @@
 -export([archive_size/4,
          archive_message/3,
          lookup_messages/3,
-         remove_archive/4]).
+         remove_archive/4,
+         remove_domain/3]).
 
 -export([get_mam_muc_gdpr_data/2]).
 
@@ -77,32 +78,25 @@ get_mam_muc_gdpr_data(Acc, #jid{luser = User, lserver = Host} = _UserJID) ->
 
 -spec start_hooks(jid:server(), _) -> ok.
 start_hooks(Host, _Opts) ->
-    case gen_mod:get_module_opt(Host, ?MODULE, no_writer, false) of
-        true ->
-            ok;
-        false ->
-            ejabberd_hooks:add(mam_muc_archive_message, Host, ?MODULE, archive_message, 50)
-    end,
-    ejabberd_hooks:add(mam_muc_archive_size, Host, ?MODULE, archive_size, 50),
-    ejabberd_hooks:add(mam_muc_lookup_messages, Host, ?MODULE, lookup_messages, 50),
-    ejabberd_hooks:add(mam_muc_remove_archive, Host, ?MODULE, remove_archive, 50),
-    ejabberd_hooks:add(get_mam_muc_gdpr_data, Host, ?MODULE, get_mam_muc_gdpr_data, 50),
-    ok.
-
+    ejabberd_hooks:add(hooks(Host)).
 
 -spec stop_hooks(jid:server()) -> ok.
 stop_hooks(Host) ->
-    case gen_mod:get_module_opt(Host, ?MODULE, no_writer, false) of
+    ejabberd_hooks:delete(hooks(Host)).
+
+hooks(Host) ->
+    NoWriter = gen_mod:get_module_opt(Host, ?MODULE, no_writer, false),
+    case NoWriter of
         true ->
-            ok;
+            [];
         false ->
-            ejabberd_hooks:delete(mam_muc_archive_message, Host, ?MODULE, archive_message, 50)
-    end,
-    ejabberd_hooks:delete(mam_muc_archive_size, Host, ?MODULE, archive_size, 50),
-    ejabberd_hooks:delete(mam_muc_lookup_messages, Host, ?MODULE, lookup_messages, 50),
-    ejabberd_hooks:delete(mam_muc_remove_archive, Host, ?MODULE, remove_archive, 50),
-    ejabberd_hooks:delete(get_mam_muc_gdpr_data, Host, ?MODULE, get_mam_muc_gdpr_data, 50),
-    ok.
+            [{mam_muc_archive_message, Host, ?MODULE, archive_message, 50}]
+    end ++
+    [{mam_muc_archive_size, Host, ?MODULE, archive_size, 50},
+     {mam_muc_lookup_messages, Host, ?MODULE, lookup_messages, 50},
+     {mam_muc_remove_archive, Host, ?MODULE, remove_archive, 50},
+     {remove_domain, Host, ?MODULE, remove_domain, 50},
+     {get_mam_muc_gdpr_data, Host, ?MODULE, get_mam_muc_gdpr_data, 50}].
 
 %% ----------------------------------------------------------------------
 %% SQL queries
@@ -112,6 +106,11 @@ register_prepared_queries() ->
     mongoose_rdbms:prepare(mam_muc_archive_remove, mam_muc_message, [room_id],
                            <<"DELETE FROM mam_muc_message "
                              "WHERE room_id = ?">>),
+    mongoose_rdbms:prepare(mam_muc_remove_domain, mam_muc_message, ['mam_server_user.server'],
+                           <<"DELETE FROM mam_muc_message "
+                             "WHERE room_id IN (SELECT id FROM mam_server_user where server = ?)">>),
+    mongoose_rdbms:prepare(mam_muc_remove_domain_users, mam_server_user, [server],
+                           <<"DELETE FROM mam_server_user WHERE server = ?">>),
     mongoose_rdbms:prepare(mam_muc_make_tombstone, mam_muc_message, [message, room_id, id],
                            <<"UPDATE mam_muc_message SET message = ?, search_body = '' "
                              "WHERE room_id = ? AND id = ?">>),
@@ -290,6 +289,25 @@ remove_archive(Acc, Host, ArcID, _ArcJID) ->
     mongoose_rdbms:execute_successfully(Host, mam_muc_archive_remove, [ArcID]),
     Acc.
 
+-spec remove_domain(mongoose_acc:t(), binary(), jid:lserver()) -> mongoose_acc:t().
+remove_domain(Acc, HostType, Domain) ->
+    SubHosts = get_subhosts(Domain),
+    {atomic, _} = mongoose_rdbms:sql_transaction(HostType, fun() ->
+            [remove_domain_trans(HostType, SubHost) || SubHost <- SubHosts]
+        end),
+    Acc.
+
+remove_domain_trans(HostType, MucHost) ->
+    mongoose_rdbms:execute_successfully(HostType, mam_muc_remove_domain, [MucHost]),
+    mongoose_rdbms:execute_successfully(HostType, mam_muc_remove_domain_users, [MucHost]).
+
+get_subhosts(Domain) ->
+    MucHost = gen_mod:get_module_opt_subhost(Domain, mod_muc_light,
+                                        mod_muc_light:default_host()),
+    LightHost = gen_mod:get_module_opt_subhost(Domain, mod_muc,
+                                        mod_muc:default_host()),
+    lists:usort([MucHost, LightHost]).
+
 %% GDPR logic
 extract_gdpr_messages(Host, SenderID) ->
     mongoose_rdbms:execute_successfully(Host, mam_muc_extract_gdpr_messages, [SenderID]).

src/mam/mod_mam_rdbms_arch.erl

@@ -23,7 +23,8 @@
 -export([archive_size/4,
          archive_message/3,
          lookup_messages/3,
-         remove_archive/4]).
+         remove_archive/4,
+         remove_domain/3]).
 
 -export([get_mam_pm_gdpr_data/2]).
 
@@ -93,32 +94,25 @@ uniform_to_gdpr({MessID, RemoteJID, Packet}) ->
 
 -spec start_hooks(jid:server(), _) -> ok.
 start_hooks(Host, _Opts) ->
-    case gen_mod:get_module_opt(Host, ?MODULE, no_writer, false) of
-        true ->
-            ok;
-        false ->
-            ejabberd_hooks:add(mam_archive_message, Host, ?MODULE, archive_message, 50)
-    end,
-    ejabberd_hooks:add(mam_archive_size, Host, ?MODULE, archive_size, 50),
-    ejabberd_hooks:add(mam_lookup_messages, Host, ?MODULE, lookup_messages, 50),
-    ejabberd_hooks:add(mam_remove_archive, Host, ?MODULE, remove_archive, 50),
-    ejabberd_hooks:add(get_mam_pm_gdpr_data, Host, ?MODULE, get_mam_pm_gdpr_data, 50),
-    ok.
-
+    ejabberd_hooks:add(hooks(Host)).
 
 -spec stop_hooks(jid:server()) -> ok.
 stop_hooks(Host) ->
-    case gen_mod:get_module_opt(Host, ?MODULE, no_writer, false) of
+    ejabberd_hooks:delete(hooks(Host)).
+
+hooks(Host) ->
+    NoWriter = gen_mod:get_module_opt(Host, ?MODULE, no_writer, false),
+    case NoWriter of
         true ->
-            ok;
+            [];
         false ->
-            ejabberd_hooks:delete(mam_archive_message, Host, ?MODULE, archive_message, 50)
-    end,
-    ejabberd_hooks:delete(mam_archive_size, Host, ?MODULE, archive_size, 50),
-    ejabberd_hooks:delete(mam_lookup_messages, Host, ?MODULE, lookup_messages, 50),
-    ejabberd_hooks:delete(mam_remove_archive, Host, ?MODULE, remove_archive, 50),
-    ejabberd_hooks:delete(get_mam_pm_gdpr_data, Host, ?MODULE, get_mam_pm_gdpr_data, 50),
-    ok.
+            [{mam_archive_message, Host, ?MODULE, archive_message, 50}]
+    end ++
+    [{mam_archive_size, Host, ?MODULE, archive_size, 50},
+     {mam_lookup_messages, Host, ?MODULE, lookup_messages, 50},
+     {mam_remove_archive, Host, ?MODULE, remove_archive, 50},
+     {remove_domain, Host, ?MODULE, remove_domain, 50},
+     {get_mam_pm_gdpr_data, Host, ?MODULE, get_mam_pm_gdpr_data, 50}].
 
 %% ----------------------------------------------------------------------
 %% SQL queries
@@ -128,6 +122,16 @@ register_prepared_queries() ->
     mongoose_rdbms:prepare(mam_archive_remove, mam_message, [user_id],
                            <<"DELETE FROM mam_message "
                              "WHERE user_id = ?">>),
+    mongoose_rdbms:prepare(mam_remove_domain, mam_message, ['mam_server_user.server'],
+                           <<"DELETE FROM mam_message "
+                             "WHERE user_id IN "
+                             "(SELECT id from mam_server_user WHERE server = ?)">>),
+    mongoose_rdbms:prepare(mam_remove_domain_prefs, mam_config, ['mam_server_user.server'],
+                           <<"DELETE FROM mam_config "
+                             "WHERE user_id IN "
+                             "(SELECT id from mam_server_user WHERE server = ?)">>),
+    mongoose_rdbms:prepare(mam_remove_domain_users, mam_server_user, [server],
+                           <<"DELETE FROM mam_server_user WHERE server = ?">>),
     mongoose_rdbms:prepare(mam_make_tombstone, mam_message, [message, user_id, id],
                            <<"UPDATE mam_message SET message = ?, search_body = '' "
                              "WHERE user_id = ? AND id = ?">>),
@@ -314,6 +318,15 @@ remove_archive(Acc, Host, ArcID, _ArcJID) ->
     mongoose_rdbms:execute_successfully(Host, mam_archive_remove, [ArcID]),
     Acc.
 
+-spec remove_domain(mongoose_acc:t(), binary(), jid:lserver()) -> mongoose_acc:t().
+remove_domain(Acc, HostType, Domain) ->
+    {atomic, _} = mongoose_rdbms:sql_transaction(HostType, fun() ->
+            mongoose_rdbms:execute_successfully(HostType, mam_remove_domain, [Domain]),
+            mongoose_rdbms:execute_successfully(HostType, mam_remove_domain_prefs, [Domain]),
+            mongoose_rdbms:execute_successfully(HostType, mam_remove_domain_users, [Domain])
+        end),
+    Acc.
+
 %% GDPR logic
 extract_gdpr_messages(Env, ArcID) ->
     Filters = [{equal, user_id, ArcID}],

src/mongoose_hooks.erl

@@ -238,7 +238,7 @@ disable_domain(HostType, Domain) ->
     Domain :: jid:lserver(),
     Result :: ok.
 remove_domain(HostType, Domain) ->
-    ejabberd_hooks:run_global(remove_domain, ok, [HostType, Domain]).
+    ejabberd_hooks:run_for_host_type(remove_domain, HostType, #{}, [HostType, Domain]).
 
 -spec node_cleanup(Node :: node()) -> Acc :: map().
 node_cleanup(Node) ->