util/datamanager.lua

   1 -- Prosody IM
   2 -- Copyright (C) 2008-2010 Matthew Wild
   3 -- Copyright (C) 2008-2010 Waqas Hussain
   4 --
   5 -- This project is MIT/X11 licensed. Please see the
   6 -- COPYING file in the source package for more information.
   7 --
   8
   9
  10 local format = string.format;
  11 local setmetatable = setmetatable;
  12 local ipairs = ipairs;
  13 local char = string.char;
  14 local pcall = pcall;
  15 local log = require "util.logger".init("datamanager");
  16 local io_open = io.open;
  17 local os_remove = os.remove;
  18 local os_rename = os.rename;
  19 local tonumber = tonumber;
  20 local next = next;
  21 local type = type;
  22 local t_insert = table.insert;
  23 local t_concat = table.concat;
  24 local envloadfile = require"util.envload".envloadfile;
  25 local serialize = require "util.serialization".serialize;
  26 local lfs = require "lfs";
  27 -- Extract directory separator from package.config (an undocumented string that comes with lua)
  28 local path_separator = assert ( package.config:match ( "^([^\n]+)" ) , "package.config not in standard form" )
  29
  30 local prosody = prosody;
  31
  32 local raw_mkdir = lfs.mkdir;
  33 local atomic_append;
  34 local ENOENT = 2;
  35 pcall(function()
  36         local pposix = require "util.pposix";
  37         raw_mkdir = pposix.mkdir or raw_mkdir; -- Doesn't trample on umask
  38         atomic_append = pposix.atomic_append;
  39         ENOENT = pposix.ENOENT or ENOENT;
  40 end);
  41
  42 local _ENV = nil;
  43 -- luacheck: std none
  44
  45 ---- utils -----
  46 local encode, decode, store_encode;
  47 do
  48         local urlcodes = setmetatable({}, { __index = function (t, k) t[k] = char(tonumber(k, 16)); return t[k]; end });
  49
  50         decode = function (s)
  51                 return s and (s:gsub("%%(%x%x)", urlcodes));
  52         end
  53
  54         encode = function (s)
  55                 return s and (s:gsub("%W", function (c) return format("%%%02x", c:byte()); end));
  56         end
  57
  58         -- Special encode function for store names, which historically were unencoded.
  59         -- All currently known stores use a-z and underscore, so this one preserves underscores.
  60         store_encode = function (s)
  61                 return s and (s:gsub("[^_%w]", function (c) return format("%%%02x", c:byte()); end));
  62         end
  63 end
  64
  65 if not atomic_append then
  66         function atomic_append(f, data)
  67                 local pos = f:seek();
  68                 if not f:write(data) or not f:flush() then
  69                         f:seek("set", pos);
  70                         f:write((" "):rep(#data));
  71                         f:flush();
  72                         return nil, "write-failed";
  73                 end
  74                 return true;
  75         end
  76 end
  77
  78 local _mkdir = {};
  79 local function mkdir(path)
  80         path = path:gsub("/", path_separator); -- TODO as an optimization, do this during path creation rather than here
  81         if not _mkdir[path] then
  82                 raw_mkdir(path);
  83                 _mkdir[path] = true;
  84         end
  85         return path;
  86 end
  87
  88 local data_path = (prosody and prosody.paths and prosody.paths.data) or ".";
  89 local callbacks = {};
  90
  91 ------- API -------------
  92
  93 local function set_data_path(path)
  94         log("debug", "Setting data path to: %s", path);
  95         data_path = path;
  96 end
  97
  98 local function callback(username, host, datastore, data)
  99         for _, f in ipairs(callbacks) do
 100                 username, host, datastore, data = f(username, host, datastore, data);
 101                 if username == false then break; end
 102         end
 103
 104         return username, host, datastore, data;
 105 end
 106 local function add_callback(func)
 107         if not callbacks[func] then -- Would you really want to set the same callback more than once?
 108                 callbacks[func] = true;
 109                 callbacks[#callbacks+1] = func;
 110                 return true;
 111         end
 112 end
 113 local function remove_callback(func)
 114         if callbacks[func] then
 115                 for i, f in ipairs(callbacks) do
 116                         if f == func then
 117                                 callbacks[i] = nil;
 118                                 callbacks[f] = nil;
 119                                 return true;
 120                         end
 121                 end
 122         end
 123 end
 124
 125 local function getpath(username, host, datastore, ext, create)
 126         ext = ext or "dat";
 127         host = (host and encode(host)) or "_global";
 128         username = username and encode(username);
 129         datastore = store_encode(datastore);
 130         if username then
 131                 if create then mkdir(mkdir(mkdir(data_path).."/"..host).."/"..datastore); end
 132                 return format("%s/%s/%s/%s.%s", data_path, host, datastore, username, ext);
 133         else
 134                 if create then mkdir(mkdir(data_path).."/"..host); end
 135                 return format("%s/%s/%s.%s", data_path, host, datastore, ext);
 136         end
 137 end
 138
 139 local function load(username, host, datastore)
 140         local data, err, errno = envloadfile(getpath(username, host, datastore), {});
 141         if not data then
 142                 if errno == ENOENT then
 143                         -- No such file, ok to ignore
 144                         return nil;
 145                 end
 146                 log("error", "Failed to load %s storage ('%s') for user: %s@%s", datastore, err, username or "nil", host or "nil");
 147                 return nil, "Error reading storage";
 148         end
 149
 150         local success, ret = pcall(data);
 151         if not success then
 152                 log("error", "Unable to load %s storage ('%s') for user: %s@%s", datastore, ret, username or "nil", host or "nil");
 153                 return nil, "Error reading storage";
 154         end
 155         return ret;
 156 end
 157
 158 local function atomic_store(filename, data)
 159         local scratch = filename.."~";
 160         local f, ok, msg, errno;
 161
 162         f, msg, errno = io_open(scratch, "w");
 163         if not f then
 164                 return nil, msg;
 165         end
 166
 167         ok, msg = f:write(data);
 168         if not ok then
 169                 f:close();
 170                 os_remove(scratch);
 171                 return nil, msg;
 172         end
 173
 174         ok, msg = f:close();
 175         if not ok then
 176                 os_remove(scratch);
 177                 return nil, msg;
 178         end
 179
 180         return os_rename(scratch, filename);
 181 end
 182
 183 if prosody and prosody.platform ~= "posix" then
 184         -- os.rename does not overwrite existing files on Windows
 185         -- TODO We could use Transactional NTFS on Vista and above
 186         function atomic_store(filename, data)
 187                 local f, err = io_open(filename, "w");
 188                 if not f then return f, err; end
 189                 local ok, msg = f:write(data);
 190                 if not ok then f:close(); return ok, msg; end
 191                 return f:close();
 192         end
 193 end
 194
 195 local function store(username, host, datastore, data)
 196         if not data then
 197                 data = {};
 198         end
 199
 200         username, host, datastore, data = callback(username, host, datastore, data);
 201         if username == false then
 202                 return true; -- Don't save this data at all
 203         end
 204
 205         -- save the datastore
 206         local d = "return " .. serialize(data) .. ";\n";
 207         local mkdir_cache_cleared;
 208         repeat
 209                 local ok, msg = atomic_store(getpath(username, host, datastore, nil, true), d);
 210                 if not ok then
 211                         if not mkdir_cache_cleared then -- We may need to recreate a removed directory
 212                                 _mkdir = {};
 213                                 mkdir_cache_cleared = true;
 214                         else
 215                                 log("error", "Unable to write to %s storage ('%s') for user: %s@%s", datastore, msg, username or "nil", host or "nil");
 216                                 return nil, "Error saving to storage";
 217                         end
 218                 end
 219                 if next(data) == nil then -- try to delete empty datastore
 220                         log("debug", "Removing empty %s datastore for user %s@%s", datastore, username or "nil", host or "nil");
 221                         os_remove(getpath(username, host, datastore));
 222                 end
 223                 -- we write data even when we are deleting because lua doesn't have a
 224                 -- platform independent way of checking for non-exisitng files
 225         until ok;
 226         return true;
 227 end
 228
 229 -- Append a blob of data to a file
 230 local function append(username, host, datastore, ext, data)
 231         if type(data) ~= "string" then return; end
 232         local filename = getpath(username, host, datastore, ext, true);
 233
 234         local f = io_open(filename, "r+");
 235         if not f then
 236                 return atomic_store(filename, data);
 237                 -- File did probably not exist, let's create it
 238         end
 239
 240         local pos = f:seek("end");
 241
 242         local ok, msg = atomic_append(f, data);
 243
 244         if not ok then
 245                 f:close();
 246                 return ok, msg, "write";
 247         end
 248
 249         ok, msg = f:close();
 250         if not ok then
 251                 return ok, msg, "close";
 252         end
 253
 254         return true, pos;
 255 end
 256
 257 local function list_append(username, host, datastore, data)
 258         if not data then return; end
 259         if callback(username, host, datastore) == false then return true; end
 260         -- save the datastore
 261
 262         data = "item(" ..  serialize(data) .. ");\n";
 263         local ok, msg, where = append(username, host, datastore, "list", data);
 264         if not ok then
 265                 log("error", "Unable to write to %s storage ('%s' in %s) for user: %s@%s",
 266                         datastore, msg, where, username or "nil", host or "nil");
 267                 return ok, msg;
 268         end
 269         return true;
 270 end
 271
 272 local function list_store(username, host, datastore, data)
 273         if not data then
 274                 data = {};
 275         end
 276         if callback(username, host, datastore) == false then return true; end
 277         -- save the datastore
 278         local d = {};
 279         for i, item in ipairs(data) do
 280                 d[i] = "item(" .. serialize(item) .. ");\n";
 281         end
 282         local ok, msg = atomic_store(getpath(username, host, datastore, "list", true), t_concat(d));
 283         if not ok then
 284                 log("error", "Unable to write to %s storage ('%s') for user: %s@%s", datastore, msg, username or "nil", host or "nil");
 285                 return;
 286         end
 287         if next(data) == nil then -- try to delete empty datastore
 288                 log("debug", "Removing empty %s datastore for user %s@%s", datastore, username or "nil", host or "nil");
 289                 os_remove(getpath(username, host, datastore, "list"));
 290         end
 291         -- we write data even when we are deleting because lua doesn't have a
 292         -- platform independent way of checking for non-exisitng files
 293         return true;
 294 end
 295
 296 local function list_load(username, host, datastore)
 297         local items = {};
 298         local data, err, errno = envloadfile(getpath(username, host, datastore, "list"), {item = function(i) t_insert(items, i); end});
 299         if not data then
 300                 if errno == ENOENT then
 301                         -- No such file, ok to ignore
 302                         return nil;
 303                 end
 304                 log("error", "Failed to load %s storage ('%s') for user: %s@%s", datastore, err, username or "nil", host or "nil");
 305                 return nil, "Error reading storage";
 306         end
 307
 308         local success, ret = pcall(data);
 309         if not success then
 310                 log("error", "Unable to load %s storage ('%s') for user: %s@%s", datastore, ret, username or "nil", host or "nil");
 311                 return nil, "Error reading storage";
 312         end
 313         return items;
 314 end
 315
 316 local type_map = {
 317         keyval = "dat";
 318         list = "list";
 319 }
 320
 321 local function users(host, store, typ) -- luacheck: ignore 431/store
 322         typ = type_map[typ or "keyval"];
 323         local store_dir = format("%s/%s/%s", data_path, encode(host), store);
 324
 325         local mode, err = lfs.attributes(store_dir, "mode");
 326         if not mode then
 327                 return function() log("debug", "%s", err or (store_dir .. " does not exist")) end
 328         end
 329         local next, state = lfs.dir(store_dir); -- luacheck: ignore 431/next 431/state
 330         return function(state) -- luacheck: ignore 431/state
 331                 for node in next, state do
 332                         local file, ext = node:match("^(.*)%.([dalist]+)$");
 333                         if file and ext == typ then
 334                                 return decode(file);
 335                         end
 336                 end
 337         end, state;
 338 end
 339
 340 local function stores(username, host, typ)
 341         typ = type_map[typ or "keyval"];
 342         local store_dir = format("%s/%s/", data_path, encode(host));
 343
 344         local mode, err = lfs.attributes(store_dir, "mode");
 345         if not mode then
 346                 return function() log("debug", err or (store_dir .. " does not exist")) end
 347         end
 348         local next, state = lfs.dir(store_dir); -- luacheck: ignore 431/next 431/state
 349         return function(state) -- luacheck: ignore 431/state
 350                 for node in next, state do
 351                         if not node:match"^%." then
 352                                 if username == true then
 353                                         if lfs.attributes(store_dir..node, "mode") == "directory" then
 354                                                 return decode(node);
 355                                         end
 356                                 elseif username then
 357                                         local store_name = decode(node);
 358                                         if lfs.attributes(getpath(username, host, store_name, typ), "mode") then
 359                                                 return store_name;
 360                                         end
 361                                 elseif lfs.attributes(node, "mode") == "file" then
 362                                         local file, ext = node:match("^(.*)%.([dalist]+)$");
 363                                         if ext == typ then
 364                                                 return decode(file)
 365                                         end
 366                                 end
 367                         end
 368                 end
 369         end, state;
 370 end
 371
 372 local function do_remove(path)
 373         local ok, err = os_remove(path);
 374         if not ok and lfs.attributes(path, "mode") then
 375                 return ok, err;
 376         end
 377         return true
 378 end
 379
 380 local function purge(username, host)
 381         local host_dir = format("%s/%s/", data_path, encode(host));
 382         local ok, iter, state, var = pcall(lfs.dir, host_dir);
 383         if not ok then
 384                 return ok, iter;
 385         end
 386         local errs = {};
 387         for file in iter, state, var do
 388                 if lfs.attributes(host_dir..file, "mode") == "directory" then
 389                         local store_name = decode(file);
 390                         local ok, err = do_remove(getpath(username, host, store_name));
 391                         if not ok then errs[#errs+1] = err; end
 392
 393                         local ok, err = do_remove(getpath(username, host, store_name, "list"));
 394                         if not ok then errs[#errs+1] = err; end
 395                 end
 396         end
 397         return #errs == 0, t_concat(errs, ", ");
 398 end
 399
 400 return {
 401         set_data_path = set_data_path;
 402         add_callback = add_callback;
 403         remove_callback = remove_callback;
 404         getpath = getpath;
 405         load = load;
 406         store = store;
 407         append_raw = append;
 408         store_raw = atomic_store;
 409         list_append = list_append;
 410         list_store = list_store;
 411         list_load = list_load;
 412         users = users;
 413         stores = stores;
 414         purge = purge;
 415         path_decode = decode;
 416         path_encode = encode;
 417 };