--- a/plugins/mod_mam/mod_mam.lua	Fri Jan 20 23:40:54 2023 +0100
+++ b/plugins/mod_mam/mod_mam.lua	Sat Jan 21 17:18:16 2023 +0100
@@ -142,9 +142,14 @@
 	local qset = rsm.get(query);
 	local qmax = m_min(qset and qset.max or default_max_items, max_max_items);
 	local reverse = qset and qset.before or false;
+
 	local before, after = qset and qset.before or qbefore, qset and qset.after or qafter;
 	if type(before) ~= "string" then before = nil; end
 
+	-- A reverse query needs to be flipped
+	local flip = reverse;
+	-- A flip-page query needs to be the opposite of that.
+	if query:get_child("flip-page") then flip = not flip end
 
 	module:log("debug", "Archive query by %s id=%s with=%s when=%s...%s rsm=%q",
 		origin.username,
@@ -154,11 +159,6 @@
 		qend and timestamp(qend) or "",
 		qset);
 
-	-- A reverse query needs to be flipped
-	local flip = reverse;
-	-- A flip-page query needs to be the opposite of that.
-	if query:get_child("flip-page") then flip = not flip end
-
 	-- Load all the data!
 	local data, err = archive:find(origin.username, {
 		start = qstart; ["end"] = qend; -- Time range

--- a/plugins/mod_muc_mam.lua	Fri Jan 20 23:40:54 2023 +0100
+++ b/plugins/mod_muc_mam.lua	Sat Jan 21 17:18:16 2023 +0100
@@ -67,6 +67,7 @@
 	module:log("info", "See https://prosody.im/doc/storage and https://prosody.im/doc/archiving for more information");
 	return false;
 end
+local use_total = module:get_option_boolean("muc_log_include_total", true);
 
 local function archiving_enabled(room)
 	if log_all_rooms then
@@ -109,10 +110,20 @@
 	{ name = "end"; type = "text-single"; datatype = "xs:dateTime" };
 };
 
+if archive.caps and archive.caps.full_id_range then
+	table.insert(query_form, { name = "before-id"; type = "text-single"; });
+	table.insert(query_form, { name = "after-id"; type = "text-single"; });
+end
+
+if archive.caps and archive.caps.ids then
+	table.insert(query_form, { name = "ids"; type = "list-multi"; });
+end
+
+
 -- Serve form
 module:hook("iq-get/bare/"..xmlns_mam..":query", function(event)
 	local origin, stanza = event.origin, event.stanza;
-	origin.send(st.reply(stanza):tag("query", { xmlns = xmlns_mam }):add_child(query_form:form()));
+	origin.send(st.reply(stanza):query(xmlns_mam):add_child(query_form:form()));
 	return true;
 end);
 
@@ -172,6 +183,7 @@
 
 	local before, after = qset and qset.before or qbefore, qset and qset.after or qafter;
 	if type(before) ~= "string" then before = nil; end
+
 	-- A reverse query needs to be flipped
 	local flip = reverse;
 	-- A flip-page query needs to be the opposite of that.
@@ -187,11 +199,12 @@
 	-- Load all the data!
 	local data, err = archive:find(room_node, {
 		start = qstart; ["end"] = qend; -- Time range
+		with = "message<groupchat";
 		limit = qmax + 1;
 		before = before; after = after;
 		ids = qids;
 		reverse = reverse;
-		with = "message<groupchat";
+		total = use_total or qmax == 0;
 	});
 
 	if not data then
@@ -216,6 +229,8 @@
 	for id, item, when in data do
 		count = count + 1;
 		if count > qmax then
+			-- We requested qmax+1 items. If that many items are retrieved then
+			-- there are more results to page through, so:
 			complete = nil;
 			break;
 		end
@@ -259,7 +274,6 @@
 		first, last = last, first;
 	end
 
-
 	origin.send(st.reply(stanza)
 		:tag("fin", { xmlns = xmlns_mam, complete = complete })
 			:add_child(rsm.generate {
@@ -551,4 +565,7 @@
 
 else
 	module:log("debug", "Archive expiry disabled");
+	-- Don't ask the backend to count the potentially unbounded number of items,
+	-- it'll get slow.
+	use_total = module:get_option_boolean("mam_include_total", false);
 end