From 17ee8ff0d933c44414a625366da07c15350a298e Mon Sep 17 00:00:00 2001
From: Chuck <chuck+github@borboggle.com>
Date: Mon, 24 May 2021 16:27:50 +0700
Subject: [PATCH] Support multiple address resolution in DNS requests

Add two new functions to the IP class that returns all addresses/aliases associated with a given address.

This is a cherry-pick merge from 010a3433df43a94fee95474360ffa6662c7441b9 which was merged in 2.1, and has been updated to build with the latest code.

This merge adds two new methods IP.resolve_hostname_addresses and IP.get_resolve_item_addresses that returns a List of all addresses returned from the DNS request.
---
 core/io/ip.cpp           | 116 ++++++++++++++++++++++++++-------------
 core/io/ip.h             |   5 +-
 doc/classes/IP.xml       |  20 +++++++
 drivers/unix/ip_unix.cpp |  22 ++++++--
 drivers/unix/ip_unix.h   |   2 +-
 5 files changed, 118 insertions(+), 47 deletions(-)

diff --git a/core/io/ip.cpp b/core/io/ip.cpp
index ecd979fb34f..99ea19f18c7 100644
--- a/core/io/ip.cpp
+++ b/core/io/ip.cpp
@@ -41,13 +41,13 @@ VARIANT_ENUM_CAST(IP::ResolverStatus);
 struct _IP_ResolverPrivate {
 	struct QueueItem {
 		SafeNumeric<IP::ResolverStatus> status;
-		IP_Address response;
+		List<IP_Address> response;
 		String hostname;
 		IP::Type type;
 
 		void clear() {
 			status.set(IP::RESOLVER_STATUS_NONE);
-			response = IP_Address();
+			response.clear();
 			type = IP::TYPE_NONE;
 			hostname = "";
 		};
@@ -80,13 +80,9 @@ struct _IP_ResolverPrivate {
 			if (queue[i].status.get() != IP::RESOLVER_STATUS_WAITING) {
 				continue;
 			}
-			queue[i].response = IP::get_singleton()->resolve_hostname(queue[i].hostname, queue[i].type);
 
-			if (!queue[i].response.is_valid()) {
-				queue[i].status.set(IP::RESOLVER_STATUS_ERROR);
-			} else {
-				queue[i].status.set(IP::RESOLVER_STATUS_DONE);
-			}
+			IP::get_singleton()->_resolve_hostname(queue[i].response, queue[i].hostname, queue[i].type);
+			queue[i].status.set(queue[i].response.empty() ? IP::RESOLVER_STATUS_ERROR : IP::RESOLVER_STATUS_DONE);
 		}
 	}
 
@@ -96,13 +92,12 @@ struct _IP_ResolverPrivate {
 		while (!ipr->thread_abort) {
 			ipr->sem.wait();
 
-			ipr->mutex.lock();
+			MutexLock lock(ipr->mutex);
 			ipr->resolve_queues();
-			ipr->mutex.unlock();
 		}
 	}
 
-	HashMap<String, IP_Address> cache;
+	HashMap<String, List<IP_Address>> cache;
 
 	static String get_cache_key(String p_hostname, IP::Type p_type) {
 		return itos(p_type) + p_hostname;
@@ -110,40 +105,63 @@ struct _IP_ResolverPrivate {
 };
 
 IP_Address IP::resolve_hostname(const String &p_hostname, IP::Type p_type) {
-	resolver->mutex.lock();
+	MutexLock lock(resolver->mutex);
+
+	List<IP_Address> res;
 
 	String key = _IP_ResolverPrivate::get_cache_key(p_hostname, p_type);
-	if (resolver->cache.has(key) && resolver->cache[key].is_valid()) {
-		IP_Address res = resolver->cache[key];
-		resolver->mutex.unlock();
-		return res;
+	if (resolver->cache.has(key)) {
+		res = resolver->cache[key];
+	} else {
+		_resolve_hostname(res, p_hostname, p_type);
+		resolver->cache[key] = res;
 	}
 
-	IP_Address res = _resolve_hostname(p_hostname, p_type);
-	resolver->cache[key] = res;
-	resolver->mutex.unlock();
-	return res;
+	for (int i = 0; i < res.size(); ++i) {
+		if (res[i].is_valid()) {
+			return res[i];
+		}
+	}
+	return IP_Address();
+}
+
+Array IP::resolve_hostname_addresses(const String &p_hostname, Type p_type) {
+	MutexLock lock(resolver->mutex);
+
+	String key = _IP_ResolverPrivate::get_cache_key(p_hostname, p_type);
+	if (!resolver->cache.has(key)) {
+		_resolve_hostname(resolver->cache[key], p_hostname, p_type);
+	}
+
+	List<IP_Address> res = resolver->cache[key];
+
+	Array result;
+	for (int i = 0; i < res.size(); ++i) {
+		if (res[i].is_valid()) {
+			result.push_back(String(res[i]));
+		}
+	}
+	return result;
 }
 
 IP::ResolverID IP::resolve_hostname_queue_item(const String &p_hostname, IP::Type p_type) {
-	resolver->mutex.lock();
+	MutexLock lock(resolver->mutex);
 
 	ResolverID id = resolver->find_empty_id();
 
 	if (id == RESOLVER_INVALID_ID) {
 		WARN_PRINT("Out of resolver queries");
-		resolver->mutex.unlock();
 		return id;
 	}
 
 	String key = _IP_ResolverPrivate::get_cache_key(p_hostname, p_type);
 	resolver->queue[id].hostname = p_hostname;
 	resolver->queue[id].type = p_type;
-	if (resolver->cache.has(key) && resolver->cache[key].is_valid()) {
+	if (resolver->cache.has(key)) {
 		resolver->queue[id].response = resolver->cache[key];
 		resolver->queue[id].status.set(IP::RESOLVER_STATUS_DONE);
 	} else {
-		resolver->queue[id].response = IP_Address();
+		resolver->queue[id].response = List<IP_Address>();
 		resolver->queue[id].status.set(IP::RESOLVER_STATUS_WAITING);
 		if (resolver->thread.is_started()) {
 			resolver->sem.post();
@@ -152,54 +170,74 @@ IP::ResolverID IP::resolve_hostname_queue_item(const String &p_hostname, IP::Typ
 		}
 	}
 
-	resolver->mutex.unlock();
 	return id;
 }
 
 IP::ResolverStatus IP::get_resolve_item_status(ResolverID p_id) const {
 	ERR_FAIL_INDEX_V(p_id, IP::RESOLVER_MAX_QUERIES, IP::RESOLVER_STATUS_NONE);
 
-	resolver->mutex.lock();
+	MutexLock lock(resolver->mutex);
+
 	if (resolver->queue[p_id].status.get() == IP::RESOLVER_STATUS_NONE) {
 		ERR_PRINT("Condition status == IP::RESOLVER_STATUS_NONE");
-		resolver->mutex.unlock();
 		return IP::RESOLVER_STATUS_NONE;
 	}
 	IP::ResolverStatus res = resolver->queue[p_id].status.get();
 
-	resolver->mutex.unlock();
 	return res;
 }
 
 IP_Address IP::get_resolve_item_address(ResolverID p_id) const {
 	ERR_FAIL_INDEX_V(p_id, IP::RESOLVER_MAX_QUERIES, IP_Address());
 
-	resolver->mutex.lock();
+	MutexLock lock(resolver->mutex);
 
 	if (resolver->queue[p_id].status.get() != IP::RESOLVER_STATUS_DONE) {
 		ERR_PRINTS("Resolve of '" + resolver->queue[p_id].hostname + "'' didn't complete yet.");
-		resolver->mutex.unlock();
 		return IP_Address();
 	}
 
-	IP_Address res = resolver->queue[p_id].response;
+	List<IP_Address> res = resolver->queue[p_id].response;
 
-	resolver->mutex.unlock();
-	return res;
+	for (int i = 0; i < res.size(); ++i) {
+		if (res[i].is_valid()) {
+			return res[i];
+		}
+	}
+	return IP_Address();
+}
+
+Array IP::get_resolve_item_addresses(ResolverID p_id) const {
+	ERR_FAIL_INDEX_V(p_id, IP::RESOLVER_MAX_QUERIES, Array());
+
+	MutexLock lock(resolver->mutex);
+
+	if (resolver->queue[p_id].status.get() != IP::RESOLVER_STATUS_DONE) {
+		ERR_PRINTS("Resolve of '" + resolver->queue[p_id].hostname + "'' didn't complete yet.");
+		return Array();
+	}
+
+	List<IP_Address> res = resolver->queue[p_id].response;
+
+	Array result;
+	for (int i = 0; i < res.size(); ++i) {
+		if (res[i].is_valid()) {
+			result.push_back(String(res[i]));
+		}
+	}
+	return result;
 }
 
 void IP::erase_resolve_item(ResolverID p_id) {
 	ERR_FAIL_INDEX(p_id, IP::RESOLVER_MAX_QUERIES);
 
-	resolver->mutex.lock();
+	MutexLock lock(resolver->mutex);
 
 	resolver->queue[p_id].status.set(IP::RESOLVER_STATUS_NONE);
-
-	resolver->mutex.unlock();
 }
 
 void IP::clear_cache(const String &p_hostname) {
-	resolver->mutex.lock();
+	MutexLock lock(resolver->mutex);
 
 	if (p_hostname.empty()) {
 		resolver->cache.clear();
@@ -209,8 +247,6 @@ void IP::clear_cache(const String &p_hostname) {
 		resolver->cache.erase(_IP_ResolverPrivate::get_cache_key(p_hostname, IP::TYPE_IPV6));
 		resolver->cache.erase(_IP_ResolverPrivate::get_cache_key(p_hostname, IP::TYPE_ANY));
 	}
-
-	resolver->mutex.unlock();
 }
 
 Array IP::_get_local_addresses() const {
@@ -259,9 +295,11 @@ void IP::get_local_addresses(List<IP_Address> *r_addresses) const {
 
 void IP::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("resolve_hostname", "host", "ip_type"), &IP::resolve_hostname, DEFVAL(IP::TYPE_ANY));
+	ClassDB::bind_method(D_METHOD("resolve_hostname_addresses", "host", "ip_type"), &IP::resolve_hostname_addresses, DEFVAL(IP::TYPE_ANY));
 	ClassDB::bind_method(D_METHOD("resolve_hostname_queue_item", "host", "ip_type"), &IP::resolve_hostname_queue_item, DEFVAL(IP::TYPE_ANY));
 	ClassDB::bind_method(D_METHOD("get_resolve_item_status", "id"), &IP::get_resolve_item_status);
 	ClassDB::bind_method(D_METHOD("get_resolve_item_address", "id"), &IP::get_resolve_item_address);
+	ClassDB::bind_method(D_METHOD("get_resolve_item_addresses", "id"), &IP::get_resolve_item_addresses);
 	ClassDB::bind_method(D_METHOD("erase_resolve_item", "id"), &IP::erase_resolve_item);
 	ClassDB::bind_method(D_METHOD("get_local_addresses"), &IP::_get_local_addresses);
 	ClassDB::bind_method(D_METHOD("get_local_interfaces"), &IP::_get_local_interfaces);
diff --git a/core/io/ip.h b/core/io/ip.h
index 14c09319933..bb3bab53412 100644
--- a/core/io/ip.h
+++ b/core/io/ip.h
@@ -71,7 +71,6 @@ protected:
 	static IP *singleton;
 	static void _bind_methods();
 
-	virtual IP_Address _resolve_hostname(const String &p_hostname, Type p_type = TYPE_ANY) = 0;
 	Array _get_local_addresses() const;
 	Array _get_local_interfaces() const;
 
@@ -86,11 +85,15 @@ public:
 	};
 
 	IP_Address resolve_hostname(const String &p_hostname, Type p_type = TYPE_ANY);
+	Array resolve_hostname_addresses(const String &p_hostname, Type p_type = TYPE_ANY);
 	// async resolver hostname
 	ResolverID resolve_hostname_queue_item(const String &p_hostname, Type p_type = TYPE_ANY);
 	ResolverStatus get_resolve_item_status(ResolverID p_id) const;
 	IP_Address get_resolve_item_address(ResolverID p_id) const;
 	virtual void get_local_addresses(List<IP_Address> *r_addresses) const;
+
+	virtual void _resolve_hostname(List<IP_Address> &r_addresses, const String &p_hostname, Type p_type = TYPE_ANY) const = 0;
+	Array get_resolve_item_addresses(ResolverID p_id) const;
 	virtual void get_local_interfaces(Map<String, Interface_Info> *r_interfaces) const = 0;
 	void erase_resolve_item(ResolverID p_id);
 
diff --git a/doc/classes/IP.xml b/doc/classes/IP.xml
index 46f5e457495..c2398867b4b 100644
--- a/doc/classes/IP.xml
+++ b/doc/classes/IP.xml
@@ -59,6 +59,15 @@
 				Returns a queued hostname's IP address, given its queue [code]id[/code]. Returns an empty string on error or if resolution hasn't happened yet (see [method get_resolve_item_status]).
 			</description>
 		</method>
+		<method name="get_resolve_item_addresses" qualifiers="const">
+			<return type="Array">
+			</return>
+			<argument index="0" name="id" type="int">
+			</argument>
+			<description>
+				Return resolved addresses, or an empty array if an error happened or resolution didn't happen yet (see [method get_resolve_item_status]).
+			</description>
+		</method>
 		<method name="get_resolve_item_status" qualifiers="const">
 			<return type="int" enum="IP.ResolverStatus">
 			</return>
@@ -79,6 +88,17 @@
 				Returns a given hostname's IPv4 or IPv6 address when resolved (blocking-type method). The address type returned depends on the [enum Type] constant given as [code]ip_type[/code].
 			</description>
 		</method>
+		<method name="resolve_hostname_addresses">
+			<return type="Array">
+			</return>
+			<argument index="0" name="host" type="String">
+			</argument>
+			<argument index="1" name="ip_type" type="int" enum="IP.Type" default="3">
+			</argument>
+			<description>
+				Resolves a given hostname in a blocking way. Addresses are returned as an [Array] of IPv4 or IPv6 depending on [code]ip_type[/code].
+			</description>
+		</method>
 		<method name="resolve_hostname_queue_item">
 			<return type="int">
 			</return>
diff --git a/drivers/unix/ip_unix.cpp b/drivers/unix/ip_unix.cpp
index 8f1a2b712e2..f4d62745f7f 100644
--- a/drivers/unix/ip_unix.cpp
+++ b/drivers/unix/ip_unix.cpp
@@ -89,7 +89,7 @@ static IP_Address _sockaddr2ip(struct sockaddr *p_addr) {
 	return ip;
 };
 
-IP_Address IP_Unix::_resolve_hostname(const String &p_hostname, Type p_type) {
+void IP_Unix::_resolve_hostname(List<IP_Address> &r_addresses, const String &p_hostname, Type p_type) const {
 	struct addrinfo hints;
 	struct addrinfo *result;
 
@@ -108,7 +108,7 @@ IP_Address IP_Unix::_resolve_hostname(const String &p_hostname, Type p_type) {
 	int s = getaddrinfo(p_hostname.utf8().get_data(), nullptr, &hints, &result);
 	if (s != 0) {
 		ERR_PRINT("getaddrinfo failed! Cannot resolve hostname.");
-		return IP_Address();
+		return;
 	};
 
 	if (result == nullptr || result->ai_addr == nullptr) {
@@ -116,14 +116,24 @@ IP_Address IP_Unix::_resolve_hostname(const String &p_hostname, Type p_type) {
 		if (result) {
 			freeaddrinfo(result);
 		}
-		return IP_Address();
+		return;
 	};
 
-	IP_Address ip = _sockaddr2ip(result->ai_addr);
+	struct addrinfo *next = result;
+
+	do {
+		if (next->ai_addr == NULL) {
+			next = next->ai_next;
+			continue;
+		}
+		IP_Address ip = _sockaddr2ip(next->ai_addr);
+		if (!r_addresses.find(ip)) {
+			r_addresses.push_back(ip);
+		}
+		next = next->ai_next;
+	} while (next);
 
 	freeaddrinfo(result);
-
-	return ip;
 }
 
 #if defined(WINDOWS_ENABLED)
diff --git a/drivers/unix/ip_unix.h b/drivers/unix/ip_unix.h
index a5224e52308..e5c4e056032 100644
--- a/drivers/unix/ip_unix.h
+++ b/drivers/unix/ip_unix.h
@@ -38,7 +38,7 @@
 class IP_Unix : public IP {
 	GDCLASS(IP_Unix, IP);
 
-	virtual IP_Address _resolve_hostname(const String &p_hostname, IP::Type p_type);
+	virtual void _resolve_hostname(List<IP_Address> &r_addresses, const String &p_hostname, Type p_type = TYPE_ANY) const;
 
 	static IP *_create_unix();