util/datamanager.lua

   1 -- Prosody IM
   2 -- Copyright (C) 2008-2009 Matthew Wild
   3 -- Copyright (C) 2008-2009 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, type = setmetatable, type;
  12 local pairs, ipairs = pairs, ipairs;
  13 local char = string.char;
  14 local loadfile, setfenv, pcall = loadfile, setfenv, pcall;
  15 local log = require "util.logger".init("datamanager");
  16 local io_open = io.open;
  17 local os_remove = os.remove;
  18 local tostring, tonumber = tostring, tonumber;
  19 local error = error;
  20 local next = next;
  21 local t_insert = table.insert;
  22 local append = require "util.serialization".append;
  23 local path_separator = "/"; if os.getenv("WINDIR") then path_separator = "\\" end
  24 local lfs_mkdir = require "lfs".mkdir;
  25
  26 module "datamanager"
  27
  28 ---- utils -----
  29 local encode, decode;
  30 do
  31         local urlcodes = setmetatable({}, { __index = function (t, k) t[k] = char(tonumber("0x"..k)); return t[k]; end });
  32
  33         decode = function (s)
  34                 return s and (s:gsub("+", " "):gsub("%%([a-fA-F0-9][a-fA-F0-9])", urlcodes));
  35         end
  36
  37         encode = function (s)
  38                 return s and (s:gsub("%W", function (c) return format("%%%02x", c:byte()); end));
  39         end
  40 end
  41
  42 local _mkdir = {};
  43 local function mkdir(path)
  44         path = path:gsub("/", path_separator); -- TODO as an optimization, do this during path creation rather than here
  45         if not _mkdir[path] then
  46                 lfs_mkdir(path);
  47                 _mkdir[path] = true;
  48         end
  49         return path;
  50 end
  51
  52 local data_path = "data";
  53 local callbacks = {};
  54
  55 ------- API -------------
  56
  57 function set_data_path(path)
  58         log("debug", "Setting data path to: %s", path);
  59         data_path = path;
  60 end
  61
  62 local function callback(username, host, datastore, data)
  63         for _, f in ipairs(callbacks) do
  64                 username, host, datastore, data = f(username, host, datastore, data);
  65                 if username == false then break; end
  66         end
  67
  68         return username, host, datastore, data;
  69 end
  70 function add_callback(func)
  71         if not callbacks[func] then -- Would you really want to set the same callback more than once?
  72                 callbacks[func] = true;
  73                 callbacks[#callbacks+1] = func;
  74                 return true;
  75         end
  76 end
  77 function remove_callback(func)
  78         if callbacks[func] then
  79                 for i, f in ipairs(callbacks) do
  80                         if f == func then
  81                                 callbacks[i] = nil;
  82                                 callbacks[f] = nil;
  83                                 return true;
  84                         end
  85                 end
  86         end
  87 end
  88
  89 function getpath(username, host, datastore, ext, create)
  90         ext = ext or "dat";
  91         host = host and encode(host);
  92         username = username and encode(username);
  93         if username then
  94                 if create then mkdir(mkdir(mkdir(data_path).."/"..host).."/"..datastore); end
  95                 return format("%s/%s/%s/%s.%s", data_path, host, datastore, username, ext);
  96         elseif host then
  97                 if create then mkdir(mkdir(data_path).."/"..host); end
  98                 return format("%s/%s/%s.%s", data_path, host, datastore, ext);
  99         else
 100                 if create then mkdir(data_path); end
 101                 return format("%s/%s.%s", data_path, datastore, ext);
 102         end
 103 end
 104
 105 function load(username, host, datastore)
 106         local data, ret = loadfile(getpath(username, host, datastore));
 107         if not data then
 108                 log("debug", "Failed to load "..datastore.." storage ('"..ret.."') for user: "..(username or "nil").."@"..(host or "nil"));
 109                 return nil;
 110         end
 111         setfenv(data, {});
 112         local success, ret = pcall(data);
 113         if not success then
 114                 log("error", "Unable to load "..datastore.." storage ('"..ret.."') for user: "..(username or "nil").."@"..(host or "nil"));
 115                 return nil;
 116         end
 117         return ret;
 118 end
 119
 120 function store(username, host, datastore, data)
 121         if not data then
 122                 data = {};
 123         end
 124
 125         username, host, datastore, data = callback(username, host, datastore, data);
 126         if username == false then
 127                 return true; -- Don't save this data at all
 128         end
 129
 130         -- save the datastore
 131         local f, msg = io_open(getpath(username, host, datastore, nil, true), "w+");
 132         if not f then
 133                 log("error", "Unable to write to "..datastore.." storage ('"..msg.."') for user: "..(username or "nil").."@"..(host or "nil"));
 134                 return;
 135         end
 136         f:write("return ");
 137         append(f, data);
 138         f:close();
 139         if next(data) == nil then -- try to delete empty datastore
 140                 log("debug", "Removing empty %s datastore for user %s@%s", datastore, username or "nil", host or "nil");
 141                 os_remove(getpath(username, host, datastore));
 142         end
 143         -- we write data even when we are deleting because lua doesn't have a
 144         -- platform independent way of checking for non-exisitng files
 145         return true;
 146 end
 147
 148 function list_append(username, host, datastore, data)
 149         if not data then return; end
 150         if callback(username, host, datastore) == false then return true; end
 151         -- save the datastore
 152         local f, msg = io_open(getpath(username, host, datastore, "list", true), "a+");
 153         if not f then
 154                 log("error", "Unable to write to "..datastore.." storage ('"..msg.."') for user: "..(username or "nil").."@"..(host or "nil"));
 155                 return;
 156         end
 157         f:write("item(");
 158         append(f, data);
 159         f:write(");\n");
 160         f:close();
 161         return true;
 162 end
 163
 164 function list_store(username, host, datastore, data)
 165         if not data then
 166                 data = {};
 167         end
 168         if callback(username, host, datastore) == false then return true; end
 169         -- save the datastore
 170         local f, msg = io_open(getpath(username, host, datastore, "list", true), "w+");
 171         if not f then
 172                 log("error", "Unable to write to "..datastore.." storage ('"..msg.."') for user: "..(username or "nil").."@"..(host or "nil"));
 173                 return;
 174         end
 175         for _, d in ipairs(data) do
 176                 f:write("item(");
 177                 append(f, d);
 178                 f:write(");\n");
 179         end
 180         f:close();
 181         if next(data) == nil then -- try to delete empty datastore
 182                 log("debug", "Removing empty %s datastore for user %s@%s", datastore, username or "nil", host or "nil");
 183                 os_remove(getpath(username, host, datastore, "list"));
 184         end
 185         -- we write data even when we are deleting because lua doesn't have a
 186         -- platform independent way of checking for non-exisitng files
 187         return true;
 188 end
 189
 190 function list_load(username, host, datastore)
 191         local data, ret = loadfile(getpath(username, host, datastore, "list"));
 192         if not data then
 193                 log("debug", "Failed to load "..datastore.." storage ('"..ret.."') for user: "..(username or "nil").."@"..(host or "nil"));
 194                 return nil;
 195         end
 196         local items = {};
 197         setfenv(data, {item = function(i) t_insert(items, i); end});
 198         local success, ret = pcall(data);
 199         if not success then
 200                 log("error", "Unable to load "..datastore.." storage ('"..ret.."') for user: "..(username or "nil").."@"..(host or "nil"));
 201                 return nil;
 202         end
 203         return items;
 204 end
 205
 206 return _M;