From 5836e02b95696a39fc643e9c62ebd9c507634242 Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Mon, 4 Mar 2019 14:15:24 +0000
Subject: [PATCH] [Project] Allow to hash and compare inet addresses
 considering port

---
 src/libmime/message.c  |   2 +-
 src/libutil/addr.c     | 115 ++++++++++++++++++++++++++++++++++++-----
 src/libutil/addr.h     |   7 +--
 src/libutil/upstream.c |   2 +-
 src/lua/lua_ip.c       |   2 +-
 5 files changed, 108 insertions(+), 20 deletions(-)

diff --git a/src/libmime/message.c b/src/libmime/message.c
index da065195d..cca134f81 100644
--- a/src/libmime/message.c
+++ b/src/libmime/message.c
@@ -1150,7 +1150,7 @@ rspamd_message_parse (struct rspamd_task *task)
 				need_recv_correction = TRUE;
 			}
 			else {
-				if (rspamd_inet_address_compare (raddr, task->from_addr) != 0) {
+				if (rspamd_inet_address_compare (raddr, task->from_addr, FALSE) != 0) {
 					need_recv_correction = TRUE;
 				}
 			}
diff --git a/src/libutil/addr.c b/src/libutil/addr.c
index fd4621aef..7fa634f6f 100644
--- a/src/libutil/addr.c
+++ b/src/libutil/addr.c
@@ -1636,7 +1636,7 @@ rspamd_inet_address_af_order (const rspamd_inet_addr_t *addr)
 
 gint
 rspamd_inet_address_compare (const rspamd_inet_addr_t *a1,
-		const rspamd_inet_addr_t *a2)
+		const rspamd_inet_addr_t *a2, gboolean compare_ports)
 {
 	g_assert (a1 != NULL);
 	g_assert (a2 != NULL);
@@ -1648,11 +1648,33 @@ rspamd_inet_address_compare (const rspamd_inet_addr_t *a1,
 	else {
 		switch (a1->af) {
 		case AF_INET:
-			return memcmp (&a1->u.in.addr.s4.sin_addr,
-					&a2->u.in.addr.s4.sin_addr, sizeof (struct in_addr));
+			if (!compare_ports) {
+				return memcmp (&a1->u.in.addr.s4.sin_addr,
+						&a2->u.in.addr.s4.sin_addr, sizeof (struct in_addr));
+			}
+			else {
+				if (a1->u.in.addr.s4.sin_port == a2->u.in.addr.s4.sin_port) {
+					return memcmp (&a1->u.in.addr.s4.sin_addr,
+							&a2->u.in.addr.s4.sin_addr, sizeof (struct in_addr));
+				}
+				else {
+					return a1->u.in.addr.s4.sin_port - a2->u.in.addr.s4.sin_port;
+				}
+			}
 		case AF_INET6:
-			return memcmp (&a1->u.in.addr.s6.sin6_addr,
-				&a2->u.in.addr.s6.sin6_addr, sizeof (struct in6_addr));
+			if (!compare_ports) {
+				return memcmp (&a1->u.in.addr.s6.sin6_addr,
+						&a2->u.in.addr.s6.sin6_addr, sizeof (struct in6_addr));
+			}
+			else {
+				if (a1->u.in.addr.s6.sin6_port == a2->u.in.addr.s6.sin6_port) {
+					return memcmp (&a1->u.in.addr.s6.sin6_addr,
+							&a2->u.in.addr.s6.sin6_addr, sizeof (struct in6_addr));
+				}
+				else {
+					return a1->u.in.addr.s6.sin6_port - a2->u.in.addr.s6.sin6_port;
+				}
+			}
 		case AF_UNIX:
 			return strncmp (a1->u.un->addr.sun_path,
 				a2->u.un->addr.sun_path, sizeof (a1->u.un->addr.sun_path));
@@ -1671,7 +1693,7 @@ rspamd_inet_address_compare_ptr (gconstpointer a1,
 	const rspamd_inet_addr_t **i1 = (const rspamd_inet_addr_t **)a1,
 			**i2 = (const rspamd_inet_addr_t **)a2;
 
-	return rspamd_inet_address_compare (*i1, *i2);
+	return rspamd_inet_address_compare (*i1, *i2, FALSE);
 }
 
 rspamd_inet_addr_t *
@@ -1724,28 +1746,85 @@ guint
 rspamd_inet_address_hash (gconstpointer a)
 {
 	const rspamd_inet_addr_t *addr = a;
-	rspamd_cryptobox_fast_hash_state_t st;
-
-	rspamd_cryptobox_fast_hash_init (&st, rspamd_hash_seed ());
-	rspamd_cryptobox_fast_hash_update (&st, &addr->af, sizeof (addr->af));
+	struct {
+		gchar buf[sizeof(struct in6_addr)]; /* 16 bytes */
+		int af;
+	} layout;
 
+	gint32 k;
 
 	if (addr->af == AF_UNIX && addr->u.un) {
+		rspamd_cryptobox_fast_hash_state_t st;
+
+		rspamd_cryptobox_fast_hash_init (&st, rspamd_hash_seed ());
+		rspamd_cryptobox_fast_hash_update (&st, &addr->af, sizeof (addr->af));
 		rspamd_cryptobox_fast_hash_update (&st, addr->u.un, sizeof (*addr->u.un));
+
+		return rspamd_cryptobox_fast_hash_final (&st);
 	}
 	else {
+		memset (&layout, 0, sizeof (layout));
+		layout.af = addr->af;
+
 		/* We ignore port part here */
 		if (addr->af == AF_INET) {
-			rspamd_cryptobox_fast_hash_update (&st, &addr->u.in.addr.s4.sin_addr,
+			memcpy (layout.buf, &addr->u.in.addr.s4.sin_addr,
+					sizeof (addr->u.in.addr.s4.sin_addr));
+		}
+		else {
+			memcpy (layout.buf, &addr->u.in.addr.s6.sin6_addr,
+					sizeof (addr->u.in.addr.s6.sin6_addr));
+		}
+
+		k = rspamd_cryptobox_fast_hash (&layout, sizeof (layout),
+				rspamd_hash_seed ());
+	}
+
+	return k;
+}
+
+guint
+rspamd_inet_address_port_hash (gconstpointer a)
+{
+	const rspamd_inet_addr_t *addr = a;
+	struct {
+		gchar buf[sizeof(struct in6_addr)]; /* 16 bytes */
+		int port;
+		int af;
+	} layout;
+
+	gint32 k;
+
+	if (addr->af == AF_UNIX && addr->u.un) {
+		rspamd_cryptobox_fast_hash_state_t st;
+
+		rspamd_cryptobox_fast_hash_init (&st, rspamd_hash_seed ());
+		rspamd_cryptobox_fast_hash_update (&st, &addr->af, sizeof (addr->af));
+		rspamd_cryptobox_fast_hash_update (&st, addr->u.un, sizeof (*addr->u.un));
+
+		return rspamd_cryptobox_fast_hash_final (&st);
+	}
+	else {
+		memset (&layout, 0, sizeof (layout));
+		layout.af = addr->af;
+
+		/* We consider port part here */
+		if (addr->af == AF_INET) {
+			memcpy (layout.buf, &addr->u.in.addr.s4.sin_addr,
 					sizeof (addr->u.in.addr.s4.sin_addr));
+			layout.port = addr->u.in.addr.s4.sin_port;
 		}
 		else {
-			rspamd_cryptobox_fast_hash_update (&st, &addr->u.in.addr.s6.sin6_addr,
+			memcpy (layout.buf, &addr->u.in.addr.s6.sin6_addr,
 					sizeof (addr->u.in.addr.s6.sin6_addr));
+			layout.port = addr->u.in.addr.s6.sin6_port;
 		}
+
+		k = rspamd_cryptobox_fast_hash (&layout, sizeof (layout),
+				rspamd_hash_seed ());
 	}
 
-	return rspamd_cryptobox_fast_hash_final (&st);
+	return k;
 }
 
 gboolean
@@ -1753,7 +1832,15 @@ rspamd_inet_address_equal (gconstpointer a, gconstpointer b)
 {
 	const rspamd_inet_addr_t *a1 = a, *a2 = b;
 
-	return rspamd_inet_address_compare (a1, a2) == 0;
+	return rspamd_inet_address_compare (a1, a2, FALSE) == 0;
+}
+
+gboolean
+rspamd_inet_address_port_equal (gconstpointer a, gconstpointer b)
+{
+	const rspamd_inet_addr_t *a1 = a, *a2 = b;
+
+	return rspamd_inet_address_compare (a1, a2, TRUE) == 0;
 }
 
 #ifndef IN6_IS_ADDR_LOOPBACK
diff --git a/src/libutil/addr.h b/src/libutil/addr.h
index 46b705a4b..bfe586ad1 100644
--- a/src/libutil/addr.h
+++ b/src/libutil/addr.h
@@ -263,7 +263,7 @@ void rspamd_inet_address_apply_mask (rspamd_inet_addr_t *addr, guint mask);
  * @return
  */
 gint rspamd_inet_address_compare (const rspamd_inet_addr_t *a1,
-		const rspamd_inet_addr_t *a2);
+		const rspamd_inet_addr_t *a2, gboolean compare_ports);
 
 /**
  * Utility function to compare addresses by in g_ptr_array
@@ -281,15 +281,16 @@ gint rspamd_inet_address_compare_ptr (gconstpointer a1,
 rspamd_inet_addr_t *rspamd_inet_address_copy (const rspamd_inet_addr_t *addr);
 
 /**
- * Returns hash for inet address
+ * Returns hash for inet address (ignoring port)
  */
 guint rspamd_inet_address_hash (gconstpointer a);
-
+guint rspamd_inet_address_port_hash (gconstpointer a);
 
 /**
  * Returns true if two address are equal
  */
 gboolean rspamd_inet_address_equal (gconstpointer a, gconstpointer b);
+gboolean rspamd_inet_address_port_equal (gconstpointer a, gconstpointer b);
 
 /**
  * Returns TRUE if an address belongs to some local address
diff --git a/src/libutil/upstream.c b/src/libutil/upstream.c
index 2ce294525..64d5291fa 100644
--- a/src/libutil/upstream.c
+++ b/src/libutil/upstream.c
@@ -289,7 +289,7 @@ rspamd_upstream_update_addrs (struct upstream *up)
 			rspamd_inet_address_set_port (cur->addr, port);
 
 			PTR_ARRAY_FOREACH (up->addrs.addr, i, addr_elt) {
-				if (rspamd_inet_address_compare (addr_elt->addr, cur->addr) == 0) {
+				if (rspamd_inet_address_compare (addr_elt->addr, cur->addr, FALSE) == 0) {
 					naddr = g_malloc0 (sizeof (*naddr));
 					naddr->addr = cur->addr;
 					naddr->errors = reset_errors ? 0 : addr_elt->errors;
diff --git a/src/lua/lua_ip.c b/src/lua/lua_ip.c
index 5ee4a39b4..7f24f7712 100644
--- a/src/lua/lua_ip.c
+++ b/src/lua/lua_ip.c
@@ -490,7 +490,7 @@ lua_ip_equal (lua_State *L)
 	gboolean res = FALSE;
 
 	if (ip1 && ip2 && ip1->addr && ip2->addr) {
-		res = rspamd_inet_address_compare (ip1->addr, ip2->addr);
+		res = rspamd_inet_address_compare (ip1->addr, ip2->addr, TRUE);
 	}
 
 	lua_pushboolean (L, res);
-- 
2.39.5