Limit result set of disco#items for mod_pubsub

The size of a list of nodes returned for disco#items request
is now controlled by option 'max_nodes_discoitems'. The default
value is 100. The name and the default value of the option is
chosen to be consistent with mod_muc's 'max_rooms_discoitems' option.

6 files changed, 101 insertions, 49 deletions

diff --git a/src/gen_pubsub_nodetree.erl b/src/gen_pubsub_nodetree.erl
index 813ed71ce..e834f8e0b 100644
--- a/src/gen_pubsub_nodetree.erl
+++ b/src/gen_pubsub_nodetree.erl
@@ -61,7 +61,7 @@
     {error, stanza_error()}.
 
 -callback get_nodes(Host :: host(),
-	From :: jid:jid())->
+		    Limit :: non_neg_integer() | infinity)->
     [pubsubNode()].
 
 -callback get_nodes(Host :: host())->
@@ -79,8 +79,8 @@
     [{0, [pubsubNode(),...]}].
 
 -callback get_subnodes(Host :: host(),
-	NodeId :: nodeId(),
-	From :: jid:jid()) ->
+		       NodeId :: nodeId(),
+		       Limit :: non_neg_integer() | infinity) ->
     [pubsubNode()].
 
 -callback get_subnodes_tree(Host :: host(),
diff --git a/src/mod_pubsub.erl b/src/mod_pubsub.erl
index 0063df8d9..48a19b99f 100644
--- a/src/mod_pubsub.erl
+++ b/src/mod_pubsub.erl
@@ -498,6 +498,7 @@ disco_sm_items(Acc, _From, _To, _Node, _Lang) -> Acc.
 
 -spec disco_items(ljid(), binary(), jid()) -> [disco_item()].
 disco_items(Host, <<>>, From) ->
+    MaxNodes = mod_pubsub_opt:max_nodes_discoitems(serverhost(Host)),
     Action =
 	fun(#pubsub_node{nodeid = {_, Node}, options = Options,
 			 type = Type, id = Nidx, owners = O}, Acc) ->
@@ -513,7 +514,7 @@ disco_items(Host, <<>>, From) ->
 		end
 	end,
     NodeBloc = fun() ->
-		       case tree_call(Host, get_nodes, [Host]) of
+		       case tree_call(Host, get_nodes, [Host, MaxNodes]) of
 			   Nodes when is_list(Nodes) ->
 			       {result, lists:foldl(Action, [], Nodes)};
 			   Error ->
@@ -1007,8 +1008,9 @@ iq_disco_info(ServerHost, Host, SNode, From, Lang) ->
 
 -spec iq_disco_items(host(), binary(), jid(), undefined | rsm_set()) ->
 			    {result, disco_items()} | {error, stanza_error()}.
-iq_disco_items(Host, <<>>, From, _RSM) ->
-    case tree_action(Host, get_subnodes, [Host, <<>>, From]) of
+iq_disco_items(Host, <<>>, _From, _RSM) ->
+    MaxNodes = mod_pubsub_opt:max_nodes_discoitems(serverhost(Host)),
+    case tree_action(Host, get_subnodes, [Host, <<>>, MaxNodes]) of
 	{error, #stanza_error{}} = Err ->
 	    Err;
 	Nodes when is_list(Nodes) ->
@@ -1039,6 +1041,7 @@ iq_disco_items(Host, Item, From, RSM) ->
 	[_Node, _ItemId] ->
 	    {result, #disco_items{}};
 	[Node] ->
+	    MaxNodes = mod_pubsub_opt:max_nodes_discoitems(serverhost(Host)),
 	    Action = fun(#pubsub_node{id = Nidx, type = Type, options = Options, owners = O}) ->
 			     Owners = node_owners_call(Host, Type, Nidx, O),
 			     {NodeItems, RsmOut} = case get_allowed_items_call(
@@ -1046,7 +1049,7 @@ iq_disco_items(Host, Item, From, RSM) ->
 						       {result, R} -> R;
 						       _ -> {[], undefined}
 						   end,
-			     case tree_call(Host, get_subnodes, [Host, Node, From]) of
+			     case tree_call(Host, get_subnodes, [Host, Node, MaxNodes]) of
 				 SubNodes when is_list(SubNodes) ->
 				     Nodes = lists:map(
 					       fun(#pubsub_node{nodeid = {_, SubNode}, options = SubOptions}) ->
@@ -3154,7 +3157,7 @@ send_last_pep(From, To) ->
     Host = host(ServerHost),
     Publisher = jid:tolower(From),
     Owner = jid:remove_resource(Publisher),
-    case tree_action(Host, get_nodes, [Owner, From]) of
+    case tree_action(Host, get_nodes, [Owner, infinity]) of
 	Nodes when is_list(Nodes) ->
 	    lists:foreach(
 	      fun(#pubsub_node{nodeid = {_, Node}, type = Type, id = Nidx, options = Options}) ->
@@ -4123,6 +4126,8 @@ mod_opt_type(last_item_cache) ->
     econf:bool();
 mod_opt_type(max_items_node) ->
     econf:non_neg_int();
+mod_opt_type(max_nodes_discoitems) ->
+    econf:non_neg_int(infinity);
 mod_opt_type(max_subscriptions_node) ->
     econf:non_neg_int();
 mod_opt_type(force_node_config) ->
@@ -4168,6 +4173,7 @@ mod_options(Host) ->
      {ignore_pep_from_offline, true},
      {last_item_cache, false},
      {max_items_node, ?MAXITEMS},
+     {max_nodes_discoitems, 100},
      {nodetree, ?STDTREE},
      {pep_mapping, []},
      {plugins, [?STDNODE]},
diff --git a/src/mod_pubsub_opt.erl b/src/mod_pubsub_opt.erl
index 68ed7f960..8db5532f6 100644
--- a/src/mod_pubsub_opt.erl
+++ b/src/mod_pubsub_opt.erl
@@ -12,6 +12,7 @@
 -export([ignore_pep_from_offline/1]).
 -export([last_item_cache/1]).
 -export([max_items_node/1]).
+-export([max_nodes_discoitems/1]).
 -export([max_subscriptions_node/1]).
 -export([name/1]).
 -export([nodetree/1]).
@@ -73,6 +74,12 @@ max_items_node(Opts) when is_map(Opts) ->
 max_items_node(Host) ->
     gen_mod:get_module_opt(Host, mod_pubsub, max_items_node).
 
+-spec max_nodes_discoitems(gen_mod:opts() | global | binary()) -> 'infinity' | non_neg_integer().
+max_nodes_discoitems(Opts) when is_map(Opts) ->
+    gen_mod:get_opt(max_nodes_discoitems, Opts);
+max_nodes_discoitems(Host) ->
+    gen_mod:get_module_opt(Host, mod_pubsub, max_nodes_discoitems).
+
 -spec max_subscriptions_node(gen_mod:opts() | global | binary()) -> 'undefined' | non_neg_integer().
 max_subscriptions_node(Opts) when is_map(Opts) ->
     gen_mod:get_opt(max_subscriptions_node, Opts);
diff --git a/src/nodetree_tree.erl b/src/nodetree_tree.erl
index 08bc3192c..c94ba197b 100644
--- a/src/nodetree_tree.erl
+++ b/src/nodetree_tree.erl
@@ -38,6 +38,7 @@
 -author('christophe.romain@process-one.net').
 
 -include_lib("stdlib/include/qlc.hrl").
+-include_lib("stdlib/include/ms_transform.hrl").
 
 -include("pubsub.hrl").
 -include("xmpp.hrl").
@@ -81,11 +82,21 @@ get_node(Nidx) ->
 	_ -> {error, xmpp:err_item_not_found(?T("Node not found"), ejabberd_option:language())}
     end.
 
-get_nodes(Host, _From) ->
-    get_nodes(Host).
-
 get_nodes(Host) ->
-    mnesia:match_object(#pubsub_node{nodeid = {Host, '_'}, _ = '_'}).
+    get_nodes(Host, infinity).
+
+get_nodes(Host, infinity) ->
+    mnesia:match_object(#pubsub_node{nodeid = {Host, '_'}, _ = '_'});
+get_nodes(Host, Limit) ->
+    case mnesia:select(
+	   pubsub_node,
+	   ets:fun2ms(
+	     fun(#pubsub_node{nodeid = {H, _}} = Node) when H == Host ->
+		     Node
+	     end), Limit, read) of
+	'$end_of_table' -> [];
+	{Nodes, _} -> Nodes
+    end.
 
 get_parentnodes(Host, Node, _From) ->
     case catch mnesia:read({pubsub_node, {Host, Node}}) of
@@ -109,25 +120,40 @@ get_parentnodes_tree(Host, Node, Level, Acc) ->
 	    Acc
     end.
 
-get_subnodes(Host, Node, _From) ->
-    get_subnodes(Host, Node).
-
-get_subnodes(Host, <<>>) ->
-    Q = qlc:q([N
-		|| #pubsub_node{nodeid = {NHost, _},
-			parents = Parents} =
-		    N
-		    <- mnesia:table(pubsub_node),
-		    Host == NHost, Parents == []]),
-    qlc:e(Q);
-get_subnodes(Host, Node) ->
+get_subnodes(Host, <<>>, infinity) ->
+    mnesia:match_object(#pubsub_node{nodeid = {Host, '_'}, parents = [], _ = '_'});
+get_subnodes(Host, <<>>, Limit) ->
+    case mnesia:select(
+	   pubsub_node,
+	   ets:fun2ms(
+	     fun(#pubsub_node{nodeid = {H, _}, parents = []} = Node) when H == Host ->
+		     Node
+	     end), Limit, read) of
+	'$end_of_table' -> [];
+	{Nodes, _} -> Nodes
+    end;
+get_subnodes(Host, Node, infinity) ->
     Q = qlc:q([N
 		|| #pubsub_node{nodeid = {NHost, _},
 			parents = Parents} =
 		    N
 		    <- mnesia:table(pubsub_node),
 		    Host == NHost, lists:member(Node, Parents)]),
-    qlc:e(Q).
+    qlc:e(Q);
+get_subnodes(Host, Node, Limit) ->
+    case mnesia:select(
+	   pubsub_node,
+	   ets:fun2ms(
+	     fun(#pubsub_node{nodeid = {H, _}, parents = Ps} = N)
+		   when H == Host andalso Ps /= [] -> N
+	     end), Limit, read) of
+	'$end_of_table' -> [];
+	{Nodes, _} ->
+	    lists:filter(
+	      fun(#pubsub_node{parents = Parents}) ->
+		      lists:member(Node, Parents)
+	      end, Nodes)
+    end.
 
 get_subnodes_tree(Host, Node, _From) ->
     get_subnodes_tree(Host, Node).
diff --git a/src/nodetree_tree_sql.erl b/src/nodetree_tree_sql.erl
index efef656c5..bb9a27a61 100644
--- a/src/nodetree_tree_sql.erl
+++ b/src/nodetree_tree_sql.erl
@@ -140,16 +140,25 @@ get_node(Nidx) ->
 	    {error, xmpp:err_item_not_found(?T("Node not found"), ejabberd_option:language())}
     end.
 
-get_nodes(Host, _From) ->
-    get_nodes(Host).
-
 get_nodes(Host) ->
+    get_nodes(Host, infinity).
+
+get_nodes(Host, Limit) ->
     H = node_flat_sql:encode_host(Host),
-    case catch
-	ejabberd_sql:sql_query_t(
-	  ?SQL("select @(node)s, @(parent)s, @(plugin)s, @(nodeid)d from pubsub_node "
-	       "where host=%(H)s"))
-    of
+    Query = fun(mssql, _) when is_integer(Limit), Limit>=0 ->
+		    ejabberd_sql:sql_query_t(
+		      ?SQL("select top %(Limit)d @(node)s, @(parent)s, @(plugin)s, @(nodeid)d "
+			   "from pubsub_node where host=%(H)s"));
+	       (_, _) when is_integer(Limit), Limit>=0 ->
+		    ejabberd_sql:sql_query_t(
+		      ?SQL("select @(node)s, @(parent)s, @(plugin)s, @(nodeid)d "
+			   "from pubsub_node where host=%(H)s limit %(Limit)d"));
+	       (_, _) ->
+		    ejabberd_sql:sql_query_t(
+		      ?SQL("select @(node)s, @(parent)s, @(plugin)s, @(nodeid)d "
+			   "from pubsub_node where host=%(H)s"))
+	    end,
+    case ejabberd_sql:sql_query_t(Query) of
 	{selected, RItems} ->
 	    [raw_to_node(Host, Item) || Item <- RItems];
 	_ ->
@@ -178,16 +187,23 @@ get_parentnodes_tree(Host, Node, Level, Acc) ->
 	    Acc
     end.
 
-get_subnodes(Host, Node, _From) ->
-    get_subnodes(Host, Node).
-
-get_subnodes(Host, Node) ->
+get_subnodes(Host, Node, Limit) ->
     H = node_flat_sql:encode_host(Host),
-    case catch
-	ejabberd_sql:sql_query_t(
-	  ?SQL("select @(node)s, @(parent)s, @(plugin)s, @(nodeid)d from pubsub_node "
-	       "where host=%(H)s and parent=%(Node)s"))
-    of
+    Query = fun(mssql, _) when is_integer(Limit), Limit>=0 ->
+		    ejabberd_sql:sql_query_t(
+		      ?SQL("select top %(Limit)d @(node)s, @(parent)s, @(plugin)s, @(nodeid)d "
+			   "from pubsub_node where host=%(H)s and parent=%(Node)s"));
+	       (_, _) when is_integer(Limit), Limit>=0 ->
+		    ejabberd_sql:sql_query_t(
+		      ?SQL("select @(node)s, @(parent)s, @(plugin)s, @(nodeid)d "
+			   "from pubsub_node where host=%(H)s and parent=%(Node)s "
+			   "limit %(Limit)d"));
+	       (_, _) ->
+		    ejabberd_sql:sql_query_t(
+		      ?SQL("select @(node)s, @(parent)s, @(plugin)s, @(nodeid)d "
+			   "from pubsub_node where host=%(H)s and parent=%(Node)s"))
+	    end,
+    case ejabberd_sql:sql_query_t(Query) of
 	{selected, RItems} ->
 	    [raw_to_node(Host, Item) || Item <- RItems];
 	_ ->
diff --git a/src/nodetree_virtual.erl b/src/nodetree_virtual.erl
index c27efe44b..627eca092 100644
--- a/src/nodetree_virtual.erl
+++ b/src/nodetree_virtual.erl
@@ -65,10 +65,10 @@ get_node(Nidx) ->
     {Host, Node} = nodeid(Nidx),
     node_record(Host, Node, Nidx).
 
-get_nodes(Host, _From) ->
-    get_nodes(Host).
+get_nodes(Host) ->
+    get_nodes(Host, infinity).
 
-get_nodes(_Host) ->
+get_nodes(_Host, _Limit) ->
     [].
 
 get_parentnodes(_Host, _Node, _From) ->
@@ -77,10 +77,7 @@ get_parentnodes(_Host, _Node, _From) ->
 get_parentnodes_tree(Host, Node, From) ->
     [{0, [get_node(Host, Node, From)]}].
 
-get_subnodes(Host, Node, _From) ->
-    get_subnodes(Host, Node).
-
-get_subnodes(_Host, _Node) ->
+get_subnodes(_Host, _Node, _From) ->
     [].
 
 get_subnodes_tree(Host, Node, _From) ->