Updated riak support

src/ejabberd_app.erl

@@ -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

src/ejabberd_riak.erl

@@ -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}.

src/ejabberd_riak_sup.erl

@@ -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
+        ejabberd_config:get_local_option(
-            I when is_integer(I) ->
+          riak_pool_size,
-                I;
+          fun(N) when is_integer(N), N >= 1 -> N end,
-	    undefined ->
+          ?DEFAULT_POOL_SIZE),
-                ?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
+        ejabberd_config:get_local_option(
-            Interval when is_integer(Interval) ->
+          riak_start_interval,
-                Interval;
+          fun(N) when is_integer(N), N >= 1 -> N end,
-            undefined ->
+          ?DEFAULT_RIAK_START_INTERVAL),
-                ?DEFAULT_RIAK_START_INTERVAL;
+    {Server, Port} =
-            _Other2 ->
+        ejabberd_config:get_local_option(
-                ?ERROR_MSG("Wrong riak_start_interval "
+          riak_server,
-                           "definition '~p', "
+          fun({S, P}) when is_list(S), is_integer(P), P >= 1 -> {S, P} end,
-                           "defaulting to ~p~n",
+          {"127.0.0.1", 8081}),
-                           [_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]},
+		     {ejabberd_riak, start_link,
+                      [Server, Port, StartInterval*1000]},
 		     transient,
                      2000,
 		     worker,

src/gen_mod.erl

@@ -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()].

src/mod_offline.erl

@@ -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}) ->

src/mod_offline_riak.erl

@@ -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}.

src/mod_private.erl

@@ -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),

src/mod_private_riak.erl

@@ -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.

src/mod_roster.erl

@@ -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.

src/mod_vcard.erl

@@ -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(),

src/mod_vcard_riak.erl

@@ -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.