From da982de896b7853fc57e50862fb395060261aca6 Mon Sep 17 00:00:00 2001
From: xiehan <52160700+Barenboim@users.noreply.github.com>
Date: Sun, 14 Apr 2024 02:14:50 +0800
Subject: [PATCH] Free 'addrinfo' directly when it's a UDS address. (#1523)

---
 src/manager/DnsCache.h           | 11 +++++++----
 src/nameservice/WFDnsResolver.cc | 30 +++++++++++++++---------------
 src/protocol/DnsUtil.cc          |  4 ++--
 3 files changed, 24 insertions(+), 21 deletions(-)

diff --git a/src/manager/DnsCache.h b/src/manager/DnsCache.h
index 5ef1a26a..bd301462 100644
--- a/src/manager/DnsCache.h
+++ b/src/manager/DnsCache.h
@@ -148,10 +148,13 @@ private:
 		{
 			struct addrinfo *ai = value.addrinfo;
 
-			if (ai && (ai->ai_flags & 1))
-				freeaddrinfo(ai);
-			else
-				protocol::DnsUtil::freeaddrinfo(ai);
+			if (ai)
+			{
+				if (ai->ai_flags)
+					freeaddrinfo(ai);
+				else
+					protocol::DnsUtil::freeaddrinfo(ai);
+			}
 		}
 	};
 
diff --git a/src/nameservice/WFDnsResolver.cc b/src/nameservice/WFDnsResolver.cc
index 961318b3..b718e0bf 100644
--- a/src/nameservice/WFDnsResolver.cc
+++ b/src/nameservice/WFDnsResolver.cc
@@ -102,7 +102,12 @@ public:
 	~DnsOutput()
 	{
 		if (addrinfo_)
-			freeaddrinfo(addrinfo_);
+		{
+			if (addrinfo_->ai_flags)
+				freeaddrinfo(addrinfo_);
+			else
+				free(addrinfo_);
+		}
 	}
 
 	int get_error() const { return error_; }
@@ -130,7 +135,12 @@ public:
 	static void create(DnsOutput *out, int error, struct addrinfo *ai)
 	{
 		if (out->addrinfo_)
-			freeaddrinfo(out->addrinfo_);
+		{
+			if (out->addrinfo_->ai_flags)
+				freeaddrinfo(out->addrinfo_);
+			else
+				free(out->addrinfo_);
+		}
 
 		out->error_ = error;
 		out->addrinfo_ = ai;
@@ -172,7 +182,7 @@ void DnsRoutine::run_local_path(const std::string& path, DnsOutput *out)
 
 void DnsRoutine::run(const DnsInput *in, DnsOutput *out)
 {
-	if (!in->host_.empty() && in->host_[0] == '/')
+	if (in->host_[0] == '/')
 	{
 		run_local_path(in->host_, out);
 		return;
@@ -191,6 +201,8 @@ void DnsRoutine::run(const DnsInput *in, DnsOutput *out)
 	snprintf(port_str, PORT_STR_MAX + 1, "%u", in->port_);
 	out->error_ = getaddrinfo(in->host_.c_str(), port_str,
 							  &hints, &out->addrinfo_);
+	if (out->error_ == 0)
+		out->addrinfo_->ai_flags = 1;
 }
 
 // Dns Thread task. For internal usage only.
@@ -322,15 +334,6 @@ static int __readaddrinfo(const char *path,
 	return ret;
 }
 
-static void __set_thread_dns_flag(struct addrinfo *ai)
-{
-	while (ai)
-	{
-		ai->ai_flags = 1;
-		ai = ai->ai_next;
-	}
-}
-
 static ThreadDnsTask *__create_thread_dns_task(const std::string& host,
 											   unsigned short port,
 											   int family,
@@ -456,7 +459,6 @@ void WFResolverTask::dispatch()
 			DnsOutput dns_out;
 
 			DnsRoutine::run(&dns_in, &dns_out);
-			__set_thread_dns_flag((struct addrinfo *)dns_out.get_addrinfo());
 			dns_callback_internal(&dns_out, (unsigned int)-1, (unsigned int)-1);
 			this->subtask_done();
 			return;
@@ -479,7 +481,6 @@ void WFResolverTask::dispatch()
 		{
 			DnsOutput out;
 			DnsRoutine::create(&out, ret, ai);
-			__set_thread_dns_flag((struct addrinfo *)out.get_addrinfo());
 			dns_callback_internal(&out, dns_ttl_default_, dns_ttl_min_);
 			this->subtask_done();
 			return;
@@ -720,7 +721,6 @@ void WFResolverTask::thread_dns_callback(void *thrd_dns_task)
 	if (dns_task->get_state() == WFT_STATE_SUCCESS)
 	{
 		DnsOutput *out = dns_task->get_output();
-		__set_thread_dns_flag((struct addrinfo *)out->get_addrinfo());
 		dns_callback_internal(out, dns_ttl_default_, dns_ttl_min_);
 	}
 	else
diff --git a/src/protocol/DnsUtil.cc b/src/protocol/DnsUtil.cc
index 6c84d103..8909a52d 100644
--- a/src/protocol/DnsUtil.cc
+++ b/src/protocol/DnsUtil.cc
@@ -93,8 +93,8 @@ int DnsUtil::getaddrinfo(const DnsResponse *resp,
 		if (ai == NULL)
 		{
 			if (res)
-				freeaddrinfo(res);
-			return EAI_MEMORY;
+				DnsUtil::freeaddrinfo(res);
+			return EAI_SYSTEM;
 		}
 
 		ai->ai_family = family;