plugins/mod_storage_memory.lua

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