From 3e58495ffe5c6f972063fc134f703f2f622e725e Mon Sep 17 00:00:00 2001
From: Mikhail Uvarov <arcusfelis@gmail.com>
Date: Tue, 15 Jun 2021 12:09:35 +0200
Subject: [PATCH] Fix disco for MAM

---
 big_tests/dynamic_domains.spec | 3 ---
 src/mam/mod_mam.erl            | 4 ++--
 src/mam/mod_mam_muc.erl        | 4 ++--
 src/mam/mod_mam_utils.erl      | 8 ++++----
 4 files changed, 8 insertions(+), 11 deletions(-)

diff --git a/big_tests/dynamic_domains.spec b/big_tests/dynamic_domains.spec
index 696f92b2c92..31eed5401d8 100644
--- a/big_tests/dynamic_domains.spec
+++ b/big_tests/dynamic_domains.spec
@@ -31,9 +31,6 @@
 {suites, "tests", inbox_extensions_SUITE}.
 
 {suites, "tests", mam_SUITE}.
-{skip_cases, "tests", mam_SUITE,
- [muc_service_discovery, mam_service_discovery],
- "at the moment mod_mam config options don't support dynamic domains"}.
 {skip_cases, "tests", mam_SUITE,
  [messages_filtered_when_prefs_default_policy_is_roster],
  "at the moment mod_roster doesn't support dynamic domains"}.
diff --git a/src/mam/mod_mam.erl b/src/mam/mod_mam.erl
index 398a5d5c309..25b6d68ec9e 100644
--- a/src/mam/mod_mam.erl
+++ b/src/mam/mod_mam.erl
@@ -245,8 +245,8 @@ process_mam_iq(Acc, From, To, IQ, _Extra) ->
     end.
 
 -spec disco_local_features(mongoose_disco:feature_acc()) -> mongoose_disco:feature_acc().
-disco_local_features(Acc = #{to_jid := #jid{lserver = LServer}, node := <<>>}) ->
-    mongoose_disco:add_features(features(?MODULE, LServer), Acc);
+disco_local_features(Acc = #{host_type := HostType, node := <<>>}) ->
+    mongoose_disco:add_features(features(?MODULE, HostType), Acc);
 disco_local_features(Acc) ->
     Acc.
 
diff --git a/src/mam/mod_mam_muc.erl b/src/mam/mod_mam_muc.erl
index 4889ec7a887..27c85b15ef6 100644
--- a/src/mam/mod_mam_muc.erl
+++ b/src/mam/mod_mam_muc.erl
@@ -162,8 +162,8 @@ stop(HostType) ->
 %% hooks and handlers for MUC
 
 -spec disco_muc_features(mongoose_disco:feature_acc()) -> mongoose_disco:feature_acc().
-disco_muc_features(Acc = #{to_jid := #jid{lserver = LServer}, node := <<>>}) ->
-    mongoose_disco:add_features(features(?MODULE, LServer), Acc);
+disco_muc_features(Acc = #{host_type := HostType, node := <<>>}) ->
+    mongoose_disco:add_features(features(?MODULE, HostType), Acc);
 disco_muc_features(Acc) ->
     Acc.
 
diff --git a/src/mam/mod_mam_utils.erl b/src/mam/mod_mam_utils.erl
index 6bb7938bfe4..f7e65dddca5 100644
--- a/src/mam/mod_mam_utils.erl
+++ b/src/mam/mod_mam_utils.erl
@@ -634,14 +634,14 @@ maybe_get_result_namespace(Packet) ->
 is_mam_namespace(NS) ->
     lists:member(NS, mam_features()).
 
-features(Module, Host) ->
-    mam_features() ++ retraction_features(Module, Host).
+features(Module, HostType) ->
+    mam_features() ++ retraction_features(Module, HostType).
 
 mam_features() ->
     [?NS_MAM_04, ?NS_MAM_06].
 
-retraction_features(Module, Host) ->
-    case has_message_retraction(Module, Host) of
+retraction_features(Module, HostType) ->
+    case has_message_retraction(Module, HostType) of
         true -> [?NS_RETRACT, ?NS_RETRACT_TOMBSTONE];
         false -> [?NS_RETRACT]
     end.