prosody/core/storagemanager.lua

local type, pairs = type, pairs;
local setmetatable = setmetatable;
local rawset = rawset;

local config = require "prosody.core.configmanager";
local datamanager = require "prosody.util.datamanager";
local modulemanager = require "prosody.core.modulemanager";
local multitable = require "prosody.util.multitable";
local log = require "prosody.util.logger".init("storagemanager");
local async = require "prosody.util.async";
local debug = debug;

local prosody = prosody;
local hosts = prosody.hosts;

local _ENV = nil;
-- luacheck: std none

local olddm = {}; -- maintain old datamanager, for backwards compatibility
for k,v in pairs(datamanager) do olddm[k] = v; end

local null_storage_method = function () return false, "no data storage active"; end
local null_storage_driver = setmetatable(
	{
		name = "null",
		open = function (self) return self; end
	}, {
		__index = function (self, method) --luacheck: ignore 212
			return null_storage_method;
		end
	}
);

local async_check = config.get("*", "storage_async_check") == true;

local stores_available = multitable.new();

local function check_async_wrapper(event)
	local store = event.store;
	event.store = setmetatable({}, {
		__index = function (t, method_name)
			local original_method = store[method_name];
			if type(original_method) ~= "function" then
				if original_method then
					rawset(t, method_name, original_method);
				end
				return original_method;
			end
			local wrapped_method = function (...)
				if not async.ready() then
					log("warn", "ASYNC-01: Attempt to access storage outside async context, "
					  .."see https://prosody.im/doc/developers/async - %s", debug.traceback());
				end
				return original_method(...);
			end
			rawset(t, method_name, wrapped_method);
			return wrapped_method;
		end;
	});
end

local function initialize_host(host)
	local host_session = hosts[host];
	host_session.events.add_handler("item-added/storage-provider", function (event)
		local item = event.item;
		stores_available:set(host, item.name, item);
	end);

	host_session.events.add_handler("item-removed/storage-provider", function (event)
		local item = event.item;
		stores_available:set(host, item.name, nil);
	end);
	if async_check then
		host_session.events.add_handler("store-opened", check_async_wrapper);
	end
end
prosody.events.add_handler("host-activated", initialize_host, 101);

local function load_driver(host, driver_name)
	if driver_name == "null" then
		return null_storage_driver;
	end
	local driver = stores_available:get(host, driver_name);
	if driver then return driver; end
	local ok, err = modulemanager.load(host, "storage_"..driver_name);
	if not ok then
		log("error", "Failed to load storage driver plugin %s on %s: %s", driver_name, host, err);
	end
	return stores_available:get(host, driver_name);
end

local function get_storage_config(host)
	-- Here used to be some some compat checks
	return config.get(host, "storage");
end

local function get_driver(host, store)
	local storage = get_storage_config(host);
	local driver_name;
	local option_type = type(storage);
	if option_type == "string" then
		driver_name = storage;
	elseif option_type == "table" then
		driver_name = storage[store];
	end
	if not driver_name then
		driver_name = config.get(host, "default_storage") or "internal";
	end

	local driver = load_driver(host, driver_name);
	if not driver then
		log("warn", "Falling back to null driver for %s storage on %s", store, host);
		driver_name = "null";
		driver = null_storage_driver;
	end
	return driver, driver_name;
end

local map_shim_mt = {
	__index = {
		get = function(self, username, key)
			local ret, err = self.keyval_store:get(username);
			if ret == nil then return nil, err end
			return ret[key];
		end;
		set = function(self, username, key, data)
			local current, err = self.keyval_store:get(username);
			if current == nil then
				if err then
					return nil, err;
				else
					current = {};
				end
			end
			current[key] = data;
			return self.keyval_store:set(username, current);
		end;
		set_keys = function (self, username, keydatas)
			local current, err = self.keyval_store:get(username);
			if current == nil then
				if err then
					return nil, err;
				end
				current = {};
			end
			for k,v in pairs(keydatas) do
				if v == self.remove then v = nil; end
				current[k] = v;
			end
			return self.keyval_store:set(username, current);
		end;
		remove = {};
		get_all = function (self, key)
			if type(key) ~= "string" or key == "" then
				return nil, "get_all only supports non-empty string keys";
			end
			local ret;
			for username in self.keyval_store:users() do
				local key_data = self:get(username, key);
				if key_data then
					if not ret then
						ret = {};
					end
					ret[username] = key_data;
				end
			end
			return ret;
		end;
		delete_all = function (self, key)
			if type(key) ~= "string" or key == "" then
				return nil, "delete_all only supports non-empty string keys";
			end
			local data = { [key] = self.remove };
			local last_err;
			for username in self.keyval_store:users() do
				local ok, err = self:set_keys(username, data);
				if not ok then
					last_err = err;
				end
			end
			if last_err then
				return nil, last_err;
			end
			return true;
		end;
	};
}

local combined_store_mt = {
	__index = {
		-- keyval
		get = function (self, name)
			return self.keyval_store:get(name);
		end;
		set = function (self, name, data)
			return self.keyval_store:set(name, data);
		end;
		items = function (self)
			return self.keyval_store:users();
		end;
		-- map
		get_key = function (self, name, key)
			return self.map_store:get(name, key);
		end;
		set_key = function (self, name, key, value)
			return self.map_store:set(name, key, value);
		end;
		set_keys = function (self, name, map)
			return self.map_store:set_keys(name, map);
		end;
		get_key_from_all = function (self, key)
			return self.map_store:get_all(key);
		end;
		delete_key_from_all = function (self, key)
			return self.map_store:delete_all(key);
		end;
	};
};

local open; -- forward declaration

local function create_map_shim(host, store)
	local keyval_store, err = open(host, store, "keyval");
	if keyval_store == nil then return nil, err end
	return setmetatable({
		keyval_store = keyval_store;
	}, map_shim_mt);
end

local function open_combined(host, store)
	local driver, driver_name = get_driver(host, store);

	-- Open keyval
	local keyval_store, err = driver:open(store, "keyval");
	if not keyval_store then
		if err == "unsupported-store" then
			log("debug", "Storage driver %s does not support store %s (keyval), falling back to null driver",
				driver_name, store);
			keyval_store, err = null_storage_driver, nil;
		end
	end

	local map_store;
	if keyval_store then
		-- Open map
		map_store, err = driver:open(store, "map");
		if not map_store then
			if err == "unsupported-store" then
				log("debug", "Storage driver %s does not support store %s (map), falling back to shim",
					driver_name, store);
				map_store, err = setmetatable({ keyval_store = keyval_store }, map_shim_mt), nil;
			end
		end
	end

	if not(keyval_store and map_store) then
		return nil, err;
	end
	local combined_store = setmetatable({
		keyval_store = keyval_store;
		map_store = map_store;
		remove = map_store.remove;
	}, combined_store_mt);
	local event_data = { host = host, store_name = store, store_type = "keyval+", store = combined_store };
	hosts[host].events.fire_event("store-opened", event_data);
	return event_data.store, event_data.store_err;
end

function open(host, store, typ)
	if typ == "keyval+" then -- TODO: default in some release?
		return open_combined(host, store);
	end
	local driver, driver_name = get_driver(host, store);
	local ret, err = driver:open(store, typ);
	if not ret then
		if err == "unsupported-store" then
			if typ == "map" then -- Use shim on top of keyval store
				log("debug", "map storage driver unavailable, using shim on top of keyval store.");
				ret, err = create_map_shim(host, store);
			else
				log("debug", "Storage driver %s does not support store %s (%s), falling back to null driver",
					driver_name, store, typ or "<nil>");
				ret, err = null_storage_driver, nil;
			end
		end
	end
	if ret then
		local event_data = { host = host, store_name = store, store_type = typ, store = ret };
		hosts[host].events.fire_event("store-opened", event_data);
		ret, err = event_data.store, event_data.store_err;
	end
	return ret, err;
end

local function purge(user, host)
	local storage = get_storage_config(host);
	if type(storage) == "table" then
		-- multiple storage backends in use that we need to purge
		local purged = {};
		for store, driver_name in pairs(storage) do
			if not purged[driver_name] then
				local driver = get_driver(host, store);
				if driver.purge then
					purged[driver_name] = driver:purge(user);
				else
					log("warn", "Storage driver %s does not support removing all user data, "
						.."you may need to delete it manually", driver_name);
				end
			end
		end
	end
	get_driver(host):purge(user); -- and the default driver

	olddm.purge(user, host); -- COMPAT list stores, like offline messages end up in the old datamanager

	return true;
end

function datamanager.load(username, host, datastore)
	return open(host, datastore):get(username);
end
function datamanager.store(username, host, datastore, data)
	return open(host, datastore):set(username, data);
end
function datamanager.users(host, datastore, typ)
	local driver = open(host, datastore, typ);
	if not driver.users then
		return function() log("warn", "Storage driver %s does not support listing users", driver.name) end
	end
	return driver:users();
end
function datamanager.stores(username, host, typ)
	return get_driver(host):stores(username, typ);
end
function datamanager.purge(username, host)
	return purge(username, host);
end

return {
	initialize_host = initialize_host;
	load_driver = load_driver;
	get_driver = get_driver;
	open = open;
	purge = purge;

	olddm = olddm;
};