Rework IP addresses in upstreams.

- Select ipv4/unix addresses if they exist and use ipv6 for ipv6 only
  upstreams (since the support of ipv6 is poor in many OSes and
  environments).
- Free IP list on upstream destruction.
- Add test cases for addresses selection.
- Allow adding of free form IP addresses to upstreams.

3 files changed, 117 insertions, 1 deletions

diff --git a/src/libutil/upstream.c b/src/libutil/upstream.c
index 7bec3729a..b5f1ed35e 100644
--- a/src/libutil/upstream.c
+++ b/src/libutil/upstream.c
@@ -105,6 +105,42 @@ rspamd_upstreams_library_init (struct rdns_resolver *resolver,
 	ev_base = base;
 }
 
+static gint
+rspamd_upstream_af_to_weight (const rspamd_inet_addr_t *addr)
+{
+	int ret;
+
+	switch (addr->af) {
+	case AF_UNIX:
+		ret = 2;
+		break;
+	case AF_INET:
+		ret = 1;
+		break;
+	default:
+		ret = 0;
+		break;
+	}
+
+	return ret;
+}
+
+/*
+ * Select IPv4 addresses before IPv6
+ */
+static gint
+rspamd_upstream_addr_sort_func (gconstpointer a, gconstpointer b)
+{
+	const rspamd_inet_addr_t *ip1 = (const rspamd_inet_addr_t *)a,
+			*ip2 = (const rspamd_inet_addr_t *)b;
+	gint w1, w2;
+
+	w1 = rspamd_upstream_af_to_weight (ip1);
+	w2 = rspamd_upstream_af_to_weight (ip2);
+
+	return w2 - w1;
+}
+
 static void
 rspamd_upstream_set_active (struct upstream_list *ls, struct upstream *up)
 {
@@ -186,6 +222,8 @@ rspamd_upstream_update_addrs (struct upstream *up)
 		up->addrs.cur = 0;
 		up->addrs.count = addr_cnt;
 		up->addrs.addr = new_addrs;
+		qsort (up->addrs.addr, up->addrs.count, sizeof (up->addrs.addr[0]),
+					rspamd_upstream_addr_sort_func);
 		g_free (old);
 	}
 
@@ -340,13 +378,36 @@ rspamd_upstream_dtor (struct upstream *up)
 
 	rspamd_mutex_free (up->lock);
 	g_free (up->name);
+	g_free (up->addrs.addr);
 	g_slice_free1 (sizeof (*up), up);
 }
 
 rspamd_inet_addr_t*
 rspamd_upstream_addr (struct upstream *up)
 {
-	return &up->addrs.addr[up->addrs.cur++ % up->addrs.count];
+	gint idx, next_idx, w1, w2;
+	/*
+	 * We know that addresses are sorted in the way that ipv4 addresses come
+	 * first. Therefore, we select only ipv4 addresses if they exist, since
+	 * many systems now has poorly supported ipv6
+	 */
+	idx = up->addrs.cur;
+	next_idx = (idx + 1) % up->addrs.count;
+	w1 = rspamd_upstream_af_to_weight (&up->addrs.addr[idx]);
+	w2 = rspamd_upstream_af_to_weight (&up->addrs.addr[next_idx]);
+
+	/*
+	 * We don't care about the exact priorities, but we prefer ipv4/unix
+	 * addresses before any ipv6 addresses
+	 */
+	if (!w1 || w2) {
+		up->addrs.cur = next_idx;
+	}
+	else {
+		up->addrs.cur = 0;
+	}
+
+	return &up->addrs.addr[up->addrs.cur];
 }
 
 const gchar*
@@ -377,6 +438,8 @@ rspamd_upstreams_add_upstream (struct upstream_list *ups,
 	up->ls = ups;
 	REF_INIT_RETAIN (up, rspamd_upstream_dtor);
 	up->lock = rspamd_mutex_new ();
+	qsort (up->addrs.addr, up->addrs.count, sizeof (up->addrs.addr[0]),
+			rspamd_upstream_addr_sort_func);
 
 	rspamd_upstream_set_active (ups, up);
 
@@ -384,6 +447,24 @@ rspamd_upstreams_add_upstream (struct upstream_list *ups,
 }
 
 gboolean
+rspamd_upstream_add_addr (struct upstream *up, const rspamd_inet_addr_t *addr)
+{
+	gint nsz;
+
+	/*
+	 * XXX: slow and inefficient
+	 */
+	nsz = ++up->addrs.count;
+	up->addrs.addr = g_realloc (up->addrs.addr,
+			nsz * sizeof (up->addrs.addr[0]));
+	memcpy (&up->addrs.addr[nsz - 1], addr, sizeof (*addr));
+	qsort (up->addrs.addr, up->addrs.count, sizeof (up->addrs.addr[0]),
+				rspamd_upstream_addr_sort_func);
+
+	return TRUE;
+}
+
+gboolean
 rspamd_upstreams_parse_line (struct upstream_list *ups,
 		const gchar *str, guint16 def_port, void *data)
 {
diff --git a/src/libutil/upstream.h b/src/libutil/upstream.h
index 812cc4bd8..cff6acfb1 100644
--- a/src/libutil/upstream.h
+++ b/src/libutil/upstream.h
@@ -117,6 +117,14 @@ gboolean rspamd_upstreams_from_ucl (struct upstream_list *ups,
 rspamd_inet_addr_t* rspamd_upstream_addr (struct upstream *up);
 
 /**
+ * Add custom address for an upstream
+ * @param up
+ * @return
+ */
+gboolean rspamd_upstream_add_addr (struct upstream *up,
+		const rspamd_inet_addr_t *addr);
+
+/**
  * Returns the symbolic name of the upstream
  * @param up
  * @return
diff --git a/test/rspamd_upstream_test.c b/test/rspamd_upstream_test.c
index 80d963963..23a450ec4 100644
--- a/test/rspamd_upstream_test.c
+++ b/test/rspamd_upstream_test.c
@@ -72,6 +72,7 @@ rspamd_upstream_test_func (void)
 	gdouble p;
 	struct event ev;
 	struct timeval tv;
+	rspamd_inet_addr_t *addr, *next_addr, paddr;
 
 	cfg = (struct rspamd_config *)g_malloc (sizeof (struct rspamd_config));
 	bzero (cfg, sizeof (struct rspamd_config));
@@ -129,6 +130,29 @@ rspamd_upstream_test_func (void)
 	msg_debug ("p value for hash consistency: %.6f", p);
 	g_assert (p > 0.9);
 
+	rspamd_upstreams_destroy (nls);
+
+	/*
+	 * Test v4/v6 priorities
+	 */
+	nls = rspamd_upstreams_create ();
+	g_assert (rspamd_upstreams_add_upstream (nls, "127.0.0.1", 0, NULL));
+	up = rspamd_upstream_get (nls, RSPAMD_UPSTREAM_RANDOM);
+	addr = g_malloc (sizeof (*addr));
+	rspamd_parse_inet_address(&paddr, "127.0.0.2");
+	g_assert (rspamd_upstream_add_addr (up, &paddr));
+	rspamd_parse_inet_address(&paddr, "::1");
+	g_assert (rspamd_upstream_add_addr (up, &paddr));
+	addr = rspamd_upstream_addr (up);
+	for (i = 0; i < 256; i ++) {
+		next_addr = rspamd_upstream_addr (up);
+		g_assert (addr->af == AF_INET);
+		g_assert (next_addr->af == AF_INET);
+		g_assert (addr != next_addr);
+		addr = next_addr;
+	}
+	rspamd_upstreams_destroy (nls);
+
 	/* Upstream fail test */
 	evtimer_set (&ev, rspamd_upstream_timeout_handler, resolver);
 	event_base_set (ev_base, &ev);
@@ -143,4 +167,7 @@ rspamd_upstream_test_func (void)
 
 	event_base_loop (ev_base, 0);
 	g_assert (rspamd_upstreams_alive (ls) == 3);
+
+	g_free (cfg);
+	rspamd_upstreams_destroy (ls);
 }