1local serialize = require "util.serialization".serialize;
2local array = require "util.array";
3local envload = require "util.envload".envload;
4local st = require "util.stanza";
5local is_stanza = st.is_stanza or function (s) return getmetatable(s) == st.stanza_mt end
6local new_id = require "util.id".medium;
7
8local auto_purge_enabled = module:get_option_boolean("storage_memory_temporary", false);
9local auto_purge_stores = module:get_option_set("storage_memory_temporary_stores", {});
10
11local memory = setmetatable({}, {
12	__index = function(t, k)
13		local store = module:shared(k)
14		t[k] = store;
15		return store;
16	end
17});
18
19local function NULL() return nil end
20
21local function _purge_store(self, username)
22	self.store[username or NULL] = nil;
23	return true;
24end
25
26local function _users(self)
27	return next, self.store, nil;
28end
29
30local keyval_store = {};
31keyval_store.__index = keyval_store;
32
33function keyval_store:get(username)
34	return (self.store[username or NULL] or NULL)();
35end
36
37function keyval_store:set(username, data)
38	if data ~= nil then
39		data = envload("return "..serialize(data), "=(data)", {});
40	end
41	self.store[username or NULL] = data;
42	return true;
43end
44
45keyval_store.purge = _purge_store;
46
47keyval_store.users = _users;
48
49local archive_store = {};
50archive_store.__index = archive_store;
51
52archive_store.users = _users;
53
54function archive_store:append(username, key, value, when, with)
55	if is_stanza(value) then
56		value = st.preserialize(value);
57		value = envload("return xml"..serialize(value), "=(stanza)", { xml = st.deserialize })
58	else
59		value = envload("return "..serialize(value), "=(data)", {});
60	end
61	local a = self.store[username or NULL];
62	if not a then
63		a = {};
64		self.store[username or NULL] = a;
65	end
66	local v = { key = key, when = when, with = with, value = value };
67	if not key then
68		key = new_id();
69		v.key = key;
70	end
71	if a[key] then
72		table.remove(a, a[key]);
73	end
74	local i = #a+1;
75	a[i] = v;
76	a[key] = i;
77	return key;
78end
79
80function archive_store:find(username, query)
81	local items = self.store[username or NULL];
82	if not items then
83		return function () end, 0;
84	end
85	local count = #items;
86	local i = 0;
87	if query then
88		items = array():append(items);
89		if query.key then
90			items:filter(function (item)
91				return item.key == query.key;
92			end);
93		end
94		if query.with then
95			items:filter(function (item)
96				return item.with == query.with;
97			end);
98		end
99		if query.start then
100			items:filter(function (item)
101				return item.when >= query.start;
102			end);
103		end
104		if query["end"] then
105			items:filter(function (item)
106				return item.when <= query["end"];
107			end);
108		end
109		count = #items;
110		if query.reverse then
111			items:reverse();
112			if query.before then
113				for j = 1, count do
114					if (items[j].key or tostring(j)) == query.before then
115						i = j;
116						break;
117					end
118				end
119			end
120		elseif query.after then
121			for j = 1, count do
122				if (items[j].key or tostring(j)) == query.after then
123					i = j;
124					break;
125				end
126			end
127		end
128		if query.limit and #items - i > query.limit then
129			items[i+query.limit+1] = nil;
130		end
131	end
132	return function ()
133		i = i + 1;
134		local item = items[i];
135		if not item then return; end
136		return item.key, item.value(), item.when, item.with;
137	end, count;
138end
139
140
141function archive_store:delete(username, query)
142	if not query or next(query) == nil then
143		self.store[username or NULL] = nil;
144		return true;
145	end
146	local items = self.store[username or NULL];
147	if not items then
148		-- Store is empty
149		return 0;
150	end
151	items = array(items);
152	local count_before = #items;
153	if query then
154		if query.key then
155			items:filter(function (item)
156				return item.key ~= query.key;
157			end);
158		end
159		if query.with then
160			items:filter(function (item)
161				return item.with ~= query.with;
162			end);
163		end
164		if query.start then
165			items:filter(function (item)
166				return item.when < query.start;
167			end);
168		end
169		if query["end"] then
170			items:filter(function (item)
171				return item.when > query["end"];
172			end);
173		end
174		if query.truncate and #items > query.truncate then
175			if query.reverse then
176				-- Before: { 1, 2, 3, 4, 5, }
177				-- After: { 1, 2, 3 }
178				for i = #items, query.truncate + 1, -1 do
179					items[i] = nil;
180				end
181			else
182				-- Before: { 1, 2, 3, 4, 5, }
183				-- After: { 3, 4, 5 }
184				local offset = #items - query.truncate;
185				for i = 1, #items do
186					items[i] = items[i+offset];
187				end
188			end
189		end
190	end
191	local count = count_before - #items;
192	if count == 0 then
193		return 0; -- No changes, skip write
194	end
195	setmetatable(items, nil);
196
197	do -- re-index by key
198		for k in pairs(items) do
199			if type(k) == "string" then
200				items[k] = nil;
201			end
202		end
203
204		for i = 1, #items do
205			items[ items[i].key ] = i;
206		end
207	end
208
209	return count;
210end
211
212archive_store.purge = _purge_store;
213
214local stores = {
215	keyval = keyval_store;
216	archive = archive_store;
217}
218
219local driver = {};
220
221function driver:open(store, typ) -- luacheck: ignore 212/self
222	local store_mt = stores[typ or "keyval"];
223	if store_mt then
224		return setmetatable({ store = memory[store] }, store_mt);
225	end
226	return nil, "unsupported-store";
227end
228
229function driver:purge(user) -- luacheck: ignore 212/self
230	for _, store in pairs(memory) do
231		store[user] = nil;
232	end
233end
234
235if auto_purge_enabled then
236	module:hook("resource-unbind", function (event)
237		local user_bare_jid = event.session.username.."@"..event.session.host;
238		if not prosody.bare_sessions[user_bare_jid] then -- User went offline
239			module:log("debug", "Clearing store for offline user %s", user_bare_jid);
240			local f, s, v;
241			if auto_purge_stores:empty() then
242				f, s, v = pairs(memory);
243			else
244				f, s, v = auto_purge_stores:items();
245			end
246
247			for store_name in f, s, v do
248				if memory[store_name] then
249					memory[store_name][event.session.username] = nil;
250				end
251			end
252		end
253	end);
254end
255
256module:provides("storage", driver);
257