From bd535f95237368be6bac4174c59e8add25ad71db Mon Sep 17 00:00:00 2001
From: Alekzander Spiridonov <sipidronov@gmail.com>
Date: Thu, 21 Aug 2014 08:20:09 -0700
Subject: [PATCH] dispatcher: option to limit the number of items in dst_avp list

---
 modules/dispatcher/api.h        |    2 +-
 modules/dispatcher/dispatch.c   |   15 +++++++---
 modules/dispatcher/dispatch.h   |    2 +-
 modules/dispatcher/dispatcher.c |   56 ++++++++++++++++++++++++++++++++-------
 4 files changed, 59 insertions(+), 16 deletions(-)

diff --git a/modules/dispatcher/api.h b/modules/dispatcher/api.h
index df4d765..0ff4585 100644
--- a/modules/dispatcher/api.h
+++ b/modules/dispatcher/api.h
@@ -34,7 +34,7 @@
 #include "../../sr_module.h"
 
 typedef int (*ds_select_dst_f)(struct sip_msg *msg, int set,
-			int alg, int mode);
+			int alg, unsigned int limit, int mode);
 typedef int (*ds_next_dst_f)(struct sip_msg *msg, int mode);
 typedef int (*ds_mark_dst_f)(struct sip_msg *msg, int mode);
 
diff --git a/modules/dispatcher/dispatch.c b/modules/dispatcher/dispatch.c
index bc945f8..f76c490 100644
--- a/modules/dispatcher/dispatch.c
+++ b/modules/dispatcher/dispatch.c
@@ -1581,7 +1581,7 @@ static inline int ds_update_dst(struct sip_msg *msg, str *uri, int mode)
 /**
  *
  */
-int ds_select_dst(struct sip_msg *msg, int set, int alg, int mode)
+int ds_select_dst(struct sip_msg *msg, int set, int alg, unsigned int limit, int mode)
 {
 	int i, cnt;
 	unsigned int hash;
@@ -1600,6 +1600,13 @@ int ds_select_dst(struct sip_msg *msg, int set, int alg, int mode)
 		return -1;
 	}
 
+	if (limit==0)
+	{
+	  LM_DBG("Limit set to 0 - forcing to unlimited\n");
+	  limit = -1;
+	}
+  --limit; /* reserving 1 slot for selected dst */
+
 	if((mode==0) && (ds_force_dst==0)
 			&& (msg->dst_uri.s!=NULL || msg->dst_uri.len>0))
 	{
@@ -1767,7 +1774,7 @@ int ds_select_dst(struct sip_msg *msg, int set, int alg, int mode)
 	if(dst_avp_name.n!=0)
 	{
 		/* add default dst to last position in AVP list */
-		if(ds_use_default!=0 && hash!=idx->nr-1)
+		if(ds_use_default!=0 && hash!=idx->nr-1 && cnt<limit)
 		{
 			avp_val.s = idx->dlist[idx->nr-1].uri;
 			if(add_avp(AVP_VAL_STR|dst_avp_type, dst_avp_name, avp_val)!=0)
@@ -1799,7 +1806,7 @@ int ds_select_dst(struct sip_msg *msg, int set, int alg, int mode)
 
 		/* add to avp */
 
-		for(i=hash-1; i>=0; i--)
+		for(i=hash-1; i>=0 && cnt<limit; i--)
 		{	
 			if(ds_skip_dst(idx->dlist[i].flags)
 					|| (ds_use_default!=0 && i==(idx->nr-1)))
@@ -1833,7 +1840,7 @@ int ds_select_dst(struct sip_msg *msg, int set, int alg, int mode)
 			cnt++;
 		}
 
-		for(i=idx->nr-1; i>hash; i--)
+		for(i=idx->nr-1; i>hash && cnt<limit; i--)
 		{	
 			if(ds_skip_dst(idx->dlist[i].flags)
 					|| (ds_use_default!=0 && i==(idx->nr-1)))
diff --git a/modules/dispatcher/dispatch.h b/modules/dispatcher/dispatch.h
index 31262b5..4384a46 100644
--- a/modules/dispatcher/dispatch.h
+++ b/modules/dispatcher/dispatch.h
@@ -107,7 +107,7 @@ void ds_disconnect_db(void);
 int ds_load_db(void);
 int ds_reload_db(void);
 int ds_destroy_list(void);
-int ds_select_dst(struct sip_msg *msg, int set, int alg, int mode);
+int ds_select_dst(struct sip_msg *msg, int set, int alg, unsigned int limit, int mode);
 int ds_next_dst(struct sip_msg *msg, int mode);
 int ds_update_state(sip_msg_t *msg, int group, str *address, int state);
 int ds_reinit_state(int group, str *address, int state);
diff --git a/modules/dispatcher/dispatcher.c b/modules/dispatcher/dispatcher.c
index 7df0c6d..7277ba1 100644
--- a/modules/dispatcher/dispatcher.c
+++ b/modules/dispatcher/dispatcher.c
@@ -106,8 +106,8 @@ pv_elem_t * hash_param_model = NULL;
 
 int probing_threshhold = 1; /* number of failed requests, before a destination
 							   is taken into probing */
-str ds_ping_method = {"OPTIONS",7};
-str ds_ping_from   = {"sip:dispatcher@localhost", 24};
+str ds_ping_method = str_init("OPTIONS");
+str ds_ping_from   = str_init("sip:dispatcher@localhost");
 static int ds_ping_interval = 0;
 int ds_probing_mode  = DS_PROBE_NONE;
 
@@ -147,7 +147,9 @@ static int ds_parse_reply_codes();
 static int ds_init_rpc(void);
 
 static int w_ds_select_dst(struct sip_msg*, char*, char*);
+static int w_ds_select_dst_limit(struct sip_msg*, char*, char*, char*);
 static int w_ds_select_domain(struct sip_msg*, char*, char*);
+static int w_ds_select_domain_limit(struct sip_msg*, char*, char*, char*);
 static int w_ds_next_dst(struct sip_msg*, char*, char*);
 static int w_ds_next_domain(struct sip_msg*, char*, char*);
 static int w_ds_mark_dst0(struct sip_msg*, char*, char*);
@@ -170,8 +172,12 @@ static int mi_child_init(void);
 static cmd_export_t cmds[]={
 	{"ds_select_dst",    (cmd_function)w_ds_select_dst,    2,
 		fixup_igp_igp, 0, REQUEST_ROUTE|FAILURE_ROUTE},
+  {"ds_select_dst",    (cmd_function)w_ds_select_dst_limit,    3,
+    fixup_igp_null, 0, REQUEST_ROUTE|FAILURE_ROUTE},
 	{"ds_select_domain", (cmd_function)w_ds_select_domain, 2,
 		fixup_igp_igp, 0, REQUEST_ROUTE|FAILURE_ROUTE},
+  {"ds_select_domain", (cmd_function)w_ds_select_domain_limit, 3,
+      fixup_igp_null, 0, REQUEST_ROUTE|FAILURE_ROUTE},
 	{"ds_next_dst",      (cmd_function)w_ds_next_dst,      0,
 		ds_warn_fixup, 0, REQUEST_ROUTE|FAILURE_ROUTE},
 	{"ds_next_domain",   (cmd_function)w_ds_next_domain,   0,
@@ -526,12 +532,13 @@ static void destroy(void)
 /**
  *
  */
-static int w_ds_select(struct sip_msg* msg, char* set, char* alg, int mode)
+static int w_ds_select(struct sip_msg* msg, char* set, char* alg, char* limit, int mode)
 {
-	unsigned int algo_flags, set_flags;
-	str s_algo = {NULL, 0};
-	str s_set = {NULL, 0};
-	int a, s;
+	unsigned int algo_flags, set_flags, limit_flags;
+	str s_algo = STR_NULL;
+	str s_set = STR_NULL;
+	str s_limit = STR_NULL;
+	int a, s, l;
 	if(msg==NULL)
 		return -1;
 
@@ -552,14 +559,35 @@ static int w_ds_select(struct sip_msg* msg, char* set, char* alg, int mode)
 		return -1;
 	}
 
-	return ds_select_dst(msg, s, a, mode);
+  if (limit) {
+    GET_VALUE("limit", limit, l, s_limit, limit_flags);
+    if (!(limit_flags&PARAM_INT)) {
+      if (limit_flags&PARAM_STR)
+        LM_ERR("unable to get dst number limit from [%.*s]\n", s_limit.len, s_limit.s);
+      else
+        LM_ERR("unable to get dst number limit\n");
+      return -1;
+    }
+  } else {
+    l = -1; /* will be casted to a rather big unsigned value */
+  }
+
+	return ds_select_dst(msg, s, a, (unsigned int)l, mode);
 }
 /**
  *
  */
 static int w_ds_select_dst(struct sip_msg* msg, char* set, char* alg)
 {
-	return w_ds_select(msg, set, alg, 0 /*set dst uri*/);
+	return w_ds_select(msg, set, alg, 0 /* limit number of dst*/, 0 /*set dst uri*/);
+}
+
+/**
+ *
+ */
+static int w_ds_select_dst_limit(struct sip_msg* msg, char* set, char* alg, char* limit)
+{
+  return w_ds_select(msg, set, alg, limit /* limit number of dst*/, 0 /*set dst uri*/);
 }
 
 /**
@@ -567,7 +595,15 @@ static int w_ds_select_dst(struct sip_msg* msg, char* set, char* alg)
  */
 static int w_ds_select_domain(struct sip_msg* msg, char* set, char* alg)
 {
-	return w_ds_select(msg, set, alg, 1 /*set host port*/);
+	return w_ds_select(msg, set, alg, 0 /* limit number of dst*/, 1 /*set host port*/);
+}
+
+/**
+ *
+ */
+static int w_ds_select_domain_limit(struct sip_msg* msg, char* set, char* alg, char* limit)
+{
+  return w_ds_select(msg, set, alg, limit /* limit number of dst*/, 1 /*set host port*/);
 }
 
 /**
-- 
1.7.1

