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