prosody/plugins/mod_storage_memory.lua

local serialize = require "prosody.util.serialization".serialize;
local array = require "prosody.util.array";
local envload = require "prosody.util.envload".envload;
local st = require "prosody.util.stanza";
local is_stanza = st.is_stanza or function (s) return getmetatable(s) == st.stanza_mt end
local new_id = require "prosody.util.id".medium;
local set = require "prosody.util.set";

local auto_purge_enabled = module:get_option_boolean("storage_memory_temporary", false);
local auto_purge_stores = module:get_option_set("storage_memory_temporary_stores", {});

local archive_item_limit = module:get_option_integer("storage_archive_item_limit", 1000, 0);

local memory = setmetatable({}, {
	__index = function(t, k)
		local store = module:shared(k)
		t[k] = store;
		return store;
	end
});

local function NULL() return nil end

local function _purge_store(self, username)
	self.store[username or NULL] = nil;
	return true;
end

local function _users(self)
	return next, self.store, nil;
end

local keyval_store = {};
keyval_store.__index = keyval_store;

function keyval_store:get(username)
	return (self.store[username or NULL] or NULL)();
end

function keyval_store:set(username, data)
	if data ~= nil then
		data = envload("return "..serialize(data), "=(data)", {});
	end
	self.store[username or NULL] = data;
	return true;
end

keyval_store.purge = _purge_store;

keyval_store.users = _users;

local archive_store = {};
archive_store.__index = archive_store;

archive_store.users = _users;

archive_store.caps = {
	total = true;
	quota = archive_item_limit;
	truncate = true;
	full_id_range = true;
	ids = true;
};

function archive_store:append(username, key, value, when, with)
	if is_stanza(value) then
		value = st.preserialize(value);
		value = envload("return xml"..serialize(value), "=(stanza)", { xml = st.deserialize })
	else
		value = envload("return "..serialize(value), "=(data)", {});
	end
	local a = self.store[username or NULL];
	if not a then
		a = {};
		self.store[username or NULL] = a;
	end
	local v = { key = key, when = when, with = with, value = value };
	if not key then
		key = new_id();
		v.key = key;
	end
	if a[key] then
		table.remove(a, a[key]);
	elseif #a >= archive_item_limit then
		return nil, "quota-limit";
	end
	local i = #a+1;
	a[i] = v;
	a[key] = i;
	return key;
end

function archive_store:find(username, query)
	local items = self.store[username or NULL];
	if not items then
		if 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 = nil;
	local i, last_key = 0;
	if query then
		items = array():append(items);
		if query.key then
			items:filter(function (item)
				return item.key == query.key;
			end);
		end
		if query.ids then
			local ids = set.new(query.ids);
			items:filter(function (item)
				return ids:contains(item.key);
			end);
		end
		if query.with then
			items:filter(function (item)
				return item.with == query.with;
			end);
		end
		if query.start then
			items:filter(function (item)
				return item.when >= query.start;
			end);
		end
		if query["end"] then
			items:filter(function (item)
				return item.when <= query["end"];
			end);
		end
		if query.total then
			count = #items;
		end
		if query.reverse then
			items:reverse();
			if query.before then
				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
			last_key = query.after;
		elseif query.after then
			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
			last_key = query.before;
		elseif query.before then
			last_key = query.before;
		end
		if query.limit and #items - i > query.limit then
			items[i+query.limit+1] = nil;
		end
	end
	return function ()
		i = i + 1;
		local item = items[i];
		if not item or (last_key and item.key == last_key) then return; end
		return item.key, item.value(), item.when, item.with;
	end, count;
end

function archive_store:get(username, wanted_key)
	local items = self.store[username or NULL];
	if not items then return nil, "item-not-found"; end
	local i = items[wanted_key];
	if not i then return nil, "item-not-found"; end
	local item = items[i];
	return item.value(), item.when, item.with;
end

function archive_store:set(username, wanted_key, new_value, new_when, new_with)
	local items = self.store[username or NULL];
	if not items then return nil, "item-not-found"; end
	local i = items[wanted_key];
	if not i then return nil, "item-not-found"; end
	local item = items[i];

	if is_stanza(new_value) then
		new_value = st.preserialize(new_value);
		item.value = envload("return xml"..serialize(new_value), "=(stanza)", { xml = st.deserialize })
	else
		item.value = envload("return "..serialize(new_value), "=(data)", {});
	end
	if new_when then
		item.when = new_when;
	end
	if new_with then
		item.with = new_when;
	end
	return true;
end

function archive_store:summary(username, query)
	local iter, err = self:find(username, query)
	if not iter then return iter, err; end
	local counts = {};
	local earliest = {};
	local latest = {};
	for _, _, when, with in iter do
		counts[with] = (counts[with] or 0) + 1;
		if earliest[with] == nil then
			earliest[with] = when;
		end
		latest[with] = when;
	end
	return {
		counts = counts;
		earliest = earliest;
		latest = latest;
	};
end


function archive_store:delete(username, query)
	if not query or next(query) == nil then
		self.store[username or NULL] = nil;
		return true;
	end
	local items = self.store[username or NULL];
	if not items then
		-- Store is empty
		return 0;
	end
	items = array(items);
	local count_before = #items;
	if query then
		if query.key then
			items:filter(function (item)
				return item.key ~= query.key;
			end);
		end
		if query.with then
			items:filter(function (item)
				return item.with ~= query.with;
			end);
		end
		if query.start then
			items:filter(function (item)
				return item.when < query.start;
			end);
		end
		if query["end"] then
			items:filter(function (item)
				return item.when > query["end"];
			end);
		end
		if query.truncate and #items > query.truncate then
			if query.reverse then
				-- Before: { 1, 2, 3, 4, 5, }
				-- After: { 1, 2, 3 }
				for i = #items, query.truncate + 1, -1 do
					items[i] = nil;
				end
			else
				-- Before: { 1, 2, 3, 4, 5, }
				-- After: { 3, 4, 5 }
				local offset = #items - query.truncate;
				for i = 1, #items do
					items[i] = items[i+offset];
				end
			end
		end
	end
	local count = count_before - #items;
	if count == 0 then
		return 0; -- No changes, skip write
	end
	setmetatable(items, nil);

	do -- re-index by key
		for k in pairs(items) do
			if type(k) == "string" then
				items[k] = nil;
			end
		end

		for i = 1, #items do
			items[ items[i].key ] = i;
		end
	end

	return count;
end

archive_store.purge = _purge_store;

local stores = {
	keyval = keyval_store;
	archive = archive_store;
}

local driver = {};

function driver:open(store, typ) -- luacheck: ignore 212/self
	local store_mt = stores[typ or "keyval"];
	if store_mt then
		return setmetatable({ store = memory[store] }, store_mt);
	end
	return nil, "unsupported-store";
end

function driver:purge(user) -- luacheck: ignore 212/self
	for _, store in pairs(memory) do
		store[user] = nil;
	end
end

if auto_purge_enabled then
	module:hook("resource-unbind", function (event)
		local user_bare_jid = event.session.username.."@"..event.session.host;
		if not prosody.bare_sessions[user_bare_jid] then -- User went offline
			module:log("debug", "Clearing store for offline user %s", user_bare_jid);
			local f, s, v;
			if auto_purge_stores:empty() then
				f, s, v = pairs(memory);
			else
				f, s, v = auto_purge_stores:items();
			end

			for store_name in f, s, v do
				if memory[store_name] then
					memory[store_name][event.session.username] = nil;
				end
			end
		end
	end);
end

module:provides("storage", driver);