Preliminary Riak support

src/ejabberd_riak.erl

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

src/ejabberd_riak_sup.erl

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

src/mod_offline_riak.erl

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

src/mod_private_riak.erl

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

src/mod_vcard_riak.erl

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