plugins/storage/mod_ejabberd.lua

   1
   2 local setmetatable = setmetatable;
   3 local error = error;
   4 local unpack = unpack;
   5 local module = module;
   6 local tostring = tostring;
   7 local pairs, next = pairs, next;
   8 local prosody = prosody;
   9 local assert = assert;
  10 local require = require;
  11 local st = require "util.stanza";
  12 local DBI = require "DBI";
  13
  14 -- connect to db
  15 local option_datastore_params = module:get_option("datastore_params") or error("Missing option: datastore_params");
  16 local database;
  17 do
  18         local driver, db = unpack(option_datastore_params);
  19         module:log("debug", "Opening database: %s", "dbi:"..driver..":"..db);
  20         prosody.unlock_globals();
  21         local dbh, err = DBI.Connect(unpack(option_datastore_params));
  22         prosody.lock_globals();
  23         assert(dbh, err);
  24         dbh:autocommit(true)
  25         database = dbh;
  26 end
  27
  28 -- initialize db
  29 local ejabberd_init = module:require("ejabberd_init");
  30 ejabberd_init.init(database);
  31
  32 local sqlcache = {};
  33 local function prepare(sql)
  34         module:log("debug", "query: %s", sql);
  35         local err;
  36         local r = sqlcache[sql];
  37         if not r then
  38                 r, err = database:prepare(sql);
  39                 if not r then error("Unable to prepare SQL statement: "..err); end
  40                 sqlcache[sql] = r;
  41         end
  42         return r;
  43 end
  44
  45 local _parse_xml = module:require("xmlparse");
  46 local function parse_xml(str)
  47         local s = _parse_xml(str);
  48         if s and not s.gsub then
  49                 return st.preserialize(s);
  50         end
  51 end
  52 local function unparse_xml(s)
  53         return tostring(st.deserialize(s));
  54 end
  55
  56
  57 local handlers = {};
  58
  59 handlers.accounts = {
  60         get = function(self, user)
  61                 local select = self:query("select password from users where username=? and host=?", user, self.host);
  62                 local row = select and select:fetch();
  63                 if row then return { password = row[1] }; end
  64         end;
  65         set = function(self, user, data)
  66                 if data and data.password then
  67                         return self:modify("update users set password=? where username=? and host=?", data.password, user, self.host)
  68                                 or self:modify("insert into users (username, host, password) values (?, ?, ?)", user, self.host, data.password);
  69                 else
  70                         return self:modify("delete from users where username=? and host=?", user, self.host);
  71                 end
  72         end;
  73 };
  74 handlers.vcard = {
  75         get = function(self, user)
  76                 local select = self:query("select vcard from vcard where username=? and host=?", user, self.host);
  77                 local row = select and select:fetch();
  78                 if row then return parse_xml(row[1]); end
  79         end;
  80         set = function(self, user, data)
  81                 if data then
  82                         data = unparse_xml(data);
  83                         return self:modify("update vcard set vcard=? where username=? and host=?", data, user, self.host)
  84                                 or self:modify("insert into vcard (username, host, vcard) values (?, ?, ?)", user, self.host, data);
  85                 else
  86                         return self:modify("delete from vcard where username=? and host=?", user, self.host);
  87                 end
  88         end;
  89 };
  90 handlers.private = {
  91         get = function(self, user)
  92                 local select = self:query("select namespace,data from private_storage where username=? and host=?", user, self.host);
  93                 if select then
  94                         local data = {};
  95                         for row in select:rows() do
  96                                 data[row[1]] = parse_xml(row[2]);
  97                         end
  98                         return data;
  99                 end
 100         end;
 101         set = function(self, user, data)
 102                 if data then
 103                         self:modify("delete from private_storage where username=? and host=?", user, self.host);
 104                         for namespace,text in pairs(data) do
 105                                 self:modify("insert into private_storage (username, host, namespace, data) values (?, ?, ?, ?)", user, self.host, namespace, unparse_xml(text));
 106                         end
 107                         return true;
 108                 else
 109                         return self:modify("delete from private_storage where username=? and host=?", user, self.host);
 110                 end
 111         end;
 112         -- TODO map_set, map_get
 113 };
 114 local subscription_map = { N = "none", B = "both", F = "from", T = "to" };
 115 local subscription_map_reverse = { none = "N", both = "B", from = "F", to = "T" };
 116 handlers.roster = {
 117         get = function(self, user)
 118                 local select = self:query("select jid,nick,subscription,ask,server,subscribe,type from rosterusers where username=?", user);
 119                 if select then
 120                         local roster = { pending = {} };
 121                         for row in select:rows() do
 122                                 local jid,nick,subscription,ask,server,subscribe,typ = unpack(row);
 123                                 local item = { groups = {} };
 124                                 if nick == "" then nick = nil; end
 125                                 item.nick = nick;
 126                                 item.subscription = subscription_map[subscription];
 127                                 if ask == "N" then ask = nil;
 128                                 elseif ask == "O" then ask = "subscribe"
 129                                 elseif ask == "I" then roster.pending[jid] = true; ask = nil;
 130                                 elseif ask == "B" then roster.pending[jid] = true; ask = "subscribe";
 131                                 else module:log("debug", "bad roster_item.ask: %s", ask); ask = nil; end
 132                                 item.ask = ask;
 133                                 roster[jid] = item;
 134                         end
 135
 136                         select = self:query("select jid,grp from rostergroups where username=?", user);
 137                         if select then
 138                                 for row in select:rows() do
 139                                         local jid,grp = unpack(row);
 140                                         if roster[jid] then roster[jid].groups[grp] = true; end
 141                                 end
 142                         end
 143                         select = self:query("select version from roster_version where username=?", user);
 144                         local row = select and select:fetch();
 145                         if row then
 146                                 roster[false] = { version = row[1]; };
 147                         end
 148                         return roster;
 149                 end
 150         end;
 151         set = function(self, user, data)
 152                 if data and next(data) ~= nil then
 153                         self:modify("delete from rosterusers where username=?", user);
 154                         self:modify("delete from rostergroups where username=?", user);
 155                         self:modify("delete from roster_version where username=?", user);
 156                         local done = {};
 157                         local pending = data.pending or {};
 158                         for jid,item in pairs(data) do
 159                                 if jid and jid ~= "pending" then
 160                                         local subscription = subscription_map_reverse[item.subscription];
 161                                         local ask;
 162                                         if pending[jid] then
 163                                                 if item.ask then ask = "B"; else ask = "I"; end
 164                                         else
 165                                                 if item.ask then ask = "O"; else ask = "N"; end
 166                                         end
 167                                         local r = self:modify("insert into rosterusers (username,jid,nick,subscription,ask,askmessage,server,subscribe) values (?, ?, ?, ?, ?, '', '', '')", user, jid, item.nick or "", subscription, ask);
 168                                         if not r then module:log("debug", "--- :( %s", tostring(r)); end
 169                                         done[jid] = true;
 170                                         for group in pairs(item.groups) do
 171                                                 self:modify("insert into rostergroups (username,jid,grp) values (?, ?, ?)", user, jid, group);
 172                                         end
 173                                 end
 174                         end
 175                         for jid in pairs(pending) do
 176                                 if not done[jid] then
 177                                         self:modify("insert into rosterusers (username,jid,nick,subscription,ask,askmessage,server,subscribe) values (?, ?, ?, ?, ?. ''. ''. '')", user, jid, "", "N", "I");
 178                                 end
 179                         end
 180                         local version = data[false] and data[false].version;
 181                         if version then
 182                                 self:modify("insert into roster_version (username,version) values (?, ?)", user, version);
 183                         end
 184                         return true;
 185                 else
 186                         self:modify("delete from rosterusers where username=?", user);
 187                         self:modify("delete from rostergroups where username=?", user);
 188                         self:modify("delete from roster_version where username=?", user);
 189                 end
 190         end;
 191 };
 192
 193 -----------------------------
 194 local driver = {};
 195 driver.__index = driver;
 196
 197 function driver:query(sql, ...)
 198         local stmt,err = prepare(sql);
 199         if not stmt then
 200                 module:log("error", "Failed to prepare SQL [[%s]], error: %s", sql, err);
 201                 return nil, err;
 202         end
 203         local ok, err = stmt:execute(...);
 204         if not ok then
 205                 module:log("error", "Failed to execute SQL [[%s]], error: %s", sql, err);
 206                 return nil, err;
 207         end
 208         return stmt;
 209 end
 210 function driver:modify(sql, ...)
 211         local stmt, err = self:query(sql, ...);
 212         if stmt and stmt:affected() > 0 then return stmt; end
 213         return nil, err;
 214 end
 215
 216 function driver:open(datastore, typ)
 217         local instance = setmetatable({ host = module.host, datastore = datastore }, self);
 218         local handler = handlers[datastore];
 219         if not handler then return nil; end
 220         for key,val in pairs(handler) do
 221                 instance[key] = val;
 222         end
 223         if instance.init then instance:init(); end
 224         return instance;
 225 end
 226
 227 -----------------------------
 228
 229 module:add_item("data-driver", driver);