Chapril
/
xmpp.chapril.org-ejabberd
mirror of https://github.com/processone/ejabberd.git
24
1
Fork 0
Code Releases Activity

Preliminary Riak support

This commit is contained in:
Alexey Shchepin 2012-01-25 12:02:16 +02:00 committed by Evgeniy Khramtsov
parent fc692ea512
commit 47763c10e3
6 changed files with 2467 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

134
src/ejabberd_riak.erl Normal file
View File

@ -0,0 +1,134 @@
%%%----------------------------------------------------------------------
%%% File : ejabberd_riak.erl
%%% Author : Alexey Shchepin <alexey@process-one.net>
%%% Purpose : Serve Riak connection
%%% Created : 29 Dec 2011 by Alexey Shchepin <alexey@process-one.net>
%%%
%%%
%%% ejabberd, Copyright (C) 2002-2011 ProcessOne
%%%
%%% This program is free software; you can redistribute it and/or
%%% modify it under the terms of the GNU General Public License as
%%% published by the Free Software Foundation; either version 2 of the
%%% License, or (at your option) any later version.
%%%
%%% This program is distributed in the hope that it will be useful,
%%% but WITHOUT ANY WARRANTY; without even the implied warranty of
%%% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
%%% General Public License for more details.
%%%
%%% You should have received a copy of the GNU General Public License
%%% along with this program; if not, write to the Free Software
%%% Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
%%% 02111-1307 USA
%%%
%%%----------------------------------------------------------------------
-module(ejabberd_riak).
-author('alexey@process-one.net').
%% External exports
-export([start_link/1,
put/4,
put/5,
get_object/3,
get/3,
get_objects_by_index/4,
get_by_index/4,
get_keys_by_index/4,
count_by_index/4,
delete/3]).
-include("ejabberd.hrl").
%%%----------------------------------------------------------------------
%%% API
%%%----------------------------------------------------------------------
start_link(StartInterval) ->
{ok, Pid} = riakc_pb_socket:start_link(
"127.0.0.1", 8081,
[auto_reconnect]),
ejabberd_riak_sup:add_pid(Pid),
{ok, Pid}.
make_bucket(Host, Table) ->
iolist_to_binary([Host, $@, Table]).
put(Host, Table, Key, Value) ->
Bucket = make_bucket(Host, Table),
Obj = riakc_obj:new(Bucket, Key, Value),
riakc_pb_socket:put(ejabberd_riak_sup:get_random_pid(), Obj).
put(Host, Table, Key, Value, Indexes) ->
Bucket = make_bucket(Host, Table),
Obj = riakc_obj:new(Bucket, Key, Value),
MetaData = dict:store(<<"index">>, Indexes, dict:new()),
Obj2 = riakc_obj:update_metadata(Obj, MetaData),
riakc_pb_socket:put(ejabberd_riak_sup:get_random_pid(), Obj2).
get_object(Host, Table, Key) ->
Bucket = make_bucket(Host, Table),
riakc_pb_socket:get(ejabberd_riak_sup:get_random_pid(), Bucket, Key).
get(Host, Table, Key) ->
case get_object(Host, Table, Key) of
{ok, Obj} ->
{ok, riakc_obj:get_value(Obj)};
Error ->
Error
end.
get_objects_by_index(Host, Table, Index, Key) ->
Bucket = make_bucket(Host, Table),
case riakc_pb_socket:mapred(
ejabberd_riak_sup:get_random_pid(),
{index, Bucket, Index, Key},
[{map, {modfun, riak_kv_mapreduce, map_identity}, none, true}]) of
{ok, [{_, Objs}]} ->
{ok, Objs};
Error ->
Error
end.
get_by_index(Host, Table, Index, Key) ->
Bucket = make_bucket(Host, Table),
case riakc_pb_socket:mapred(
ejabberd_riak_sup:get_random_pid(),
{index, Bucket, Index, Key},
[{map, {modfun, riak_kv_mapreduce, map_object_value},
none, true}]) of
{ok, [{_, Objs}]} ->
{ok, Objs};
Error ->
Error
end.
get_keys_by_index(Host, Table, Index, Key) ->
Bucket = make_bucket(Host, Table),
case riakc_pb_socket:mapred(
ejabberd_riak_sup:get_random_pid(),
{index, Bucket, Index, Key},
[]) of
{ok, [{_, Ls}]} ->
{ok, [K || {_, K} <- Ls]};
Error ->
Error
end.
count_by_index(Host, Table, Index, Key) ->
Bucket = make_bucket(Host, Table),
case riakc_pb_socket:mapred(
ejabberd_riak_sup:get_random_pid(),
{index, Bucket, Index, Key},
[{reduce, {modfun, riak_kv_mapreduce, reduce_count_inputs},
none, true}]) of
{ok, [{_, [Cnt]}]} ->
{ok, Cnt};
Error ->
Error
end.
delete(Host, Table, Key) ->
Bucket = make_bucket(Host, Table),
riakc_pb_socket:delete(ejabberd_riak_sup:get_random_pid(), Bucket, Key).

142
src/ejabberd_riak_sup.erl Normal file
View File

@ -0,0 +1,142 @@
%%%----------------------------------------------------------------------
%%% File : ejabberd_riak_sup.erl
%%% Author : Alexey Shchepin <alexey@process-one.net>
%%% Purpose : Riak connections supervisor
%%% Created : 29 Dec 2011 by Alexey Shchepin <alexey@process-one.net>
%%%
%%%
%%% ejabberd, Copyright (C) 2002-2011 ProcessOne
%%%
%%% This program is free software; you can redistribute it and/or
%%% modify it under the terms of the GNU General Public License as
%%% published by the Free Software Foundation; either version 2 of the
%%% License, or (at your option) any later version.
%%%
%%% This program is distributed in the hope that it will be useful,
%%% but WITHOUT ANY WARRANTY; without even the implied warranty of
%%% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
%%% General Public License for more details.
%%%
%%% You should have received a copy of the GNU General Public License
%%% along with this program; if not, write to the Free Software
%%% Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
%%% 02111-1307 USA
%%%
%%%----------------------------------------------------------------------
-module(ejabberd_riak_sup).
-author('alexey@process-one.net').
%% API
-export([start/0,
start_link/0,
init/1,
add_pid/1,
remove_pid/1,
get_pids/0,
get_random_pid/0
]).
-include("ejabberd.hrl").
-define(DEFAULT_POOL_SIZE, 10).
-define(DEFAULT_RIAK_START_INTERVAL, 30). % 30 seconds
% time to wait for the supervisor to start its child before returning
% a timeout error to the request
-define(CONNECT_TIMEOUT, 500). % milliseconds
-record(riak_pool, {undefined, pid}).
start() ->
SupervisorName = ?MODULE,
ChildSpec =
{SupervisorName,
{?MODULE, start_link, []},
transient,
infinity,
supervisor,
[?MODULE]},
case supervisor:start_child(ejabberd_sup, ChildSpec) of
{ok, _PID} ->
ok;
_Error ->
?ERROR_MSG("Start of supervisor ~p failed:~n~p~nRetrying...~n",
[SupervisorName, _Error]),
timer:sleep(5000),
start()
end.
start_link() ->
mnesia:create_table(riak_pool,
[{ram_copies, [node()]},
{type, bag},
{local_content, true},
{attributes, record_info(fields, riak_pool)}]),
mnesia:add_table_copy(riak_pool, node(), ram_copies),
F = fun() ->
mnesia:delete({riak_pool, undefined})
end,
mnesia:ets(F),
supervisor:start_link({local, ?MODULE}, ?MODULE, []).
init([]) ->
PoolSize =
case ejabberd_config:get_local_option(riak_pool_size) of
I when is_integer(I) ->
I;
undefined ->
?DEFAULT_POOL_SIZE;
Other ->
?ERROR_MSG("Wrong riak_pool_size definition '~p' "
"default to ~p~n",
[Other, ?DEFAULT_POOL_SIZE]),
?DEFAULT_POOL_SIZE
end,
StartInterval =
case ejabberd_config:get_local_option(riak_start_interval) of
Interval when is_integer(Interval) ->
Interval;
undefined ->
?DEFAULT_RIAK_START_INTERVAL;
_Other2 ->
?ERROR_MSG("Wrong riak_start_interval "
"definition '~p', "
"defaulting to ~p~n",
[_Other2,
?DEFAULT_RIAK_START_INTERVAL]),
?DEFAULT_RIAK_START_INTERVAL
end,
{ok, {{one_for_one, PoolSize*10, 1},
lists:map(
fun(I) ->
{I,
{ejabberd_riak, start_link, [StartInterval*1000]},
transient,
2000,
worker,
[?MODULE]}
end, lists:seq(1, PoolSize))}}.
get_pids() ->
Rs = mnesia:dirty_read(riak_pool, undefined),
[R#riak_pool.pid || R <- Rs].
get_random_pid() ->
Pids = get_pids(),
lists:nth(erlang:phash(now(), length(Pids)), Pids).
add_pid(Pid) ->
F = fun() ->
mnesia:write(
#riak_pool{pid = Pid})
end,
mnesia:ets(F).
remove_pid(Pid) ->
F = fun() ->
mnesia:delete_object(
#riak_pool{pid = Pid})
end,
mnesia:ets(F).

533
src/mod_offline_riak.erl Normal file
View File

@ -0,0 +1,533 @@
%%%----------------------------------------------------------------------
%%% File : mod_offline_riak.erl
%%% Author : Alexey Shchepin <alexey@process-one.net>
%%% Purpose : Store and manage offline messages in Riak.
%%% Created : 4 Jan 2012 by Alexey Shchepin <alexey@process-one.net>
%%%
%%%
%%% ejabberd, Copyright (C) 2002-2011 ProcessOne
%%%
%%% This program is free software; you can redistribute it and/or
%%% modify it under the terms of the GNU General Public License as
%%% published by the Free Software Foundation; either version 2 of the
%%% License, or (at your option) any later version.
%%%
%%% This program is distributed in the hope that it will be useful,
%%% but WITHOUT ANY WARRANTY; without even the implied warranty of
%%% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
%%% General Public License for more details.
%%%
%%% You should have received a copy of the GNU General Public License
%%% along with this program; if not, write to the Free Software
%%% Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
%%% 02111-1307 USA
%%%
%%%----------------------------------------------------------------------
-module(mod_offline_riak).
-author('alexey@process-one.net').
-behaviour(gen_mod).
-export([count_offline_messages/2]).
-export([start/2,
init/2,
stop/1,
store_packet/3,
pop_offline_messages/3,
remove_user/2,
webadmin_page/3,
webadmin_user/4,
webadmin_user_parse_query/5,
count_offline_messages/3]).
-include("ejabberd.hrl").
-include("jlib.hrl").
-include("web/ejabberd_http.hrl").
-include("web/ejabberd_web_admin.hrl").
-record(offline_msg, {user, timestamp, expire, from, to, packet}).
-define(PROCNAME, ejabberd_offline).
-define(OFFLINE_TABLE_LOCK_THRESHOLD, 1000).
start(Host, Opts) ->
ejabberd_hooks:add(offline_message_hook, Host,
?MODULE, store_packet, 50),
ejabberd_hooks:add(resend_offline_messages_hook, Host,
?MODULE, pop_offline_messages, 50),
ejabberd_hooks:add(remove_user, Host,
?MODULE, remove_user, 50),
ejabberd_hooks:add(anonymous_purge_hook, Host,
?MODULE, remove_user, 50),
ejabberd_hooks:add(webadmin_page_host, Host,
?MODULE, webadmin_page, 50),
ejabberd_hooks:add(webadmin_user, Host,
?MODULE, webadmin_user, 50),
ejabberd_hooks:add(webadmin_user_parse_query, Host,
?MODULE, webadmin_user_parse_query, 50),
ejabberd_hooks:add(count_offline_messages, Host,
?MODULE, count_offline_messages, 50),
MaxOfflineMsgs = gen_mod:get_opt(user_max_messages, Opts, infinity),
register(gen_mod:get_module_proc(Host, ?PROCNAME),
spawn(?MODULE, init, [Host, MaxOfflineMsgs])).
%% MaxOfflineMsgs is either infinity of integer > 0
init(Host, infinity) ->
loop(Host, infinity);
init(Host, MaxOfflineMsgs)
when is_integer(MaxOfflineMsgs), MaxOfflineMsgs > 0 ->
loop(Host, MaxOfflineMsgs).
loop(Host, MaxOfflineMsgs) ->
receive
#offline_msg{user = User} = Msg ->
Msgs = receive_all(User, [Msg]),
Len = length(Msgs),
%% Only count existing messages if needed:
Count = if MaxOfflineMsgs =/= infinity ->
Len + count_offline_messages(User, Host);
true -> 0
end,
if
Count > MaxOfflineMsgs ->
discard_warn_sender(Msgs);
true ->
lists:foreach(
fun(M) ->
Username = list_to_binary(User),
From = M#offline_msg.from,
To = M#offline_msg.to,
{xmlelement, Name, Attrs, Els} =
M#offline_msg.packet,
Attrs2 = jlib:replace_from_to_attrs(
jlib:jid_to_string(From),
jlib:jid_to_string(To),
Attrs),
Packet = {xmlelement, Name, Attrs2,
Els ++
[jlib:timestamp_to_xml(
calendar:now_to_universal_time(
M#offline_msg.timestamp))]},
XML =
iolist_to_binary(
xml:element_to_string(Packet)),
{MegaSecs, Secs, MicroSecs} =
M#offline_msg.timestamp,
TS =
iolist_to_binary(
io_lib:format("~6..0w~6..0w.~6..0w",
[MegaSecs, Secs, MicroSecs])),
ejabberd_riak:put(
Host, <<"offline">>,
undefined, XML,
[{<<"user_bin">>, Username},
{<<"timestamp_bin">>, TS}
])
end, Msgs)
end,
loop(Host, MaxOfflineMsgs);
_ ->
loop(Host, MaxOfflineMsgs)
end.
receive_all(Username, Msgs) ->
receive
#offline_msg{user=Username} = Msg ->
receive_all(Username, [Msg | Msgs])
after 0 ->
lists:reverse(Msgs)
end.
stop(Host) ->
ejabberd_hooks:delete(offline_message_hook, Host,
?MODULE, store_packet, 50),
ejabberd_hooks:delete(resend_offline_messages_hook, Host,
?MODULE, pop_offline_messages, 50),
ejabberd_hooks:delete(remove_user, Host,
?MODULE, remove_user, 50),
ejabberd_hooks:delete(anonymous_purge_hook, Host,
?MODULE, remove_user, 50),
ejabberd_hooks:delete(webadmin_page_host, Host,
?MODULE, webadmin_page, 50),
ejabberd_hooks:delete(webadmin_user, Host,
?MODULE, webadmin_user, 50),
ejabberd_hooks:delete(webadmin_user_parse_query, Host,
?MODULE, webadmin_user_parse_query, 50),
Proc = gen_mod:get_module_proc(Host, ?PROCNAME),
exit(whereis(Proc), stop),
ok.
store_packet(From, To, Packet) ->
Type = xml:get_tag_attr_s("type", Packet),
if
(Type /= "error") and (Type /= "groupchat") and
(Type /= "headline") ->
case check_event(From, To, Packet) of
true ->
#jid{luser = LUser} = To,
TimeStamp = now(),
{xmlelement, _Name, _Attrs, Els} = Packet,
Expire = find_x_expire(TimeStamp, Els),
gen_mod:get_module_proc(To#jid.lserver, ?PROCNAME) !
#offline_msg{user = LUser,
timestamp = TimeStamp,
expire = Expire,
from = From,
to = To,
packet = Packet},
stop;
_ ->
ok
end;
true ->
ok
end.
check_event(From, To, Packet) ->
{xmlelement, Name, Attrs, Els} = Packet,
case find_x_event(Els) of
false ->
true;
El ->
case xml:get_subtag(El, "id") of
false ->
case xml:get_subtag(El, "offline") of
false ->
true;
_ ->
ID = case xml:get_tag_attr_s("id", Packet) of
"" ->
{xmlelement, "id", [], []};
S ->
{xmlelement, "id", [],
[{xmlcdata, S}]}
end,
ejabberd_router:route(
To, From, {xmlelement, Name, Attrs,
[{xmlelement, "x",
[{"xmlns", ?NS_EVENT}],
[ID,
{xmlelement, "offline", [], []}]}]
}),
true
end;
_ ->
false
end
end.
find_x_event([]) ->
false;
find_x_event([{xmlcdata, _} | Els]) ->
find_x_event(Els);
find_x_event([El | Els]) ->
case xml:get_tag_attr_s("xmlns", El) of
?NS_EVENT ->
El;
_ ->
find_x_event(Els)
end.
find_x_expire(_, []) ->
never;
find_x_expire(TimeStamp, [{xmlcdata, _} | Els]) ->
find_x_expire(TimeStamp, Els);
find_x_expire(TimeStamp, [El | Els]) ->
case xml:get_tag_attr_s("xmlns", El) of
?NS_EXPIRE ->
Val = xml:get_tag_attr_s("seconds", El),
case catch list_to_integer(Val) of
{'EXIT', _} ->
never;
Int when Int > 0 ->
{MegaSecs, Secs, MicroSecs} = TimeStamp,
S = MegaSecs * 1000000 + Secs + Int,
MegaSecs1 = S div 1000000,
Secs1 = S rem 1000000,
{MegaSecs1, Secs1, MicroSecs};
_ ->
never
end;
_ ->
find_x_expire(TimeStamp, Els)
end.
pop_offline_messages(Ls, User, Server) ->
LUser = jlib:nodeprep(User),
LServer = jlib:nameprep(Server),
Username = list_to_binary(LUser),
case ejabberd_riak:get_objects_by_index(
LServer, <<"offline">>, <<"user_bin">>, Username) of
{ok, Rs} ->
SortedRs =
lists:sort(fun(X, Y) ->
MX = riak_object:get_metadata(X),
{ok, IX} = dict:find(<<"index">>, MX),
{value, TSX} = lists:keysearch(
<<"timestamp_bin">>, 1,
IX),
MY = riak_object:get_metadata(Y),
{ok, IY} = dict:find(<<"index">>, MY),
{value, TSY} = lists:keysearch(
<<"timestamp_bin">>, 1,
IY),
TSX =< TSY
end, Rs),
Ls ++ lists:flatmap(
fun(R) ->
Key = riak_object:key(R),
ejabberd_riak:delete(LServer, <<"offline">>, Key),
XML = riak_object:get_value(R),
case xml_stream:parse_element(XML) of
{error, _Reason} ->
[];
El ->
To = jlib:string_to_jid(
xml:get_tag_attr_s("to", El)),
From = jlib:string_to_jid(
xml:get_tag_attr_s("from", El)),
if
(To /= error) and
(From /= error) ->
[{route, From, To, El}];
true ->
[]
end
end
end, SortedRs);
_ ->
Ls
end.
remove_user(User, Server) ->
LUser = jlib:nodeprep(User),
LServer = jlib:nameprep(Server),
Username = list_to_binary(LUser),
case ejabberd_riak:get_keys_by_index(
LServer, <<"offline">>, <<"user_bin">>, Username) of
{ok, Keys} ->
lists:foreach(
fun(Key) ->
ejabberd_riak:delete(LServer, <<"offline">>, Key)
end, Keys);
_ ->
ok
end.
%% Helper functions:
%% TODO: Warning - This function is a duplicate from mod_offline.erl
%% It is duplicate to stay consistent (many functions are duplicated
%% in this module). It will be refactored later on.
%% Warn senders that their messages have been discarded:
discard_warn_sender(Msgs) ->
lists:foreach(
fun(#offline_msg{from=From, to=To, packet=Packet}) ->
ErrText = "Your contact offline message queue is full. The message has been discarded.",
Lang = xml:get_tag_attr_s("xml:lang", Packet),
Err = jlib:make_error_reply(
Packet, ?ERRT_RESOURCE_CONSTRAINT(Lang, ErrText)),
ejabberd_router:route(
To,
From, Err)
end, Msgs).
webadmin_page(_, Host,
#request{us = _US,
path = ["user", U, "queue"],
q = Query,
lang = Lang} = _Request) ->
Res = user_queue(U, Host, Query, Lang),
{stop, Res};
webadmin_page(Acc, _, _) -> Acc.
user_queue(User, Server, Query, Lang) ->
LUser = jlib:nodeprep(User),
LServer = jlib:nameprep(Server),
Username = ejabberd_odbc:escape(LUser),
US = {LUser, LServer},
Res = user_queue_parse_query(Username, LServer, Query),
Msgs = case catch ejabberd_odbc:sql_query(
LServer,
["select username, xml from spool"
" where username='", Username, "'"
" order by seq;"]) of
{selected, ["username", "xml"], Rs} ->
lists:flatmap(
fun({_, XML}) ->
case xml_stream:parse_element(XML) of
{error, _Reason} ->
[];
El ->
[El]
end
end, Rs);
_ ->
[]
end,
FMsgs =
lists:map(
fun({xmlelement, _Name, _Attrs, _Els} = Msg) ->
ID = jlib:encode_base64(binary_to_list(term_to_binary(Msg))),
Packet = Msg,
FPacket = ejabberd_web_admin:pretty_print_xml(Packet),
?XE("tr",
[?XAE("td", [{"class", "valign"}], [?INPUT("checkbox", "selected", ID)]),
?XAE("td", [{"class", "valign"}], [?XC("pre", FPacket)])]
)
end, Msgs),
[?XC("h1", io_lib:format(?T("~s's Offline Messages Queue"),
[us_to_list(US)]))] ++
case Res of
ok -> [?XREST("Submitted")];
nothing -> []
end ++
[?XAE("form", [{"action", ""}, {"method", "post"}],
[?XE("table",
[?XE("thead",
[?XE("tr",
[?X("td"),
?XCT("td", "Packet")
])]),
?XE("tbody",
if
FMsgs == [] ->
[?XE("tr",
[?XAC("td", [{"colspan", "4"}], " ")]
)];
true ->
FMsgs
end
)]),
?BR,
?INPUTT("submit", "delete", "Delete Selected")
])].
user_queue_parse_query(Username, LServer, Query) ->
case lists:keysearch("delete", 1, Query) of
{value, _} ->
Msgs = case catch ejabberd_odbc:sql_query(
LServer,
["select xml, seq from spool"
" where username='", Username, "'"
" order by seq;"]) of
{selected, ["xml", "seq"], Rs} ->
lists:flatmap(
fun({XML, Seq}) ->
case xml_stream:parse_element(XML) of
{error, _Reason} ->
[];
El ->
[{El, Seq}]
end
end, Rs);
_ ->
[]
end,
F = fun() ->
lists:foreach(
fun({Msg, Seq}) ->
ID = jlib:encode_base64(
binary_to_list(term_to_binary(Msg))),
case lists:member({"selected", ID}, Query) of
true ->
SSeq = ejabberd_odbc:escape(Seq),
catch ejabberd_odbc:sql_query(
LServer,
["delete from spool"
" where username='", Username, "'"
" and seq='", SSeq, "';"]);
false ->
ok
end
end, Msgs)
end,
mnesia:transaction(F),
ok;
false ->
nothing
end.
us_to_list({User, Server}) ->
jlib:jid_to_string({User, Server, ""}).
webadmin_user(Acc, User, Server, Lang) ->
LUser = jlib:nodeprep(User),
LServer = jlib:nameprep(Server),
Username = ejabberd_odbc:escape(LUser),
QueueLen = case catch ejabberd_odbc:sql_query(
LServer,
["select count(*) from spool"
" where username='", Username, "';"]) of
{selected, [_], [{SCount}]} ->
SCount;
_ ->
0
end,
FQueueLen = [?AC("queue/", QueueLen)],
Acc ++ [?XCT("h3", "Offline Messages:")] ++ FQueueLen ++ [?C(" "), ?INPUTT("submit", "removealloffline", "Remove All Offline Messages")].
webadmin_user_parse_query(_, "removealloffline", User, Server, _Query) ->
case catch odbc_queries:del_spool_msg(Server, User) of
{'EXIT', Reason} ->
?ERROR_MSG("Failed to remove offline messages: ~p", [Reason]),
{stop, error};
{error, Reason} ->
?ERROR_MSG("Failed to remove offline messages: ~p", [Reason]),
{stop, error};
_ ->
?INFO_MSG("Removed all offline messages for ~s@~s", [User, Server]),
{stop, ok}
end;
webadmin_user_parse_query(Acc, _Action, _User, _Server, _Query) ->
Acc.
%% ------------------------------------------------
%% mod_offline: number of messages quota management
%% Returns as integer the number of offline messages for a given user
count_offline_messages(LUser, LServer) ->
Username = list_to_binary([LUser, $@, LServer]),
case catch ejabberd_riak:count_by_index(
LServer, <<"offline">>, <<"user_bin">>, Username) of
{ok, Res} when is_integer(Res) ->
Res;
_ ->
0
end.
count_offline_messages(_Acc, User, Server) ->
LUser = jlib:nodeprep(User),
LServer = jlib:nameprep(Server),
Num = case catch ejabberd_odbc:sql_query(
LServer,
["select xml from spool"
" where username='", LUser, "';"]) of
{selected, ["xml"], Rs} ->
lists:foldl(
fun({XML}, Acc) ->
case xml_stream:parse_element(XML) of
{error, _Reason} ->
Acc;
El ->
case xml:get_subtag(El, "body") of
false ->
Acc;
_ ->
Acc + 1
end
end
end, 0, Rs);
_ ->
0
end,
{stop, Num}.

139
src/mod_private_riak.erl Normal file
View File

@ -0,0 +1,139 @@
%%%----------------------------------------------------------------------
%%% File : mod_private_riak.erl
%%% Author : Alexey Shchepin <alexey@process-one.net>
%%% Purpose : Private storage support
%%% Created : 6 Jan 2012 by Alexey Shchepin <alexey@process-one.net>
%%%
%%%
%%% ejabberd, Copyright (C) 2002-2011 ProcessOne
%%%
%%% This program is free software; you can redistribute it and/or
%%% modify it under the terms of the GNU General Public License as
%%% published by the Free Software Foundation; either version 2 of the
%%% License, or (at your option) any later version.
%%%
%%% This program is distributed in the hope that it will be useful,
%%% but WITHOUT ANY WARRANTY; without even the implied warranty of
%%% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
%%% General Public License for more details.
%%%
%%% You should have received a copy of the GNU General Public License
%%% along with this program; if not, write to the Free Software
%%% Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
%%% 02111-1307 USA
%%%
%%%----------------------------------------------------------------------
-module(mod_private_riak).
-author('alexey@process-one.net').
-behaviour(gen_mod).
-export([start/2,
stop/1,
process_sm_iq/3,
remove_user/2]).
-include("ejabberd.hrl").
-include("jlib.hrl").
start(Host, Opts) ->
IQDisc = gen_mod:get_opt(iqdisc, Opts, one_queue),
ejabberd_hooks:add(remove_user, Host,
?MODULE, remove_user, 50),
gen_iq_handler:add_iq_handler(ejabberd_sm, Host, ?NS_PRIVATE,
?MODULE, process_sm_iq, IQDisc).
stop(Host) ->
ejabberd_hooks:delete(remove_user, Host,
?MODULE, remove_user, 50),
gen_iq_handler:remove_iq_handler(ejabberd_sm, Host, ?NS_PRIVATE).
process_sm_iq(From, _To, #iq{type = Type, sub_el = SubEl} = IQ) ->
#jid{luser = LUser, lserver = LServer} = From,
case lists:member(LServer, ?MYHOSTS) of
true ->
{xmlelement, Name, Attrs, Els} = SubEl,
case Type of
set ->
lists:foreach(
fun(El) ->
set_data(LUser, LServer, El)
end, Els),
IQ#iq{type = result,
sub_el = [{xmlelement, Name, Attrs, []}]};
get ->
case catch get_data(LUser, LServer, Els) of
{'EXIT', _Reason} ->
IQ#iq{type = error,
sub_el = [SubEl,
?ERR_INTERNAL_SERVER_ERROR]};
Res ->
IQ#iq{type = result,
sub_el = [{xmlelement, Name, Attrs, Res}]}
end
end;
false ->
IQ#iq{type = error, sub_el = [SubEl, ?ERR_NOT_ALLOWED]}
end.
set_data(LUser, LServer, El) ->
case El of
{xmlelement, _Name, Attrs, _Els} ->
XMLNS = xml:get_attr_s("xmlns", Attrs),
case XMLNS of
"" ->
ignore;
_ ->
Username = list_to_binary(LUser),
Key = list_to_binary([LUser, $@, LServer, $@, XMLNS]),
SData = xml:element_to_binary(El),
ejabberd_riak:put(
LServer, <<"private">>, Key, SData,
[{<<"user_bin">>, Username}]),
ok
end;
_ ->
ignore
end.
get_data(LUser, LServer, Els) ->
get_data(LUser, LServer, Els, []).
get_data(_LUser, _LServer, [], Res) ->
lists:reverse(Res);
get_data(LUser, LServer, [El | Els], Res) ->
case El of
{xmlelement, _Name, Attrs, _} ->
XMLNS = xml:get_attr_s("xmlns", Attrs),
Key = list_to_binary([LUser, $@, LServer, $@, XMLNS]),
case ejabberd_riak:get(LServer, <<"private">>, Key) of
{ok, SData} ->
case xml_stream:parse_element(SData) of
Data when element(1, Data) == xmlelement ->
get_data(LUser, LServer, Els,
[Data | Res])
end;
_ ->
get_data(LUser, LServer, Els, [El | Res])
end;
_ ->
get_data(LUser, LServer, Els, Res)
end.
remove_user(User, Server) ->
LUser = jlib:nodeprep(User),
LServer = jlib:nameprep(Server),
Username = list_to_binary(LUser),
case ejabberd_riak:get_keys_by_index(
LServer, <<"private">>, <<"user_bin">>, Username) of
{ok, Keys} ->
lists:foreach(
fun(Key) ->
ejabberd_riak:delete(LServer, <<"private">>, Key)
end, Keys);
_ ->
ok
end.

1310
src/mod_roster_riak.erl Normal file
View File

File diff suppressed because it is too large Load Diff

209
src/mod_vcard_riak.erl Normal file
View File

@ -0,0 +1,209 @@
%%%----------------------------------------------------------------------
%%% File : mod_vcard_riak.erl
%%% Author : Alexey Shchepin <alexey@process-one.net>
%%% Purpose : vCard support via Riak
%%% Created : 6 Jan 2012 by Alexey Shchepin <alexey@process-one.net>
%%%
%%%
%%% ejabberd, Copyright (C) 2002-2011 ProcessOne
%%%
%%% This program is free software; you can redistribute it and/or
%%% modify it under the terms of the GNU General Public License as
%%% published by the Free Software Foundation; either version 2 of the
%%% License, or (at your option) any later version.
%%%
%%% This program is distributed in the hope that it will be useful,
%%% but WITHOUT ANY WARRANTY; without even the implied warranty of
%%% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
%%% General Public License for more details.
%%%
%%% You should have received a copy of the GNU General Public License
%%% along with this program; if not, write to the Free Software
%%% Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
%%% 02111-1307 USA
%%%
%%%----------------------------------------------------------------------
-module(mod_vcard_riak).
-author('alexey@process-one.net').
-behaviour(gen_mod).
-export([start/2, stop/1,
get_sm_features/5,
process_local_iq/3,
process_sm_iq/3,
remove_user/2]).
-include("ejabberd.hrl").
-include("jlib.hrl").
-define(JUD_MATCHES, 30).
-define(PROCNAME, ejabberd_mod_vcard).
start(Host, Opts) ->
ejabberd_hooks:add(remove_user, Host,
?MODULE, remove_user, 50),
IQDisc = gen_mod:get_opt(iqdisc, Opts, one_queue),
gen_iq_handler:add_iq_handler(ejabberd_local, Host, ?NS_VCARD,
?MODULE, process_local_iq, IQDisc),
gen_iq_handler:add_iq_handler(ejabberd_sm, Host, ?NS_VCARD,
?MODULE, process_sm_iq, IQDisc),
ejabberd_hooks:add(disco_sm_features, Host, ?MODULE, get_sm_features, 50),
ok.
stop(Host) ->
ejabberd_hooks:delete(remove_user, Host,
?MODULE, remove_user, 50),
gen_iq_handler:remove_iq_handler(ejabberd_local, Host, ?NS_VCARD),
gen_iq_handler:remove_iq_handler(ejabberd_sm, Host, ?NS_VCARD),
ejabberd_hooks:delete(disco_sm_features, Host, ?MODULE, get_sm_features, 50),
ok.
get_sm_features({error, _Error} = Acc, _From, _To, _Node, _Lang) ->
Acc;
get_sm_features(Acc, _From, _To, Node, _Lang) ->
case Node of
[] ->
case Acc of
{result, Features} ->
{result, [?NS_VCARD | Features]};
empty ->
{result, [?NS_VCARD]}
end;
_ ->
Acc
end.
process_local_iq(_From, _To, #iq{type = Type, lang = Lang, sub_el = SubEl} = IQ) ->
case Type of
set ->
IQ#iq{type = error, sub_el = [SubEl, ?ERR_NOT_ALLOWED]};
get ->
IQ#iq{type = result,
sub_el = [{xmlelement, "vCard",
[{"xmlns", ?NS_VCARD}],
[{xmlelement, "FN", [],
[{xmlcdata, "ejabberd"}]},
{xmlelement, "URL", [],
[{xmlcdata, ?EJABBERD_URI}]},
{xmlelement, "DESC", [],
[{xmlcdata,
translate:translate(
Lang,
"Erlang Jabber Server") ++
"\nCopyright (c) 2002-2011 ProcessOne"}]},
{xmlelement, "BDAY", [],
[{xmlcdata, "2002-11-16"}]}
]}]}
end.
process_sm_iq(From, To, #iq{type = Type, sub_el = SubEl} = IQ) ->
case Type of
set ->
#jid{user = User, lserver = LServer} = From,
case lists:member(LServer, ?MYHOSTS) of
true ->
set_vcard(User, LServer, SubEl),
IQ#iq{type = result, sub_el = []};
false ->
IQ#iq{type = error, sub_el = [SubEl, ?ERR_NOT_ALLOWED]}
end;
get ->
#jid{luser = LUser, lserver = LServer} = To,
Username = list_to_binary(LUser),
case catch ejabberd_riak:get(LServer, <<"vcard">>, Username) of
{ok, SVCARD} ->
case xml_stream:parse_element(SVCARD) of
{error, _Reason} ->
IQ#iq{type = error,
sub_el = [SubEl, ?ERR_SERVICE_UNAVAILABLE]};
VCARD ->
IQ#iq{type = result, sub_el = [VCARD]}
end;
{error, notfound} ->
IQ#iq{type = result, sub_el = []};
_ ->
IQ#iq{type = error,
sub_el = [SubEl, ?ERR_INTERNAL_SERVER_ERROR]}
end
end.
set_vcard(User, LServer, VCARD) ->
FN = xml:get_path_s(VCARD, [{elem, "FN"}, cdata]),
Family = xml:get_path_s(VCARD, [{elem, "N"}, {elem, "FAMILY"}, cdata]),
Given = xml:get_path_s(VCARD, [{elem, "N"}, {elem, "GIVEN"}, cdata]),
Middle = xml:get_path_s(VCARD, [{elem, "N"}, {elem, "MIDDLE"}, cdata]),
Nickname = xml:get_path_s(VCARD, [{elem, "NICKNAME"}, cdata]),
BDay = xml:get_path_s(VCARD, [{elem, "BDAY"}, cdata]),
CTRY = xml:get_path_s(VCARD, [{elem, "ADR"}, {elem, "CTRY"}, cdata]),
Locality = xml:get_path_s(VCARD, [{elem, "ADR"}, {elem, "LOCALITY"},cdata]),
EMail1 = xml:get_path_s(VCARD, [{elem, "EMAIL"}, {elem, "USERID"},cdata]),
EMail2 = xml:get_path_s(VCARD, [{elem, "EMAIL"}, cdata]),
OrgName = xml:get_path_s(VCARD, [{elem, "ORG"}, {elem, "ORGNAME"}, cdata]),
OrgUnit = xml:get_path_s(VCARD, [{elem, "ORG"}, {elem, "ORGUNIT"}, cdata]),
EMail = case EMail1 of
"" ->
EMail2;
_ ->
EMail1
end,
LUser = jlib:nodeprep(User),
LFN = stringprep:tolower(FN),
LFamily = stringprep:tolower(Family),
LGiven = stringprep:tolower(Given),
LMiddle = stringprep:tolower(Middle),
LNickname = stringprep:tolower(Nickname),
LBDay = stringprep:tolower(BDay),
LCTRY = stringprep:tolower(CTRY),
LLocality = stringprep:tolower(Locality),
LEMail = stringprep:tolower(EMail),
LOrgName = stringprep:tolower(OrgName),
LOrgUnit = stringprep:tolower(OrgUnit),
if
(LUser == error) or
(LFN == error) or
(LFamily == error) or
(LGiven == error) or
(LMiddle == error) or
(LNickname == error) or
(LBDay == error) or
(LCTRY == error) or
(LLocality == error) or
(LEMail == error) or
(LOrgName == error) or
(LOrgUnit == error) ->
{error, badarg};
true ->
Username = list_to_binary(LUser),
SVCARD = xml:element_to_binary(VCARD),
ejabberd_riak:put(
LServer, <<"vcard">>, Username, SVCARD,
[{<<"bday_bin">>, list_to_binary(LBDay)},
{<<"ctry_bin">>, list_to_binary(LCTRY)},
{<<"email_bin">>, list_to_binary(LEMail)},
{<<"fn_bin">>, list_to_binary(LFN)},
{<<"family_bin">>, list_to_binary(LFamily)},
{<<"given_bin">>, list_to_binary(LGiven)},
{<<"locality_bin">>, list_to_binary(LLocality)},
{<<"middle_bin">>, list_to_binary(LMiddle)},
{<<"nickname_bin">>, list_to_binary(LNickname)},
{<<"orgname_bin">>, list_to_binary(LOrgName)},
{<<"orgunit_bin">>, list_to_binary(LOrgUnit)},
{<<"user_bin">>, Username}]),
ejabberd_hooks:run(vcard_set, LServer, [LUser, LServer, VCARD])
end.
remove_user(User, Server) ->
LUser = jlib:nodeprep(User),
LServer = jlib:nameprep(Server),
Username = list_to_binary(LUser),
ejabberd_riak:delete(LServer, <<"vcard">>, Username),
ok.