From c8b19587f0bb34a4f60fb782b641748d9f38b5b5 Mon Sep 17 00:00:00 2001
From: Valentin Gosu <valentin.gosu@gmail.com>
Date: Thu, 22 Feb 2024 12:46:58 +0000
Subject: [PATCH] Bug 1881437 - Use a thread local __res_state for HTTPS
 res_nquery r=necko-reviewers,kershaw

This makes it so each thread has its own __res_state instead of using
the global _res which could cause data races.

Differential Revision: https://phabricator.services.mozilla.com/D202400
---
 netwerk/dns/GetAddrInfo.h          |  2 ++
 netwerk/dns/PlatformDNSAndroid.cpp |  2 ++
 netwerk/dns/PlatformDNSUnix.cpp    | 35 +++++++++++++++++++++++++++---
 netwerk/dns/PlatformDNSWin.cpp     |  2 ++
 netwerk/dns/nsHostResolver.cpp     | 21 ++++++++++++++++++
 5 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/netwerk/dns/GetAddrInfo.h b/netwerk/dns/GetAddrInfo.h
index e8dc09541574c..c8b4ebae69737 100644
--- a/netwerk/dns/GetAddrInfo.h
+++ b/netwerk/dns/GetAddrInfo.h
@@ -67,6 +67,8 @@ nsresult GetAddrInfoInit();
  */
 nsresult GetAddrInfoShutdown();
 
+void DNSThreadShutdown();
+
 /**
  * Resolves a HTTPS record. Will check overrides before calling the
  * native OS implementation.
diff --git a/netwerk/dns/PlatformDNSAndroid.cpp b/netwerk/dns/PlatformDNSAndroid.cpp
index 78a944777f041..171797b938202 100644
--- a/netwerk/dns/PlatformDNSAndroid.cpp
+++ b/netwerk/dns/PlatformDNSAndroid.cpp
@@ -151,4 +151,6 @@ nsresult ResolveHTTPSRecordImpl(const nsACString& aHost, uint16_t aFlags,
   return NS_OK;
 }
 
+void DNSThreadShutdown() {}
+
 }  // namespace mozilla::net
diff --git a/netwerk/dns/PlatformDNSUnix.cpp b/netwerk/dns/PlatformDNSUnix.cpp
index dedbe337df883..c7f57fcddad18 100644
--- a/netwerk/dns/PlatformDNSUnix.cpp
+++ b/netwerk/dns/PlatformDNSUnix.cpp
@@ -9,6 +9,7 @@
 #include "nsIDNSService.h"
 #include "mozilla/Maybe.h"
 #include "mozilla/StaticPrefs_network.h"
+#include "mozilla/ThreadLocal.h"
 
 #include <stdio.h>
 #include <stdlib.h>
@@ -18,6 +19,10 @@
 
 namespace mozilla::net {
 
+#if defined(HAVE_RES_NINIT)
+MOZ_THREAD_LOCAL(struct __res_state*) sThreadRes;
+#endif
+
 #define LOG(msg, ...) \
   MOZ_LOG(gGetAddrInfoLog, LogLevel::Debug, ("[DNS]: " msg, ##__VA_ARGS__))
 
@@ -33,16 +38,28 @@ nsresult ResolveHTTPSRecordImpl(const nsACString& aHost, uint16_t aFlags,
     return NS_ERROR_UNKNOWN_HOST;
   }
 
+#if defined(HAVE_RES_NINIT)
+  if (!sThreadRes.get()) {
+    UniquePtr<struct __res_state> resState(new struct __res_state);
+    memset(resState.get(), 0, sizeof(struct __res_state));
+    if (int ret = res_ninit(resState.get())) {
+      LOG("res_ninit failed: %d", ret);
+      return NS_ERROR_UNKNOWN_HOST;
+    }
+    sThreadRes.set(resState.release());
+  }
+#endif
+
   LOG("resolving %s\n", host.get());
   // Perform the query
   rv = packet.FillBuffer(
       [&](unsigned char response[DNSPacket::MAX_SIZE]) -> int {
         int len = 0;
-#if defined(XP_LINUX)
-        len = res_nquery(&_res, host.get(), ns_c_in,
+#if defined(HAVE_RES_NINIT)
+        len = res_nquery(sThreadRes.get(), host.get(), ns_c_in,
                          nsIDNSService::RESOLVE_TYPE_HTTPSSVC, response,
                          DNSPacket::MAX_SIZE);
-#elif defined(XP_MACOSX)
+#else
         len =
             res_query(host.get(), ns_c_in, nsIDNSService::RESOLVE_TYPE_HTTPSSVC,
                       response, DNSPacket::MAX_SIZE);
@@ -60,4 +77,16 @@ nsresult ResolveHTTPSRecordImpl(const nsACString& aHost, uint16_t aFlags,
   return ParseHTTPSRecord(host, packet, aResult, aTTL);
 }
 
+void DNSThreadShutdown() {
+#if defined(HAVE_RES_NINIT)
+  auto* res = sThreadRes.get();
+  if (!res) {
+    return;
+  }
+
+  sThreadRes.set(nullptr);
+  res_nclose(res);
+#endif
+}
+
 }  // namespace mozilla::net
diff --git a/netwerk/dns/PlatformDNSWin.cpp b/netwerk/dns/PlatformDNSWin.cpp
index 8b5539d7de8de..42f1483ec3274 100644
--- a/netwerk/dns/PlatformDNSWin.cpp
+++ b/netwerk/dns/PlatformDNSWin.cpp
@@ -137,4 +137,6 @@ nsresult ResolveHTTPSRecordImpl(const nsACString& aHost, uint16_t aFlags,
   return NS_OK;
 }
 
+void DNSThreadShutdown() {}
+
 }  // namespace mozilla::net
diff --git a/netwerk/dns/nsHostResolver.cpp b/netwerk/dns/nsHostResolver.cpp
index ccd6dca57b759..ad28cbb284faa 100644
--- a/netwerk/dns/nsHostResolver.cpp
+++ b/netwerk/dns/nsHostResolver.cpp
@@ -3,6 +3,7 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
 
+#include "nsIThreadPool.h"
 #if defined(HAVE_RES_NINIT)
 #  include <sys/types.h>
 #  include <netinet/in.h>
@@ -137,6 +138,24 @@ class nsResState {
 
 #endif  // RES_RETRY_ON_FAILURE
 
+class DnsThreadListener final : public nsIThreadPoolListener {
+  NS_DECL_THREADSAFE_ISUPPORTS
+  NS_DECL_NSITHREADPOOLLISTENER
+ private:
+  virtual ~DnsThreadListener() = default;
+};
+
+NS_IMETHODIMP
+DnsThreadListener::OnThreadCreated() { return NS_OK; }
+
+NS_IMETHODIMP
+DnsThreadListener::OnThreadShuttingDown() {
+  DNSThreadShutdown();
+  return NS_OK;
+}
+
+NS_IMPL_ISUPPORTS(DnsThreadListener, nsIThreadPoolListener)
+
 //----------------------------------------------------------------------------
 
 static const char kPrefGetTtl[] = "network.dns.get-ttl";
@@ -252,6 +271,8 @@ nsresult nsHostResolver::Init() MOZ_NO_THREAD_SAFETY_ANALYSIS {
   MOZ_ALWAYS_SUCCEEDS(
       threadPool->SetThreadStackSize(nsIThreadManager::kThreadPoolStackSize));
   MOZ_ALWAYS_SUCCEEDS(threadPool->SetName("DNS Resolver"_ns));
+  nsCOMPtr<nsIThreadPoolListener> listener = new DnsThreadListener();
+  threadPool->SetListener(listener);
   mResolverThreads = ToRefPtr(std::move(threadPool));
 
   return NS_OK;