binary refactoring

1 files changed, 192 insertions, 171 deletions

diff --git a/src/mod_private.erl b/src/mod_private.erl
index 5ca6204f3..51a8b78d0 100644
--- a/src/mod_private.erl
+++ b/src/mod_private.erl
@@ -25,233 +25,254 @@
 %%%----------------------------------------------------------------------
 
 -module(mod_private).
+
 -author('alexey@process-one.net').
 
 -behaviour(gen_mod).
 
--export([start/2,
-	 stop/1,
-	 process_sm_iq/3,
-	 remove_user/2]).
+-export([start/2, stop/1, process_sm_iq/3,
+	 remove_user/2, get_data/2, export/1]).
 
 -include("ejabberd.hrl").
+
 -include("jlib.hrl").
 
--record(private_storage, {usns, xml}).
+-record(private_storage,
+        {usns = {<<"">>, <<"">>, <<"">>} :: {binary(), binary(), binary() |
+                                             '$1' | '_'},
+         xml = #xmlel{} :: xmlel() | '_' | '$1'}).
 
 -define(Xmlel_Query(Attrs, Children),
-(
-    {xmlelement, "query", Attrs, Children}
-)).
+	#xmlel{name = <<"query">>, attrs = Attrs,
+	       children = Children}).
 
 start(Host, Opts) ->
-    IQDisc = gen_mod:get_opt(iqdisc, Opts, one_queue),
+    IQDisc = gen_mod:get_opt(iqdisc, Opts, fun gen_iq_handler:check_type/1,
+                             one_queue),
     case gen_mod:db_type(Opts) of
-        mnesia ->
-            mnesia:create_table(private_storage,
-                                [{disc_only_copies, [node()]},
-                                 {attributes,
-                                  record_info(fields, private_storage)}]),
-            update_table();
-        _ ->
-            ok
+      mnesia ->
+	  mnesia:create_table(private_storage,
+			      [{disc_only_copies, [node()]},
+			       {attributes,
+				record_info(fields, private_storage)}]),
+	  update_table();
+      _ -> ok
     end,
-    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).
+    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).
-
+    ejabberd_hooks:delete(remove_user, Host, ?MODULE,
+			  remove_user, 50),
+    gen_iq_handler:remove_iq_handler(ejabberd_sm, Host,
+				     ?NS_PRIVATE).
 
 process_sm_iq(#jid{luser = LUser, lserver = LServer},
-  #jid{luser = LUser, lserver = LServer}, IQ)
-  when IQ#iq.type == 'set' ->
+	      #jid{luser = LUser, lserver = LServer}, IQ)
+    when IQ#iq.type == set ->
     case IQ#iq.sub_el of
-        {xmlelement, "query", _, Xmlels} ->
-            case filter_xmlels(Xmlels) of
-                [] ->
-                    IQ#iq{
-                      type   = error,
-                      sub_el = [IQ#iq.sub_el, ?ERR_NOT_ACCEPTABLE]
-                    };
-                Data ->
-                    DBType = gen_mod:db_type(LServer, ?MODULE),
-                    F = fun() ->
-                                lists:foreach(
-                                  fun
-                                      (Datum) ->
-                                          set_data(LUser, LServer,
-                                                   Datum, DBType)
-                                  end, Data)
-                        end,
-                    case DBType of
-                        odbc ->
-                            ejabberd_odbc:sql_transaction(LServer, F);
-                        mnesia ->
-                            mnesia:transaction(F)
-                    end,
-                    IQ#iq{type = result, sub_el = []}
-            end;
-        _ ->
-            IQ#iq{type = error, sub_el = [IQ#iq.sub_el, ?ERR_NOT_ACCEPTABLE]}
+      #xmlel{name = <<"query">>, children = Xmlels} ->
+	  case filter_xmlels(Xmlels) of
+	    [] ->
+		IQ#iq{type = error,
+		      sub_el = [IQ#iq.sub_el, ?ERR_NOT_ACCEPTABLE]};
+	    Data ->
+		DBType = gen_mod:db_type(LServer, ?MODULE),
+		F = fun () ->
+			    lists:foreach(fun (Datum) ->
+						  set_data(LUser, LServer,
+							   Datum, DBType)
+					  end,
+					  Data)
+		    end,
+		case DBType of
+		  odbc -> ejabberd_odbc:sql_transaction(LServer, F);
+		  mnesia -> mnesia:transaction(F)
+		end,
+		IQ#iq{type = result, sub_el = []}
+	  end;
+      _ ->
+	  IQ#iq{type = error,
+		sub_el = [IQ#iq.sub_el, ?ERR_NOT_ACCEPTABLE]}
     end;
 %%
 process_sm_iq(#jid{luser = LUser, lserver = LServer},
-  #jid{luser = LUser, lserver = LServer}, IQ)
-  when IQ#iq.type == 'get' ->
+	      #jid{luser = LUser, lserver = LServer}, IQ)
+    when IQ#iq.type == get ->
     case IQ#iq.sub_el of
-        {xmlelement, "query", Attrs, Xmlels} ->
-            case filter_xmlels(Xmlels) of
-                [] ->
-                    IQ#iq{
-                      type   = error,
-                      sub_el = [IQ#iq.sub_el, ?ERR_BAD_FORMAT]
-                    };
-                Data ->
-                    case catch get_data(LUser, LServer, Data) of
-                        {'EXIT', _Reason} ->
-                            IQ#iq{
-                                type   = error,
-                                sub_el = [IQ#iq.sub_el, ?ERR_INTERNAL_SERVER_ERROR]
-                            };
-                        Storage_Xmlels ->
-                            IQ#iq{
-                                type   = result,
-                                sub_el = [?Xmlel_Query(Attrs, Storage_Xmlels)]
-                            }
-                    end
-            end;
-        _ ->
-            IQ#iq{type = error, sub_el = [IQ#iq.sub_el, ?ERR_BAD_FORMAT]}
+      #xmlel{name = <<"query">>, attrs = Attrs,
+	     children = Xmlels} ->
+	  case filter_xmlels(Xmlels) of
+	    [] ->
+		IQ#iq{type = error,
+		      sub_el = [IQ#iq.sub_el, ?ERR_BAD_FORMAT]};
+	    Data ->
+		case catch get_data(LUser, LServer, Data) of
+		  {'EXIT', _Reason} ->
+		      IQ#iq{type = error,
+			    sub_el =
+				[IQ#iq.sub_el, ?ERR_INTERNAL_SERVER_ERROR]};
+		  Storage_Xmlels ->
+		      IQ#iq{type = result,
+			    sub_el = [?Xmlel_Query(Attrs, Storage_Xmlels)]}
+		end
+	  end;
+      _ ->
+	  IQ#iq{type = error,
+		sub_el = [IQ#iq.sub_el, ?ERR_BAD_FORMAT]}
     end;
 %%
 process_sm_iq(_From, _To, IQ) ->
-    IQ#iq{type = error, sub_el = [IQ#iq.sub_el, ?ERR_FORBIDDEN]}.
+    IQ#iq{type = error,
+	  sub_el = [IQ#iq.sub_el, ?ERR_FORBIDDEN]}.
 
+filter_xmlels(Xmlels) -> filter_xmlels(Xmlels, []).
 
-filter_xmlels(Xmlels) ->
-    filter_xmlels(Xmlels, []).
-
-filter_xmlels([], Data) ->
-    lists:reverse(Data);
-filter_xmlels([{xmlelement, _, Attrs, _} = Xmlel | Xmlels], Data) ->
-    case xml:get_attr_s("xmlns", Attrs) of
-        ""    -> [];
-        XmlNS -> filter_xmlels(Xmlels, [{XmlNS, Xmlel} | Data])
+filter_xmlels([], Data) -> lists:reverse(Data);
+filter_xmlels([#xmlel{attrs = Attrs} = Xmlel | Xmlels],
+	      Data) ->
+    case xml:get_attr_s(<<"xmlns">>, Attrs) of
+      <<"">> -> [];
+      XmlNS -> filter_xmlels(Xmlels, [{XmlNS, Xmlel} | Data])
     end;
 filter_xmlels([_ | Xmlels], Data) ->
     filter_xmlels(Xmlels, Data).
 
-
 set_data(LUser, LServer, {XmlNS, Xmlel}, mnesia) ->
-    mnesia:write(#private_storage{
-                    usns = {LUser, LServer, XmlNS},
-                    xml  = Xmlel});
+    mnesia:write(#private_storage{usns =
+				      {LUser, LServer, XmlNS},
+				  xml = Xmlel});
 set_data(LUser, LServer, {XMLNS, El}, odbc) ->
     Username = ejabberd_odbc:escape(LUser),
     LXMLNS = ejabberd_odbc:escape(XMLNS),
-    SData = ejabberd_odbc:escape(
-              xml:element_to_binary(El)),
-    odbc_queries:set_private_data(LServer, Username, LXMLNS, SData).
+    SData = ejabberd_odbc:escape(xml:element_to_binary(El)),
+    odbc_queries:set_private_data(LServer, Username, LXMLNS,
+				  SData).
 
 get_data(LUser, LServer, Data) ->
-    get_data(LUser, LServer, gen_mod:db_type(LServer, ?MODULE), Data, []).
+    get_data(LUser, LServer,
+	     gen_mod:db_type(LServer, ?MODULE), Data, []).
 
-get_data(_LUser, _LServer, _DBType, [], Storage_Xmlels) ->
+get_data(_LUser, _LServer, _DBType, [],
+	 Storage_Xmlels) ->
     lists:reverse(Storage_Xmlels);
-get_data(LUser, LServer, mnesia, [{XmlNS, Xmlel} | Data], Storage_Xmlels) ->
-    case mnesia:dirty_read(private_storage, {LUser, LServer, XmlNS}) of
-        [#private_storage{xml = Storage_Xmlel}] ->
-            get_data(LUser, LServer, mnesia, Data,
-                     [Storage_Xmlel | Storage_Xmlels]);
-        _ ->
-            get_data(LUser, LServer, mnesia, Data,
-                     [Xmlel | Storage_Xmlels])
+get_data(LUser, LServer, mnesia,
+	 [{XmlNS, Xmlel} | Data], Storage_Xmlels) ->
+    case mnesia:dirty_read(private_storage,
+			   {LUser, LServer, XmlNS})
+	of
+      [#private_storage{xml = Storage_Xmlel}] ->
+	  get_data(LUser, LServer, mnesia, Data,
+		   [Storage_Xmlel | Storage_Xmlels]);
+      _ ->
+	  get_data(LUser, LServer, mnesia, Data,
+		   [Xmlel | Storage_Xmlels])
     end;
-get_data(LUser, LServer, odbc, [{XMLNS, El} | Els], Res) ->
+get_data(LUser, LServer, odbc, [{XMLNS, El} | Els],
+	 Res) ->
     Username = ejabberd_odbc:escape(LUser),
     LXMLNS = ejabberd_odbc:escape(XMLNS),
-    case catch odbc_queries:get_private_data(LServer, Username, LXMLNS) of
-        {selected, ["data"], [{SData}]} ->
-            case xml_stream:parse_element(SData) of
-                Data when element(1, Data) == xmlelement ->
-                    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])
+    case catch odbc_queries:get_private_data(LServer,
+					     Username, LXMLNS)
+	of
+      {selected, [<<"data">>], [[SData]]} ->
+	  case xml_stream:parse_element(SData) of
+	    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) ->
+    get_all_data(LUser, LServer,
+                 gen_mod:db_type(LServer, ?MODULE)).
+
+get_all_data(LUser, LServer, mnesia) ->
+    lists:flatten(
+      mnesia:dirty_select(private_storage,
+                          [{#private_storage{usns = {LUser, LServer, '_'},
+                                             xml = '$1'},
+                            [], ['$1']}]));
+get_all_data(LUser, LServer, odbc) ->
+    Username = ejabberd_odbc:escape(LUser),
+    case catch odbc_queries:get_private_data(LServer, Username) of
+        {selected, [<<"namespace">>, <<"data">>], Res} ->
+            lists:flatmap(
+              fun([_, SData]) ->
+                      case xml_stream:parse_element(SData) of
+                          #xmlel{} = El ->
+                              [El];
+                          _ ->
+                              []
+                      end
+              end, Res);
+        _ ->
+            []
     end.
 
 remove_user(User, Server) ->
     LUser = jlib:nodeprep(User),
     LServer = jlib:nameprep(Server),
-    remove_user(LUser, LServer, gen_mod:db_type(Server, ?MODULE)).
+    remove_user(LUser, LServer,
+		gen_mod:db_type(Server, ?MODULE)).
 
 remove_user(LUser, LServer, mnesia) ->
-    F = fun() ->
-		Namespaces = mnesia:select(
-			    private_storage,
-			    [{#private_storage{usns={LUser, LServer, '$1'},
-					       _ = '_'},
-			     [],
-			     ['$$']}]),
-		lists:foreach(
-		  fun([Namespace]) ->
-			  mnesia:delete({private_storage,
-					 {LUser, LServer, Namespace}})
-		     end, Namespaces)
-        end,
+    F = fun () ->
+		Namespaces = mnesia:select(private_storage,
+					   [{#private_storage{usns =
+								  {LUser,
+								   LServer,
+								   '$1'},
+							      _ = '_'},
+					     [], ['$$']}]),
+		lists:foreach(fun ([Namespace]) ->
+				      mnesia:delete({private_storage,
+						     {LUser, LServer,
+						      Namespace}})
+			      end,
+			      Namespaces)
+	end,
     mnesia:transaction(F);
 remove_user(LUser, LServer, odbc) ->
     Username = ejabberd_odbc:escape(LUser),
-    odbc_queries:del_user_private_storage(LServer, Username).
+    odbc_queries:del_user_private_storage(LServer,
+					  Username).
 
 update_table() ->
     Fields = record_info(fields, private_storage),
     case mnesia:table_info(private_storage, attributes) of
-	Fields ->
-	    ok;
-	[userns, xml] ->
-	    ?INFO_MSG("Converting private_storage table from "
-		      "{user, default, lists} format", []),
-	    Host = ?MYNAME,
-	    {atomic, ok} = mnesia:create_table(
-			     mod_private_tmp_table,
-			     [{disc_only_copies, [node()]},
-			      {type, bag},
-			      {local_content, true},
-			      {record_name, private_storage},
-			      {attributes, record_info(fields, private_storage)}]),
-	    mnesia:transform_table(private_storage, ignore, Fields),
-	    F1 = fun() ->
-			 mnesia:write_lock_table(mod_private_tmp_table),
-			 mnesia:foldl(
-			   fun(#private_storage{usns = {U, NS}} = R, _) ->
-				   mnesia:dirty_write(
-				     mod_private_tmp_table,
-				     R#private_storage{usns = {U, Host, NS}})
-			   end, ok, private_storage)
-		 end,
-	    mnesia:transaction(F1),
-	    mnesia:clear_table(private_storage),
-	    F2 = fun() ->
-			 mnesia:write_lock_table(private_storage),
-			 mnesia:foldl(
-			   fun(R, _) ->
-				   mnesia:dirty_write(R)
-			   end, ok, mod_private_tmp_table)
-		 end,
-	    mnesia:transaction(F2),
-	    mnesia:delete_table(mod_private_tmp_table);
-	_ ->
-	    ?INFO_MSG("Recreating private_storage table", []),
-	    mnesia:transform_table(private_storage, ignore, Fields)
+      Fields ->
+          ejabberd_config:convert_table_to_binary(
+            private_storage, Fields, set,
+            fun(#private_storage{usns = {U, _, _}}) -> U end,
+            fun(#private_storage{usns = {U, S, NS}, xml = El} = R) ->
+                    R#private_storage{usns = {iolist_to_binary(U),
+                                              iolist_to_binary(S),
+                                              iolist_to_binary(NS)},
+                                      xml = xml:to_xmlel(El)}
+            end);
+      _ ->
+	  ?INFO_MSG("Recreating private_storage table", []),
+	  mnesia:transform_table(private_storage, ignore, Fields)
     end.
 
-
+export(_Server) ->
+    [{private_storage,
+      fun(Host, #private_storage{usns = {LUser, LServer, XMLNS},
+                                 xml = Data})
+            when LServer == Host ->
+              Username = ejabberd_odbc:escape(LUser),
+              LXMLNS = ejabberd_odbc:escape(XMLNS),
+              SData =
+                  ejabberd_odbc:escape(xml:element_to_binary(Data)),
+              odbc_queries:set_private_data_sql(Username, LXMLNS,
+                                                SData);
+         (_Host, _R) ->
+              []
+      end}].