Use ets and maps instead of dict

2024-11-24 16:23:40 +01:00 · 2019-07-08 09:46:50 +03:00 · 2019-07-08 09:46:50 +03:00 · 92ab59a581
commit 92ab59a581
parent 80beb6d6f6
2 changed files with 108 additions and 80 deletions
--- a/src/mod_delegation.erl
+++ b/src/mod_delegation.erl
@ -45,9 +45,9 @@
 -include("translate.hrl").

 -type disco_acc() :: {error, stanza_error()} | {result, [binary()]} | empty.
-record(state, {server_host = <<"">> :: binary(),
-		delegations = dict:new() :: dict:dict()}).
-type state() :: #state{}.
+-type route_type() :: ejabberd_sm | ejabberd_local.
+-type delegations() :: #{{binary(), route_type()} => {binary(), disco_info()}}.
+-record(state, {server_host = <<"">> :: binary()}).

 %%%===================================================================
 %%% API
@ -136,6 +136,9 @@ disco_sm_identity(Acc, From, To, Node, Lang) ->
 %%%===================================================================
 init([Host, _Opts]) ->
    process_flag(trap_exit, true),
+    catch ets:new(?MODULE,
+                  [named_table, public,
+                   {heir, erlang:group_leader(), none}]),
    ejabberd_hooks:add(component_connected, ?MODULE,
 		       component_connected, 50),
    ejabberd_hooks:add(component_disconnected, ?MODULE,
@ -150,11 +153,9 @@ init([Host, _Opts]) ->
 		       disco_sm_identity, 50),
    {ok, #state{server_host = Host}}.

-handle_call(get_delegations, _From, State) ->
-    {reply, {ok, State#state.delegations}, State};
-handle_call(_Request, _From, State) ->
-    Reply = ok,
-    {reply, Reply, State}.
+handle_call(Request, From, State) ->
+    ?WARNING_MSG("Unexpected call from ~p: ~p", [From, Request]),
+    {noreply, State}.

 handle_cast({component_connected, Host}, State) ->
    ServerHost = State#state.server_host,
@ -166,15 +167,14 @@ handle_cast({component_connected, Host}, State) ->
 		  allow ->
 		      send_disco_queries(ServerHost, Host, NS);
 		  deny ->
-		      ?DEBUG("Denied delegation for ~s on ~s", [Host, NS]),
-		      ok
+		      ?DEBUG("Denied delegation for ~s on ~s", [Host, NS])
 	      end
      end, NSAttrsAccessList),
    {noreply, State};
 handle_cast({component_disconnected, Host}, State) ->
    ServerHost = State#state.server_host,
    Delegations =
-	dict:filter(
+	maps:filter(
 	  fun({NS, Type}, {H, _}) when H == Host ->
 		  ?INFO_MSG("Remove delegation of namespace '~s' "
 			    "from external component '~s'",
@ -183,24 +183,27 @@ handle_cast({component_disconnected, Host}, State) ->
 		  false;
 	     (_, _) ->
 		  true
-	  end, State#state.delegations),
-    {noreply, State#state{delegations = Delegations}};
-handle_cast(_Msg, State) ->
+	  end, get_delegations(ServerHost)),
+    set_delegations(ServerHost, Delegations),
+    {noreply, State};
+handle_cast(Msg, State) ->
+    ?WARNING_MSG("Unexpected cast: ~p", [Msg]),
    {noreply, State}.

 handle_info({iq_reply, ResIQ, {disco_info, Type, Host, NS}}, State) ->
-    {noreply,
-     case ResIQ of
-	 #iq{type = result, sub_els = [SubEl]} ->
-	     try xmpp:decode(SubEl) of
-		 #disco_info{} = Info ->
-		     process_disco_info(State, Type, Host, NS, Info)
-		 catch _:{xmpp_codec, _} ->
-			 State
-		 end;
-	 _ ->
-	     State
-     end};
+    case ResIQ of
+	#iq{type = result, sub_els = [SubEl]} ->
+	    try xmpp:decode(SubEl) of
+		#disco_info{} = Info ->
+		    ServerHost = State#state.server_host,
+		    process_disco_info(ServerHost, Type, Host, NS, Info)
+	    catch _:{xmpp_codec, _} ->
+		    ok
+	    end;
+	_ ->
+	    ok
+    end,
+    {noreply, State};
 handle_info({iq_reply, ResIQ, #iq{} = IQ}, State) ->
    process_iq_result(IQ, ResIQ),
    {noreply, State};
@ -223,7 +226,8 @@ terminate(_Reason, State) ->
    lists:foreach(
      fun({NS, Type}) ->
 	      gen_iq_handler:remove_iq_handler(Type, ServerHost, NS)
-      end, dict:fetch_keys(State#state.delegations)).
+      end, maps:keys(get_delegations(ServerHost))),
+    ets:delete(?MODULE, ServerHost).

 code_change(_OldVsn, State, _Extra) ->
    {ok, State}.
@ -231,22 +235,25 @@ code_change(_OldVsn, State, _Extra) ->
 %%%===================================================================
 %%% Internal functions
 %%%===================================================================
-spec get_delegations(binary()) -> dict:dict().
+-spec get_delegations(binary()) -> delegations().
 get_delegations(Host) ->
-    Proc = gen_mod:get_module_proc(Host, ?MODULE),
-    try gen_server:call(Proc, get_delegations) of
-	{ok, Delegations} -> Delegations
-    catch exit:{noproc, _} ->
-	    %% No module is loaded for this virtual host
-	    dict:new()
+    try ets:lookup_element(?MODULE, Host, 2)
+    catch _:badarg -> #{}
    end.

-spec process_iq(iq(), ejabberd_local | ejabberd_sm) -> ignore | iq().
+-spec set_delegations(binary(), delegations()) -> true.
+set_delegations(ServerHost, Delegations) ->
+    case maps:size(Delegations) of
+	0 -> ets:delete(?MODULE, ServerHost);
+	_ -> ets:insert(?MODULE, {ServerHost, Delegations})
+    end.
+
+-spec process_iq(iq(), route_type()) -> ignore | iq().
 process_iq(#iq{to = To, lang = Lang, sub_els = [SubEl]} = IQ, Type) ->
    LServer = To#jid.lserver,
    NS = xmpp:get_ns(SubEl),
    Delegations = get_delegations(LServer),
-    case dict:find({NS, Type}, Delegations) of
+    case maps:find({NS, Type}, Delegations) of
 	{ok, {Host, _}} ->
 	    Delegation = #delegation{
 			    forwarded = #forwarded{sub_els = [IQ]}},
@ -291,28 +298,27 @@ process_iq_result(#iq{lang = Lang} = IQ, timeout) ->
    Err = xmpp:err_internal_server_error(Txt, Lang),
    ejabberd_router:route_error(IQ, Err).

-spec process_disco_info(state(), ejabberd_local | ejabberd_sm,
-			 binary(), binary(), disco_info()) -> state().
-process_disco_info(State, Type, Host, NS, Info) ->
-    From = jid:make(State#state.server_host),
+-spec process_disco_info(binary(), route_type(),
+			 binary(), binary(), disco_info()) -> ok.
+process_disco_info(ServerHost, Type, Host, NS, Info) ->
+    From = jid:make(ServerHost),
    To = jid:make(Host),
-    case dict:find({NS, Type}, State#state.delegations) of
+    Delegations = get_delegations(ServerHost),
+    case maps:find({NS, Type}, Delegations) of
 	error ->
 	    Msg = #message{from = From, to = To,
 			   sub_els = [#delegation{delegated = [#delegated{ns = NS}]}]},
-	    Delegations = dict:store({NS, Type}, {Host, Info}, State#state.delegations),
-	    gen_iq_handler:add_iq_handler(Type, State#state.server_host, NS,
-					  ?MODULE, Type),
+	    Delegations1 = maps:put({NS, Type}, {Host, Info}, Delegations),
+	    gen_iq_handler:add_iq_handler(Type, ServerHost, NS, ?MODULE, Type),
 	    ejabberd_router:route(Msg),
+	    set_delegations(ServerHost, Delegations1),
 	    ?INFO_MSG("Namespace '~s' is delegated to external component '~s'",
-		      [NS, Host]),
-	    State#state{delegations = Delegations};
+		      [NS, Host]);
 	{ok, {AnotherHost, _}} ->
 	    ?WARNING_MSG("Failed to delegate namespace '~s' to "
 			 "external component '~s' because it's already "
 			 "delegated to '~s'",
-			 [NS, Host, AnotherHost]),
-	    State
+			 [NS, Host, AnotherHost])
    end.

 -spec send_disco_queries(binary(), binary(), binary()) -> ok.
@ -330,7 +336,7 @@ send_disco_queries(LServer, Host, NS) ->
 	    {ejabberd_sm, <<(?NS_DELEGATION)/binary, ":bare:", NS/binary>>}]).

 -spec disco_features(disco_acc(), jid(), jid(), binary(), binary(),
-		     ejabberd_local | ejabberd_sm) -> disco_acc().
+		     route_type()) -> disco_acc().
 disco_features(Acc, _From, To, <<"">>, _Lang, Type) ->
    Delegations = get_delegations(To#jid.lserver),
    Features = my_features(Type) ++
@ -339,7 +345,7 @@ disco_features(Acc, _From, To, <<"">>, _Lang, Type) ->
 		  Info#disco_info.features;
 	     (_) ->
 		  []
-	  end, dict:to_list(Delegations)),
+	  end, maps:to_list(Delegations)),
    case Acc of
 	empty when Features /= [] -> {result, Features};
 	{result, Fs} -> {result, Fs ++ Features};
@ -349,7 +355,7 @@ disco_features(Acc, _, _, _, _, _) ->
    Acc.

 -spec disco_identity(disco_acc(), jid(), jid(), binary(), binary(),
-		     ejabberd_local | ejabberd_sm) -> disco_acc().
+		     route_type()) -> disco_acc().
 disco_identity(Acc, _From, To, <<"">>, _Lang, Type) ->
    Delegations = get_delegations(To#jid.lserver),
    Identities = lists:flatmap(
@ -357,7 +363,7 @@ disco_identity(Acc, _From, To, <<"">>, _Lang, Type) ->
 			   Info#disco_info.identities;
 		      (_) ->
 			   []
-		   end, dict:to_list(Delegations)),
+		   end, maps:to_list(Delegations)),
    case Acc of
 	empty when Identities /= [] -> {result, Identities};
 	{result, Ids} -> {result, Ids ++ Identities};
--- a/src/mod_privilege.erl
+++ b/src/mod_privilege.erl
@ -43,8 +43,17 @@
 -include("xmpp.hrl").
 -include("translate.hrl").

-record(state, {server_host = <<"">> :: binary(),
-		permissions = dict:new() :: dict:dict()}).
+-type roster_permission() :: both | get | set.
+-type presence_permission() :: managed_entity | roster.
+-type message_permission() :: outgoing.
+-type roster_permissions() :: [{roster_permission(), acl:acl()}].
+-type presence_permissions() :: [{presence_permission(), acl:acl()}].
+-type message_permissions() :: [{message_permission(), acl:acl()}].
+-type access() :: [{roster, roster_permissions()} |
+		   {presence, presence_permissions()} |
+		   {message, message_permissions()}].
+-type permissions() :: #{binary() => access()}.
+-record(state, {server_host = <<"">> :: binary()}).

 %%%===================================================================
 %%% API
@ -99,7 +108,7 @@ process_message(#message{from = #jid{luser = <<"">>, lresource = <<"">>} = From,
    Host = From#jid.lserver,
    ServerHost = To#jid.lserver,
    Permissions = get_permissions(ServerHost),
-    case dict:find(Host, Permissions) of
+    case maps:find(Host, Permissions) of
 	{ok, Access} ->
 	    case proplists:get_value(message, Access, none) of
 		outgoing ->
@ -124,7 +133,7 @@ roster_access(false, #iq{from = From, to = To, type = Type}) ->
    Host = From#jid.lserver,
    ServerHost = To#jid.lserver,
    Permissions = get_permissions(ServerHost),
-    case dict:find(Host, Permissions) of
+    case maps:find(Host, Permissions) of
 	{ok, Access} ->
 	    Permission = proplists:get_value(roster, Access, none),
 	    (Permission == both)
@ -153,7 +162,7 @@ process_presence_out({#presence{
 		 true ->
 		      ok
 	      end
-      end, dict:to_list(Permissions)),
+      end, maps:to_list(Permissions)),
    {Pres, C2SState};
 process_presence_out(Acc) ->
    Acc.
@ -181,7 +190,7 @@ process_presence_in({#presence{
 		 _ ->
 		      ok
 	      end
-      end, dict:to_list(Permissions)),
+      end, maps:to_list(Permissions)),
    {Pres, C2SState};
 process_presence_in(Acc) ->
    Acc.
@ -191,6 +200,9 @@ process_presence_in(Acc) ->
 %%%===================================================================
 init([Host, _Opts]) ->
    process_flag(trap_exit, true),
+    catch ets:new(?MODULE,
+                  [named_table, public,
+                   {heir, erlang:group_leader(), none}]),
    ejabberd_hooks:add(component_connected, ?MODULE,
                       component_connected, 50),
    ejabberd_hooks:add(component_disconnected, ?MODULE,
@ -205,11 +217,9 @@ init([Host, _Opts]) ->
 		       process_presence_in, 50),
    {ok, #state{server_host = Host}}.

-handle_call(get_permissions, _From, State) ->
-    {reply, {ok, State#state.permissions}, State};
-handle_call(_Request, _From, State) ->
-    Reply = ok,
-    {reply, Reply, State}.
+handle_call(Request, From, State) ->
+    ?WARNING_MSG("Unexpected call from ~p: ~p", [From, Request]),
+    {noreply, State}.

 handle_cast({component_connected, Host}, State) ->
    ServerHost = State#state.server_host,
@ -231,23 +241,31 @@ handle_cast({component_connected, Host}, State) ->
 		      [Host, RosterPerm, PresencePerm, MessagePerm]),
 	    Msg = #message{from = From, to = To,  sub_els = [Priv]},
 	    ejabberd_router:route(Msg),
-	    Permissions = dict:store(Host, [{roster, RosterPerm},
-					    {presence, PresencePerm},
-					    {message, MessagePerm}],
-				     State#state.permissions),
-	    {noreply, State#state{permissions = Permissions}};
+	    Permissions = maps:put(Host, [{roster, RosterPerm},
+					  {presence, PresencePerm},
+					  {message, MessagePerm}],
+				   get_permissions(ServerHost)),
+	    ets:insert(?MODULE, {ServerHost, Permissions}),
+	    {noreply, State};
       true ->
 	    ?INFO_MSG("Granting no permissions to external component '~s'",
 		      [Host]),
 	    {noreply, State}
    end;
 handle_cast({component_disconnected, Host}, State) ->
-    Permissions = dict:erase(Host, State#state.permissions),
-    {noreply, State#state{permissions = Permissions}};
-handle_cast(_Msg, State) ->
+    ServerHost = State#state.server_host,
+    Permissions = maps:remove(Host, get_permissions(ServerHost)),
+    case maps:size(Permissions) of
+	0 -> ets:delete(?MODULE, ServerHost);
+	_ -> ets:insert(?MODULE, {ServerHost, Permissions})
+    end,
+    {noreply, State};
+handle_cast(Msg, State) ->
+    ?WARNING_MSG("Unexpected cast: ~p", [Msg]),
    {noreply, State}.

-handle_info(_Info, State) ->
+handle_info(Info, State) ->
+    ?WARNING_MSG("Unexpected info: ~p", [Info]),
    {noreply, State}.

 terminate(_Reason, State) ->
@ -261,7 +279,8 @@ terminate(_Reason, State) ->
    ejabberd_hooks:delete(user_send_packet, Host, ?MODULE,
 			  process_presence_out, 50),
    ejabberd_hooks:delete(user_receive_packet, Host, ?MODULE,
-			  process_presence_in, 50).
+			  process_presence_in, 50),
+    ets:delete(?MODULE, Host).

 code_change(_OldVsn, State, _Extra) ->
    {ok, State}.
@ -269,16 +288,13 @@ code_change(_OldVsn, State, _Extra) ->
 %%%===================================================================
 %%% Internal functions
 %%%===================================================================
+-spec get_permissions(binary()) -> permissions().
 get_permissions(ServerHost) ->
-    Proc = gen_mod:get_module_proc(ServerHost, ?MODULE),
-    try gen_server:call(Proc, get_permissions) of
-	{ok, Permissions} ->
-	    Permissions
-    catch exit:{noproc, _} ->
-	    %% No module is loaded for this virtual host
-	    dict:new()
+    try ets:lookup_element(?MODULE, ServerHost, 2)
+    catch _:badarg -> #{}
    end.

+-spec forward_message(message()) -> ok.
 forward_message(#message{to = To} = Msg) ->
    ServerHost = To#jid.lserver,
    Lang = xmpp:get_lang(Msg),
@ -315,6 +331,7 @@ forward_message(#message{to = To} = Msg) ->
 	    ejabberd_router:route_error(Msg, Err)
    end.

+-spec get_roster_permission(binary(), binary()) -> roster_permission() | none.
 get_roster_permission(ServerHost, Host) ->
    Perms = mod_privilege_opt:roster(ServerHost),
    case match_rule(ServerHost, Host, Perms, both) of
@ -330,6 +347,7 @@ get_roster_permission(ServerHost, Host) ->
 	    end
    end.

+-spec get_message_permission(binary(), binary()) -> message_permission() | none.
 get_message_permission(ServerHost, Host) ->
    Perms = mod_privilege_opt:message(ServerHost),
    case match_rule(ServerHost, Host, Perms, outgoing) of
@ -337,6 +355,7 @@ get_message_permission(ServerHost, Host) ->
 	deny -> none
    end.

+-spec get_presence_permission(binary(), binary()) -> presence_permission() | none.
 get_presence_permission(ServerHost, Host) ->
    Perms = mod_privilege_opt:presence(ServerHost),
    case match_rule(ServerHost, Host, Perms, roster) of
@ -349,6 +368,9 @@ get_presence_permission(ServerHost, Host) ->
 	    end
    end.

+-spec match_rule(binary(), binary(), roster_permissions(), roster_permission()) -> allow | deny;
+		(binary(), binary(), presence_permissions(), presence_permission()) -> allow | deny;
+		(binary(), binary(), message_permissions(), message_permission()) -> allow | deny.
 match_rule(ServerHost, Host, Perms, Type) ->
    Access = proplists:get_value(Type, Perms, none),
    acl:match_rule(ServerHost, Access, jid:make(Host)).