From 3f7a850ae8b5ae752059c95d57241ff4345542c4 Mon Sep 17 00:00:00 2001
From: Evgeny Khramtsov <ekhramtsov@process-one.net>
Date: Thu, 23 May 2019 11:32:55 +0300
Subject: [PATCH] Use different cache tables per auth module

Fixes #2322
---
 src/ejabberd_auth.erl | 61 +++++++++++++++++++++++++++----------------
 1 file changed, 38 insertions(+), 23 deletions(-)

diff --git a/src/ejabberd_auth.erl b/src/ejabberd_auth.erl
index 7d6591135..a3256364e 100644
--- a/src/ejabberd_auth.erl
+++ b/src/ejabberd_auth.erl
@@ -52,7 +52,6 @@
 -include("scram.hrl").
 -include("logger.hrl").
 
--define(AUTH_CACHE, auth_cache).
 -define(SALT_LENGTH, 16).
 
 -record(state, {host_modules = #{} :: map()}).
@@ -546,7 +545,7 @@ db_try_register(User, Server, Password, Mod) ->
 	    case use_cache(Mod, Server) of
 		true ->
 		    case ets_cache:update(
-			   ?AUTH_CACHE, {User, Server}, {ok, Password},
+			   cache_tab(Mod), {User, Server}, {ok, Password},
 			   fun() -> Mod:try_register(User, Server, Password1) end,
 			   cache_nodes(Mod, Server)) of
 			{ok, _} -> ok;
@@ -569,7 +568,7 @@ db_set_password(User, Server, Password, Mod) ->
 	    case use_cache(Mod, Server) of
 		true ->
 		    case ets_cache:update(
-			   ?AUTH_CACHE, {User, Server}, {ok, Password},
+			   cache_tab(Mod), {User, Server}, {ok, Password},
 			   fun() -> Mod:set_password(User, Server, Password1) end,
 			   cache_nodes(Mod, Server)) of
 			{ok, _} -> ok;
@@ -586,7 +585,7 @@ db_get_password(User, Server, Mod) ->
     UseCache = use_cache(Mod, Server),
     case erlang:function_exported(Mod, get_password, 2) of
 	false when UseCache ->
-	    case ets_cache:lookup(?AUTH_CACHE, {User, Server}) of
+	    case ets_cache:lookup(cache_tab(Mod), {User, Server}) of
 		{ok, exists} -> error;
 		Other -> Other
 	    end;
@@ -594,7 +593,7 @@ db_get_password(User, Server, Mod) ->
 	    error;
 	true when UseCache ->
 	    ets_cache:lookup(
-	      ?AUTH_CACHE, {User, Server},
+	      cache_tab(Mod), {User, Server},
 	      fun() -> Mod:get_password(User, Server) end);
 	true ->
 	    ets_cache:untag(Mod:get_password(User, Server))
@@ -608,7 +607,7 @@ db_user_exists(User, Server, Mod) ->
 	    case {Mod:store_type(Server), use_cache(Mod, Server)} of
 		{external, true} ->
 		    case ets_cache:lookup(
-			   ?AUTH_CACHE, {User, Server},
+			   cache_tab(Mod), {User, Server},
 			   fun() ->
 				   case Mod:user_exists(User, Server) of
 				       true -> {ok, exists};
@@ -642,7 +641,7 @@ db_check_password(User, AuthzId, Server, ProvidedPassword,
 	    case {Mod:store_type(Server), use_cache(Mod, Server)} of
 		{external, true} ->
 		    case ets_cache:update(
-			   ?AUTH_CACHE, {User, Server}, {ok, ProvidedPassword},
+			   cache_tab(Mod), {User, Server}, {ok, ProvidedPassword},
 			   fun() ->
 				   case Mod:check_password(
 					  User, AuthzId, Server, ProvidedPassword) of
@@ -672,7 +671,7 @@ db_remove_user(User, Server, Mod) ->
 		ok ->
 		    case use_cache(Mod, Server) of
 			true ->
-			    ets_cache:delete(?AUTH_CACHE, {User, Server},
+			    ets_cache:delete(cache_tab(Mod), {User, Server},
 					     cache_nodes(Mod, Server));
 			false ->
 			    ok
@@ -696,7 +695,7 @@ db_get_users(Server, Opts, Mod) ->
 			      [{User, Server}|Users];
 			 (_, _, Users) ->
 			      Users
-		      end, [], ?AUTH_CACHE);
+		      end, [], cache_tab(Mod));
 		false ->
 		    []
 	    end
@@ -714,7 +713,7 @@ db_count_users(Server, Opts, Mod) ->
 			      Num + 1;
 			 (_, _, Num) ->
 			      Num
-		      end, 0, ?AUTH_CACHE);
+		      end, 0, cache_tab(Mod));
 		false ->
 		    0
 	    end
@@ -755,12 +754,16 @@ password_to_scram(Password, IterationCount) ->
 %%%----------------------------------------------------------------------
 -spec init_cache(map()) -> ok.
 init_cache(HostModules) ->
-    case use_cache(HostModules) of
-	true ->
-	    ets_cache:new(?AUTH_CACHE, cache_opts());
-	false ->
-	    ets_cache:delete(?AUTH_CACHE)
-    end.
+    CacheOpts = cache_opts(),
+    {True, False} = use_cache(HostModules),
+    lists:foreach(
+      fun(Module) ->
+	      ets_cache:new(cache_tab(Module), CacheOpts)
+      end, True),
+    lists:foreach(
+      fun(Module) ->
+	      ets_cache:delete(cache_tab(Module))
+      end, False).
 
 -spec cache_opts() -> [proplists:property()].
 cache_opts() ->
@@ -778,14 +781,22 @@ cache_opts() ->
 	       end,
     [{max_size, MaxSize}, {cache_missed, CacheMissed}, {life_time, LifeTime}].
 
--spec use_cache(map()) -> boolean().
+-spec use_cache(map()) -> {True :: [module()], False :: [module()]}.
 use_cache(HostModules) ->
-    lists:any(
-      fun({Host, Modules}) ->
-	      lists:any(fun(Module) ->
-				use_cache(Module, Host)
-			end, Modules)
-      end, maps:to_list(HostModules)).
+    {Enabled, Disabled} =
+	maps:fold(
+	  fun(Host, Modules, Acc) ->
+		  lists:foldl(
+		    fun(Module, {True, False}) ->
+			    case use_cache(Module, Host) of
+				true ->
+				    {sets:add_element(Module, True), False};
+				false ->
+				    {True, sets:add_element(Module, False)}
+			    end
+		    end, Acc, Modules)
+	  end, {sets:new(), sets:new()}, HostModules),
+    {sets:to_list(Enabled), sets:to_list(sets:subtract(Disabled, Enabled))}.
 
 -spec use_cache(module(), binary()) -> boolean().
 use_cache(Mod, LServer) ->
@@ -804,6 +815,10 @@ cache_nodes(Mod, LServer) ->
 	false -> ejabberd_cluster:get_nodes()
     end.
 
+-spec cache_tab(module()) -> atom().
+cache_tab(Mod) ->
+    list_to_atom(atom_to_list(Mod) ++ "_cache").
+
 %%%----------------------------------------------------------------------
 %%% Internal functions
 %%%----------------------------------------------------------------------
-- 
2.40.0