Updated riak support

This commit is contained in:
Alexey Shchepin 2012-11-06 17:58:08 +02:00 committed by Evgeniy Khramtsov
parent 47763c10e3
commit a4b02c38db
12 changed files with 596 additions and 2244 deletions

View File

@ -57,6 +57,7 @@ start(normal, _Args) ->
connect_nodes(),
Sup = ejabberd_sup:start_link(),
ejabberd_rdbms:start(),
ejabberd_riak_sup:start(),
ejabberd_auth:start(),
cyrsasl:start(),
% Profiling

View File

@ -28,7 +28,7 @@
-author('alexey@process-one.net').
%% External exports
-export([start_link/1,
-export([start_link/3,
put/4,
put/5,
get_object/3,
@ -44,9 +44,9 @@
%%%----------------------------------------------------------------------
%%% API
%%%----------------------------------------------------------------------
start_link(StartInterval) ->
start_link(Server, Port, StartInterval) ->
{ok, Pid} = riakc_pb_socket:start_link(
"127.0.0.1", 8081,
Server, Port,
[auto_reconnect]),
ejabberd_riak_sup:add_pid(Pid),
{ok, Pid}.

View File

@ -50,6 +50,16 @@
-record(riak_pool, {undefined, pid}).
start() ->
StartRiak = ejabberd_config:get_local_option(
riak_server, fun(_) -> true end, false),
if
StartRiak ->
do_start();
true ->
ok
end.
do_start() ->
SupervisorName = ?MODULE,
ChildSpec =
{SupervisorName,
@ -83,36 +93,26 @@ start_link() ->
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,
ejabberd_config:get_local_option(
riak_pool_size,
fun(N) when is_integer(N), N >= 1 -> N end,
?DEFAULT_POOL_SIZE),
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,
ejabberd_config:get_local_option(
riak_start_interval,
fun(N) when is_integer(N), N >= 1 -> N end,
?DEFAULT_RIAK_START_INTERVAL),
{Server, Port} =
ejabberd_config:get_local_option(
riak_server,
fun({S, P}) when is_list(S), is_integer(P), P >= 1 -> {S, P} end,
{"127.0.0.1", 8081}),
{ok, {{one_for_one, PoolSize*10, 1},
lists:map(
fun(I) ->
{I,
{ejabberd_riak, start_link, [StartInterval*1000]},
{ejabberd_riak, start_link,
[Server, Port, StartInterval*1000]},
transient,
2000,
worker,

View File

@ -209,22 +209,26 @@ get_opt_host(Host, Opts, Default) ->
Val = get_opt(host, Opts, fun iolist_to_binary/1, Default),
ejabberd_regexp:greplace(Val, <<"@HOST@">>, Host).
-spec db_type(opts()) -> odbc | mnesia.
-spec db_type(opts()) -> odbc | mnesia | riak.
db_type(Opts) ->
get_opt(db_type, Opts,
fun(odbc) -> odbc;
(internal) -> mnesia;
(mnesia) -> mnesia end,
(mnesia) -> mnesia;
(riak) -> riak
end,
mnesia).
-spec db_type(binary(), atom()) -> odbc | mnesia.
-spec db_type(binary(), atom()) -> odbc | mnesia | riak.
db_type(Host, Module) ->
get_module_opt(Host, Module, db_type,
fun(odbc) -> odbc;
(internal) -> mnesia;
(mnesia) -> mnesia end,
(mnesia) -> mnesia;
(riak) -> riak
end,
mnesia).
-spec loaded_modules(binary()) -> [atom()].

View File

@ -175,6 +175,56 @@ store_offline_msg(Host, {User, _Server}, Msgs, Len, MaxOfflineMsgs, odbc) ->
end,
Msgs),
odbc_queries:add_spool(Host, Query)
end;
store_offline_msg(Host, {User, _}, Msgs, Len, MaxOfflineMsgs,
riak) ->
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 = User,
From = M#offline_msg.from,
To = M#offline_msg.to,
#xmlel{name = Name, attrs = Attrs,
children = Els} =
M#offline_msg.packet,
Attrs2 = jlib:replace_from_to_attrs(
jlib:jid_to_string(From),
jlib:jid_to_string(To),
Attrs),
Packet = #xmlel{name = Name,
attrs = Attrs2,
children =
Els ++
[jlib:timestamp_to_xml(
calendar:now_to_universal_time(
M#offline_msg.timestamp),
utc,
jlib:make_jid(<<"">>, Host, <<"">>),
<<"Offline Storage">>),
jlib:timestamp_to_xml(
calendar:now_to_universal_time(
M#offline_msg.timestamp))]},
XML = xml:element_to_binary(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.
%% Function copied from ejabberd_sm.erl:
@ -193,7 +243,8 @@ receive_all(US, Msgs, DBType) ->
after 0 ->
case DBType of
mnesia -> Msgs;
odbc -> lists:reverse(Msgs)
odbc -> lists:reverse(Msgs);
riak -> lists:reverse(Msgs)
end
end.
@ -421,6 +472,45 @@ pop_offline_messages(Ls, LUser, LServer, odbc) ->
end,
Rs);
_ -> Ls
end;
pop_offline_messages(Ls, LUser, LServer, riak) ->
Username = 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 ->
case offline_msg_to_route(LServer, El) of
error ->
[];
RouteMsg ->
[RouteMsg]
end
end
end, SortedRs);
_ ->
Ls
end.
remove_expired_messages(Server) ->
@ -445,7 +535,8 @@ remove_expired_messages(_LServer, mnesia) ->
ok, offline_msg)
end,
mnesia:transaction(F);
remove_expired_messages(_LServer, odbc) -> {atomic, ok}.
remove_expired_messages(_LServer, odbc) -> {atomic, ok};
remove_expired_messages(_LServer, riak) -> {atomic, ok}.
remove_old_messages(Days, Server) ->
LServer = jlib:nameprep(Server),
@ -470,6 +561,8 @@ remove_old_messages(Days, _LServer, mnesia) ->
end,
mnesia:transaction(F);
remove_old_messages(_Days, _LServer, odbc) ->
{atomic, ok};
remove_old_messages(_Days, _LServer, riak) ->
{atomic, ok}.
remove_user(User, Server) ->
@ -484,7 +577,19 @@ remove_user(LUser, LServer, mnesia) ->
mnesia:transaction(F);
remove_user(LUser, LServer, odbc) ->
Username = ejabberd_odbc:escape(LUser),
odbc_queries:del_spool_msg(LServer, Username).
odbc_queries:del_spool_msg(LServer, Username);
remove_user(LUser, LServer, riak) ->
Username = 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.
jid_to_binary(#jid{user = U, server = S, resource = R,
luser = LU, lserver = LS, lresource = LR}) ->

View File

@ -1,533 +0,0 @@
%%%----------------------------------------------------------------------
%%% 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}.

View File

@ -89,7 +89,8 @@ process_sm_iq(#jid{luser = LUser, lserver = LServer},
end,
case DBType of
odbc -> ejabberd_odbc:sql_transaction(LServer, F);
mnesia -> mnesia:transaction(F)
mnesia -> mnesia:transaction(F);
riak -> F()
end,
IQ#iq{type = result, sub_el = []}
end;
@ -149,7 +150,15 @@ set_data(LUser, LServer, {XMLNS, El}, odbc) ->
LXMLNS = ejabberd_odbc:escape(XMLNS),
SData = ejabberd_odbc:escape(xml:element_to_binary(El)),
odbc_queries:set_private_data(LServer, Username, LXMLNS,
SData).
SData);
set_data(LUser, LServer, {XMLNS, El}, riak) ->
Username = LUser,
Key = <<LUser/binary, $@, LServer/binary, $@, XMLNS/binary>>,
SData = xml:element_to_binary(El),
ejabberd_riak:put(
LServer, <<"private">>, Key, SData,
[{<<"user_bin">>, Username}]),
ok.
get_data(LUser, LServer, Data) ->
get_data(LUser, LServer,
@ -182,10 +191,19 @@ get_data(LUser, LServer, odbc, [{XMLNS, El} | Els],
Data when is_record(Data, xmlel) ->
get_data(LUser, LServer, odbc, Els, [Data | Res])
end;
%% MREMOND: I wonder when the query could return a vcard ?
{selected, [<<"vcard">>], []} ->
get_data(LUser, LServer, odbc, Els, [El | Res]);
_ -> get_data(LUser, LServer, odbc, Els, [El | Res])
end;
get_data(LUser, LServer, riak, [{XMLNS, El} | Els],
Res) ->
Key = <<LUser/binary, $@, LServer/binary, $@, XMLNS/binary>>,
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, riak, Els, [Data | Res])
end;
_ ->
get_data(LUser, LServer, riak, Els, [El | Res])
end.
@ -214,6 +232,23 @@ get_all_data(LUser, LServer, odbc) ->
end, Res);
_ ->
[]
end;
get_all_data(LUser, LServer, riak) ->
Username = LUser,
case ejabberd_riak:get_by_index(
LServer, <<"private">>, <<"user_bin">>, Username) of
{ok, Res} ->
lists:flatmap(
fun(SData) ->
case xml_stream:parse_element(SData) of
#xmlel{} = El ->
[El];
_ ->
[]
end
end, Res);
_ ->
[]
end.
remove_user(User, Server) ->
@ -242,7 +277,19 @@ remove_user(LUser, LServer, mnesia) ->
remove_user(LUser, LServer, odbc) ->
Username = ejabberd_odbc:escape(LUser),
odbc_queries:del_user_private_storage(LServer,
Username).
Username);
remove_user(LUser, LServer, riak) ->
Username = 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.
update_table() ->
Fields = record_info(fields, private_storage),

View File

@ -1,139 +0,0 @@
%%%----------------------------------------------------------------------
%%% 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.

View File

@ -204,6 +204,13 @@ read_roster_version(LUser, LServer, odbc) ->
of
{selected, [<<"version">>], [[Version]]} -> Version;
{selected, [<<"version">>], []} -> error
end;
read_roster_version(LServer, LUser, riak) ->
Username = LUser,
case ejabberd_riak:get(LServer, <<"roster_version">>,
Username) of
{ok, Version} -> Version;
{error, notfound} -> error
end.
write_roster_version(LUser, LServer) ->
@ -239,7 +246,11 @@ write_roster_version(LUser, LServer, InTransaction, Ver,
odbc_queries:set_roster_version(Username,
EVer)
end)
end.
end;
write_roster_version(LUser, LServer, _InTransaction, Ver,
riak) ->
Username = LUser,
riak_set_roster_version(LServer, Username, Ver).
%% Load roster from DB only if neccesary.
%% It is neccesary if
@ -388,6 +399,37 @@ get_roster(LUser, LServer, odbc) ->
Items),
RItems;
_ -> []
end;
get_roster(LUser, LServer, riak) ->
Username = LUser,
case catch riak_get_roster(LServer, Username) of
{ok, Items} when is_list(Items) ->
JIDGroups = case riak_get_roster_jid_groups(LServer, Username) of
{ok, JGrps} when is_list(JGrps) ->
JGrps;
_ ->
[]
end,
GroupsDict = dict:from_list(JIDGroups),
RItems = lists:flatmap(
fun(I) ->
case riak_raw_to_record(LServer, I) of
%% Bad JID in database:
error ->
[];
R ->
SJID = jlib:jid_to_string(R#roster.jid),
Groups =
case dict:find(SJID, GroupsDict) of
{ok, Gs} -> Gs;
error -> []
end,
[R#roster{groups = Groups}]
end
end, Items),
RItems;
_ ->
[]
end.
item_to_xml(Item) ->
@ -455,6 +497,31 @@ get_roster_by_jid_t(LUser, LServer, LJID, odbc) ->
R#roster{usj = {LUser, LServer, LJID},
us = {LUser, LServer}, jid = LJID, name = <<"">>}
end
end;
get_roster_by_jid_t(LUser, LServer, LJID, riak) ->
Username = LUser,
SJID = jlib:jid_to_string(LJID),
Res = riak_get_roster_by_jid(LServer, Username, SJID),
case Res of
{error, _} ->
#roster{usj = {LUser, LServer, LJID},
us = {LUser, LServer},
jid = LJID};
{ok, I} ->
R = riak_raw_to_record(LServer, I),
case R of
%% Bad JID in database:
error ->
#roster{usj = {LUser, LServer, LJID},
us = {LUser, LServer},
jid = LJID};
_ ->
R#roster{
usj = {LUser, LServer, LJID},
us = {LUser, LServer},
jid = LJID,
name = ""}
end
end.
try_process_iq_set(From, To, #iq{sub_el = SubEl} = IQ) ->
@ -631,8 +698,16 @@ get_subscription_lists(_, LUser, LServer, odbc) ->
<<"server">>, <<"subscribe">>, <<"type">>],
Items}
when is_list(Items) ->
Items;
lists:map(fun(I) -> raw_to_record(LServer, I) end, Items);
_ -> []
end;
get_subscription_lists(_, LUser, LServer, riak) ->
Username = LUser,
case catch riak_get_roster(LServer, Username) of
{ok, Items} when is_list(Items) ->
lists:map(fun(I) -> riak_raw_to_record(LServer, I) end, Items);
_ ->
[]
end.
fill_subscription_lists(LServer, [#roster{} = I | Is],
@ -671,12 +746,18 @@ roster_subscribe_t(LUser, LServer, LJID, Item, odbc) ->
Username = ejabberd_odbc:escape(LUser),
SJID = ejabberd_odbc:escape(jlib:jid_to_string(LJID)),
odbc_queries:roster_subscribe(LServer, Username, SJID,
ItemVals).
ItemVals);
roster_subscribe_t(LUser, LServer, LJID, Item, riak) ->
ItemVals = riak_record_to_string(Item),
Username = LUser,
SJID = jlib:jid_to_string(LJID),
riak_roster_subscribe(LServer, Username, SJID, ItemVals).
transaction(LServer, F) ->
case gen_mod:db_type(LServer, ?MODULE) of
mnesia -> mnesia:transaction(F);
odbc -> ejabberd_odbc:sql_transaction(LServer, F)
odbc -> ejabberd_odbc:sql_transaction(LServer, F);
riak -> {atomic, F()}
end.
in_subscription(_, User, Server, JID, Type, Reason) ->
@ -727,6 +808,25 @@ get_roster_by_jid_with_groups_t(LUser, LServer, LJID,
[]} ->
#roster{usj = {LUser, LServer, LJID},
us = {LUser, LServer}, jid = LJID}
end;
get_roster_by_jid_with_groups_t(LUser, LServer, LJID, riak) ->
Username = LUser,
SJID = jlib:jid_to_string(LJID),
case riak_get_roster_by_jid(LServer, Username, SJID) of
{ok, I} ->
R = riak_raw_to_record(LServer, I),
Groups =
case riak_get_roster_groups(LServer, Username, SJID) of
{ok, JGrps} when is_list(JGrps) ->
JGrps;
_ ->
[]
end,
R#roster{groups = Groups};
{error, _} ->
#roster{usj = {LUser, LServer, LJID},
us = {LUser, LServer},
jid = LJID}
end.
process_subscription(Direction, User, Server, JID1,
@ -924,12 +1024,12 @@ in_auto_reply(_, _, _) -> none.
remove_user(User, Server) ->
LUser = jlib:nodeprep(User),
LServer = jlib:nameprep(Server),
send_unsubscription_to_rosteritems(LUser, LServer),
remove_user(LUser, LServer,
gen_mod:db_type(LServer, ?MODULE)).
remove_user(LUser, LServer, mnesia) ->
US = {LUser, LServer},
send_unsubscription_to_rosteritems(LUser, LServer),
F = fun () ->
lists:foreach(fun (R) -> mnesia:delete_object(R) end,
mnesia:index_read(roster, US, #roster.us))
@ -937,8 +1037,11 @@ remove_user(LUser, LServer, mnesia) ->
mnesia:transaction(F);
remove_user(LUser, LServer, odbc) ->
Username = ejabberd_odbc:escape(LUser),
send_unsubscription_to_rosteritems(LUser, LServer),
odbc_queries:del_user_roster_t(LServer, Username),
ok;
remove_user(LUser, LServer, riak) ->
Username = LUser,
riak_del_user_roster(LServer, Username),
ok.
%% For each contact with Subscription:
@ -1009,7 +1112,15 @@ update_roster_t(LUser, LServer, LJID, Item, odbc) ->
SJID = ejabberd_odbc:escape(jlib:jid_to_string(LJID)),
ItemVals = record_to_string(Item),
ItemGroups = groups_to_string(Item),
odbc_queries:update_roster(LServer, Username, SJID, ItemVals, ItemGroups).
odbc_queries:update_roster(LServer, Username, SJID, ItemVals,
ItemGroups);
update_roster_t(LUser, LServer, LJID, Item, riak) ->
Username = LUser,
SJID = jlib:jid_to_string(LJID),
ItemVals = riak_record_to_string(Item),
ItemGroups = riak_groups_to_binary(Item),
riak_update_roster(
LServer, Username, SJID, ItemVals, ItemGroups).
del_roster_t(LUser, LServer, LJID) ->
DBType = gen_mod:db_type(LServer, ?MODULE),
@ -1020,7 +1131,11 @@ del_roster_t(LUser, LServer, LJID, mnesia) ->
del_roster_t(LUser, LServer, LJID, odbc) ->
Username = ejabberd_odbc:escape(LUser),
SJID = ejabberd_odbc:escape(jlib:jid_to_string(LJID)),
odbc_queries:del_roster(LServer, Username, SJID).
odbc_queries:del_roster(LServer, Username, SJID);
del_roster_t(LUser, LServer, LJID, riak) ->
Username = LUser,
SJID = jlib:jid_to_string(LJID),
riak_del_roster(LServer, Username, SJID).
process_item_set_t(LUser, LServer,
#xmlel{attrs = Attrs, children = Els}) ->
@ -1161,6 +1276,44 @@ get_in_pending_subscriptions(Ls, User, Server, odbc) ->
end,
Items));
_ -> Ls
end;
get_in_pending_subscriptions(Ls, User, Server, riak) ->
JID = jlib:make_jid(User, Server, <<"">>),
LUser = JID#jid.luser,
LServer = JID#jid.lserver,
Username = LUser,
case catch riak_get_roster(LServer, Username) of
{ok, Items} when is_list(Items) ->
Ls ++ lists:map(
fun(R) ->
Message = R#roster.askmessage,
#xmlel{name = <<"presence">>,
attrs = [{<<"from">>,
jlib:jid_to_string(R#roster.jid)},
{<<"to">>, jlib:jid_to_string(JID)},
{<<"type">>, <<"subscribe">>}],
children = [#xmlel{name = <<"status">>,
attrs = [],
children =
[{xmlcdata, Message}]}]}
end,
lists:flatmap(
fun(I) ->
case riak_raw_to_record(LServer, I) of
%% Bad JID in database:
error ->
[];
R ->
case R#roster.ask of
in -> [R];
both -> [R];
_ -> []
end
end
end,
Items));
_ ->
Ls
end.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
@ -1205,6 +1358,21 @@ read_subscription_and_groups(LUser, LServer, LJID,
end,
{Subscription, Groups};
_ -> error
end;
read_subscription_and_groups(LUser, LServer, LJID,
riak) ->
Username = LUser,
SJID = jlib:jid_to_string(LJID),
case catch riak_get_subscription(LServer, Username, SJID) of
{ok, Subscription} ->
Groups = case riak_get_roster_jid_groups(LServer, Username) of
{ok, JGrps} when is_list(JGrps) ->
JGrps;
_ ->
[]
end,
{Subscription, Groups};
_ -> error
end.
get_jid_info(_, User, Server, JID) ->
@ -1695,3 +1863,182 @@ import(_LServer, mnesia, #roster_version{} = RV) ->
mnesia:dirty_write(RV);
import(_, _, _) ->
pass.
riak_get_roster(LServer, Username) ->
ejabberd_riak:get_by_index(
LServer, <<"roster">>, <<"user_bin">>, Username).
riak_get_roster_jid_groups(LServer, Username) ->
case ejabberd_riak:get_by_index(
LServer, <<"roster_groups">>, <<"user_bin">>, Username) of
{ok, JGs} ->
Res = lists:map(fun riak_binary_to_groups/1, JGs),
{ok, Res};
Error -> Error
end.
riak_get_roster_groups(LServer, Username, SJID) ->
Key = <<Username/binary, $/, SJID/binary>>,
case ejabberd_riak:get(LServer, <<"roster_groups">>, Key) of
{ok, Gs} ->
{_, Res} = riak_binary_to_groups(Gs),
{ok, Res};
{error, notfound} ->
{ok, []};
Error -> Error
end.
riak_get_roster_by_jid(LServer, Username, SJID) ->
Key = <<Username/binary, $/, SJID/binary>>,
ejabberd_riak:get(LServer, <<"roster">>, Key).
riak_del_roster(LServer, Username, SJID) ->
Key = <<Username/binary, $/, SJID/binary>>,
ejabberd_riak:delete(LServer, <<"roster">>, Key).
riak_update_roster(LServer, Username, SJID, ItemVals, ItemGroups) ->
Key = <<Username/binary, $/, SJID/binary>>,
ejabberd_riak:put(
LServer, <<"roster">>, Key, ItemVals,
[{<<"user_bin">>, Username}]),
ejabberd_riak:put(
LServer, <<"roster_groups">>, Key, ItemGroups,
[{<<"user_bin">>, Username}]).
riak_roster_subscribe(LServer, Username, SJID, ItemVals) ->
Key = <<Username/binary, $/, SJID/binary>>,
ejabberd_riak:put(
LServer, <<"roster">>, Key, ItemVals,
[{<<"user_bin">>, Username}]).
riak_get_subscription(LServer, Username, SJID) ->
case riak_get_roster_by_jid(LServer, Username, SJID) of
{ok, SR} ->
case riak_raw_to_record(LServer, SR) of
error ->
{error, bad_record};
R ->
{ok, R#roster.subscription}
end;
Error ->
Error
end.
riak_set_roster_version(LServer, Username, RosterVersion) ->
ejabberd_riak:put(LServer, <<"roster_version">>,
Username, RosterVersion).
riak_del_user_roster(LServer, Username) ->
case ejabberd_riak:get_keys_by_index(
LServer, <<"roster">>, <<"user_bin">>, Username) of
{ok, Keys} ->
lists:foreach(
fun(Key) ->
ejabberd_riak:delete(LServer, <<"roster">>, Key)
end, Keys);
_ ->
ok
end,
case ejabberd_riak:get_keys_by_index(
LServer, <<"roster_groups">>, <<"user_bin">>, Username) of
{ok, GKeys} ->
lists:foreach(
fun(Key) ->
ejabberd_riak:delete(LServer, <<"roster_groups">>, Key)
end, GKeys);
_ ->
ok
end,
ejabberd_riak:delete(LServer, <<"roster_version">>, Username).
riak_raw_to_record(LServer,
<<UsernameLen:16, Username:UsernameLen/binary,
SJIDLen:16, SJID:SJIDLen/binary,
NickLen:16, Nick:NickLen/binary,
SSubscription, SAsk,
SAskMessageLen:16, SAskMessage:SAskMessageLen/binary>>) ->
User = Username,
case jlib:string_to_jid(SJID) of
error ->
error;
JID ->
LJID = jlib:jid_tolower(JID),
Subscription = case SSubscription of
$B -> both;
$T -> to;
$F -> from;
_ -> none
end,
Ask = case SAsk of
$S -> subscribe;
$U -> unsubscribe;
$B -> both;
$O -> out;
$I -> in;
_ -> none
end,
#roster{usj = {User, LServer, LJID},
us = {User, LServer},
jid = LJID,
name = Nick,
subscription = Subscription,
ask = Ask,
askmessage = SAskMessage}
end.
riak_record_to_string(#roster{us = {User, _Server},
jid = JID,
name = Name,
subscription = Subscription,
ask = Ask,
askmessage = AskMessage}) ->
Username = User,
UsernameLen = size(Username),
SJID = jlib:jid_to_string(jlib:jid_tolower(JID)),
SJIDLen = size(SJID),
Nick = Name,
NickLen = size(Nick),
SSubscription = case Subscription of
both -> $B;
to -> $T;
from -> $F;
none -> $N
end,
SAsk = case Ask of
subscribe -> $S;
unsubscribe -> $U;
both -> $B;
out -> $O;
in -> $I;
none -> $N
end,
SAskMessage = iolist_to_binary(AskMessage),
SAskMessageLen = size(SAskMessage),
<<UsernameLen:16, Username/binary,
SJIDLen:16, SJID/binary,
NickLen:16, Nick/binary,
SSubscription, SAsk,
SAskMessageLen:16, SAskMessage/binary>>.
riak_groups_to_binary(#roster{jid = JID, groups = Groups}) ->
SJID = jlib:jid_to_string(jlib:jid_tolower(JID)),
SJIDLen = size(SJID),
%% Empty groups do not need to be converted to string to be inserted in
%% the database
lists:foldl(
fun([], Acc) ->
Acc;
(Group, Acc) ->
G = Group,
Len = size(G),
<<Acc/binary, Len:16, G/binary>>
end, <<SJIDLen:16, SJID/binary>>, Groups).
riak_binary_to_groups(<<Len:16, SJID:Len/binary, Rest/binary>>) ->
{SJID, riak_binary_to_groups(Rest, [])}.
riak_binary_to_groups(<<Len:16, G:Len/binary, Rest/binary>>, Res) ->
riak_binary_to_groups(Rest, [G | Res]);
riak_binary_to_groups(_, Res) ->
Res.

File diff suppressed because it is too large Load Diff

View File

@ -212,6 +212,19 @@ get_vcard(LUser, LServer, odbc) ->
end;
{selected, [<<"vcard">>], []} -> [];
_ -> error
end;
get_vcard(LUser, LServer, riak) ->
Username = LUser,
case catch ejabberd_riak:get(LServer, <<"vcard">>, Username) of
{ok, SVCARD} ->
case xml_stream:parse_element(SVCARD) of
{error, _Reason} -> error;
VCARD -> [VCARD]
end;
{error, notfound} ->
[];
_ ->
error
end.
set_vcard(User, LServer, VCARD) ->
@ -322,7 +335,25 @@ set_vcard(User, LServer, VCARD) ->
SLGiven, SLLocality, SLMiddle,
SLNickname, SLOrgName, SLOrgUnit,
SLocality, SMiddle, SNickname, SOrgName,
SOrgUnit, SVCARD, Username)
SOrgUnit, SVCARD, Username);
riak ->
Username = LUser,
SVCARD = xml:element_to_binary(VCARD),
ejabberd_riak:put(
LServer, <<"vcard">>, Username, SVCARD,
[{<<"bday_bin">>, LBDay},
{<<"ctry_bin">>, LCTRY},
{<<"email_bin">>, LEMail},
{<<"fn_bin">>, LFN},
{<<"family_bin">>, LFamily},
{<<"given_bin">>, LGiven},
{<<"locality_bin">>, LLocality},
{<<"middle_bin">>, LMiddle},
{<<"nickname_bin">>, LNickname},
{<<"orgname_bin">>, LOrgName},
{<<"orgunit_bin">>, LOrgUnit},
{<<"user_bin">>, Username}])
end,
ejabberd_hooks:run(vcard_set, LServer,
[LUser, LServer, VCARD])
@ -687,14 +718,18 @@ search(LServer, MatchSpec, AllowReturnAll, odbc) ->
Rs;
Error -> ?ERROR_MSG("~p", [Error]), []
end
end.
end;
search(_LServer, _MatchSpec, _AllowReturnAll, riak) ->
[].
make_matchspec(LServer, Data, mnesia) ->
GlobMatch = #vcard_search{_ = '_'},
Match = filter_fields(Data, GlobMatch, LServer, mnesia),
Match;
make_matchspec(LServer, Data, odbc) ->
filter_fields(Data, <<"">>, LServer, odbc).
filter_fields(Data, <<"">>, LServer, odbc);
make_matchspec(_LServer, _Data, riak) ->
[].
filter_fields([], Match, _LServer, mnesia) -> Match;
filter_fields([], Match, _LServer, odbc) ->
@ -884,7 +919,11 @@ remove_user(LUser, LServer, odbc) ->
[[<<"delete from vcard where username='">>,
Username, <<"';">>],
[<<"delete from vcard_search where lusername='">>,
Username, <<"';">>]]).
Username, <<"';">>]]);
remove_user(LUser, LServer, riak) ->
Username = LUser,
ejabberd_riak:delete(LServer, <<"vcard">>, Username),
ok.
update_tables() ->
update_vcard_table(),

View File

@ -1,209 +0,0 @@
%%%----------------------------------------------------------------------
%%% 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.