--- a/plugins/mod_storage_internal.lua	Sat Nov 16 16:45:33 2019 +0100
+++ b/plugins/mod_storage_internal.lua	Sat Nov 16 16:52:31 2019 +0100
@@ -1,12 +1,17 @@
+local cache = require "util.cache";
 local datamanager = require "core.storagemanager".olddm;
 local array = require "util.array";
 local datetime = require "util.datetime";
 local st = require "util.stanza";
 local now = require "util.time".now;
 local id = require "util.id".medium;
+local jid_join = require "util.jid".join;
 
 local host = module.host;
 
+local archive_item_limit = module:get_option_number("storage_archive_item_limit", 10000);
+local archive_item_count_cache = cache.new(module:get_option("storage_archive_item_limit_cache_size", 1000));
+
 local driver = {};
 
 function driver:open(store, typ)
@@ -43,6 +48,12 @@
 local archive = {};
 driver.archive = { __index = archive };
 
+archive.caps = {
+	total = true;
+	quota = archive_item_limit;
+	truncate = true;
+};
+
 function archive:append(username, key, value, when, with)
 	when = when or now();
 	if not st.is_stanza(value) then
@@ -54,28 +65,57 @@
 	value.attr.stamp = datetime.datetime(when);
 	value.attr.stamp_legacy = datetime.legacy(when);
 
+	local cache_key = jid_join(username, host, self.store);
+	local item_count = archive_item_count_cache:get(cache_key);
+
 	if key then
 		local items, err = datamanager.list_load(username, host, self.store);
 		if not items and err then return items, err; end
+
+		-- Check the quota
+		item_count = items and #items or 0;
+		archive_item_count_cache:set(cache_key, item_count);
+		if item_count >= archive_item_limit then
+			module:log("debug", "%s reached or over quota, not adding to store", username);
+			return nil, "quota-limit";
+		end
+
 		if items then
+			-- Filter out any item with the same key as the one being added
 			items = array(items);
 			items:filter(function (item)
 				return item.key ~= key;
 			end);
+
 			value.key = key;
 			items:push(value);
 			local ok, err = datamanager.list_store(username, host, self.store, items);
 			if not ok then return ok, err; end
+			archive_item_count_cache:set(cache_key, #items);
 			return key;
 		end
 	else
+		if not item_count then -- Item count not cached?
+			-- We need to load the list to get the number of items currently stored
+			local items, err = datamanager.list_load(username, host, self.store);
+			if not items and err then return items, err; end
+			item_count = items and #items or 0;
+			archive_item_count_cache:set(cache_key, item_count);
+		end
+		if item_count >= archive_item_limit then
+			module:log("debug", "%s reached or over quota, not adding to store", username);
+			return nil, "quota-limit";
+		end
 		key = id();
 	end
 
+	module:log("debug", "%s has %d items out of %d limit in store %s", username, item_count, archive_item_limit, self.store);
+
 	value.key = key;
 
 	local ok, err = datamanager.list_append(username, host, self.store, value);
 	if not ok then return ok, err; end
+	archive_item_count_cache:set(cache_key, item_count+1);
 	return key;
 end
 
@@ -84,11 +124,17 @@
 	if not items then
 		if err then
 			return items, err;
-		else
-			return function () end, 0;
+		elseif query then
+			if query.before or query.after then
+				return nil, "item-not-found";
+			end
+			if query.total then
+				return function () end, 0;
+			end
 		end
+		return function () end;
 	end
-	local count = #items;
+	local count = nil;
 	local i = 0;
 	if query then
 		items = array(items);
@@ -112,24 +158,36 @@
 				return item.when <= query["end"];
 			end);
 		end
-		count = #items;
+		if query.total then
+			count = #items;
+		end
 		if query.reverse then
 			items:reverse();
 			if query.before then
-				for j = 1, count do
+				local found = false;
+				for j = 1, #items do
 					if (items[j].key or tostring(j)) == query.before then
+						found = true;
 						i = j;
 						break;
 					end
 				end
+				if not found then
+					return nil, "item-not-found";
+				end
 			end
 		elseif query.after then
-			for j = 1, count do
+			local found = false;
+			for j = 1, #items do
 				if (items[j].key or tostring(j)) == query.after then
+					found = true;
 					i = j;
 					break;
 				end
 			end
+			if not found then
+				return nil, "item-not-found";
+			end
 		end
 		if query.limit and #items - i > query.limit then
 			items[i+query.limit+1] = nil;
@@ -156,8 +214,37 @@
 	return array(items):pluck("when"):map(datetime.date):unique();
 end
 
+function archive:summary(username, query)
+	local iter, err = self:find(username, query)
+	if not iter then return iter, err; end
+	local counts = {};
+	local earliest = {};
+	local latest = {};
+	local body = {};
+	for _, stanza, when, with in iter do
+		counts[with] = (counts[with] or 0) + 1;
+		if earliest[with] == nil then
+			earliest[with] = when;
+		end
+		latest[with] = when;
+		body[with] = stanza:get_child_text("body") or body[with];
+	end
+	return {
+		counts = counts;
+		earliest = earliest;
+		latest = latest;
+		body = body;
+	};
+end
+
+function archive:users()
+	return datamanager.users(host, self.store, "list");
+end
+
 function archive:delete(username, query)
+	local cache_key = jid_join(username, host, self.store);
 	if not query or next(query) == nil then
+		archive_item_count_cache:set(cache_key, nil);
 		return datamanager.list_store(username, host, self.store, nil);
 	end
 	local items, err = datamanager.list_load(username, host, self.store);
@@ -165,6 +252,7 @@
 		if err then
 			return items, err;
 		end
+		archive_item_count_cache:set(cache_key, 0);
 		-- Store is empty
 		return 0;
 	end
@@ -214,6 +302,7 @@
 	end
 	local ok, err = datamanager.list_store(username, host, self.store, items);
 	if not ok then return ok, err; end
+	archive_item_count_cache:set(cache_key, #items);
 	return count;
 end