Chapril/xmpp.chapril.org-ejabberd
25
1
Fork 0
mirror of https://github.com/processone/ejabberd.git synced 2024-11-24 16:23:40 +01:00
Code Releases Activity

Updated riak support

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
src/ejabberd_app.erl
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

6
src/ejabberd_riak.erl
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}.

50
src/ejabberd_riak_sup.erl
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,

12
src/gen_mod.erl
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()].

111
src/mod_offline.erl
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}) ->

533
src/mod_offline_riak.erl
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}.

59
src/mod_private.erl
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),

139
src/mod_private_riak.erl
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.

363
src/mod_roster.erl
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.

1310
src/mod_roster_riak.erl
View File

File diff suppressed because it is too large Load Diff

47
src/mod_vcard.erl
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(),

209
src/mod_vcard_riak.erl
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.