From e32d22615165d0c6b746c66b93bccb2c247a3a11 Mon Sep 17 00:00:00 2001
From: HoneyryderChuck <cardoso_tiago@hotmail.com>
Date: Wed, 30 Oct 2024 11:29:41 +0000
Subject: [PATCH] refactor of internal resolver cache lookup access to make it
 a bit safer

---
 lib/httpx/resolver.rb | 35 +++++++++++++++++++++--------------
 sig/resolver.rbs      |  6 +++++-
 2 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/lib/httpx/resolver.rb b/lib/httpx/resolver.rb
index 72bcfba0..4837be6b 100644
--- a/lib/httpx/resolver.rb
+++ b/lib/httpx/resolver.rb
@@ -53,8 +53,8 @@ module HTTPX
 
     def cached_lookup(hostname)
       now = Utils.now
-      @lookup_mutex.synchronize do
-        lookup(hostname, now)
+      lookup_synchronize do |lookups|
+        lookup(hostname, lookups, now)
       end
     end
 
@@ -63,37 +63,37 @@ module HTTPX
       entries.each do |entry|
         entry["TTL"] += now
       end
-      @lookup_mutex.synchronize do
+      lookup_synchronize do |lookups|
         case family
         when Socket::AF_INET6
-          @lookups[hostname].concat(entries)
+          lookups[hostname].concat(entries)
         when Socket::AF_INET
-          @lookups[hostname].unshift(*entries)
+          lookups[hostname].unshift(*entries)
         end
         entries.each do |entry|
           next unless entry["name"] != hostname
 
           case family
           when Socket::AF_INET6
-            @lookups[entry["name"]] << entry
+            lookups[entry["name"]] << entry
           when Socket::AF_INET
-            @lookups[entry["name"]].unshift(entry)
+            lookups[entry["name"]].unshift(entry)
           end
         end
       end
     end
 
     # do not use directly!
-    def lookup(hostname, ttl)
-      return unless @lookups.key?(hostname)
+    def lookup(hostname, lookups, ttl)
+      return unless lookups.key?(hostname)
 
-      entries = @lookups[hostname] = @lookups[hostname].select do |address|
+      entries = lookups[hostname] = lookups[hostname].select do |address|
         address["TTL"] > ttl
       end
 
       ips = entries.flat_map do |address|
         if address.key?("alias")
-          lookup(address["alias"], ttl)
+          lookup(address["alias"], lookups, ttl)
         else
           IPAddr.new(address["data"])
         end
@@ -103,12 +103,11 @@ module HTTPX
     end
 
     def generate_id
-      @identifier_mutex.synchronize { @identifier = (@identifier + 1) & 0xFFFF }
+      id_synchronize { @identifier = (@identifier + 1) & 0xFFFF }
     end
 
     def encode_dns_query(hostname, type: Resolv::DNS::Resource::IN::A, message_id: generate_id)
-      Resolv::DNS::Message.new.tap do |query|
-        query.id = message_id
+      Resolv::DNS::Message.new(message_id).tap do |query|
         query.rd = 1
         query.add_question(hostname, type)
       end.encode
@@ -150,5 +149,13 @@ module HTTPX
 
       [:ok, addresses]
     end
+
+    def lookup_synchronize
+      @lookup_mutex.synchronize { yield(@lookups) }
+    end
+
+    def id_synchronize(&block)
+      @identifier_mutex.synchronize(&block)
+    end
   end
 end
diff --git a/sig/resolver.rbs b/sig/resolver.rbs
index 7193184e..8ec7afaa 100644
--- a/sig/resolver.rbs
+++ b/sig/resolver.rbs
@@ -28,12 +28,16 @@ module HTTPX
 
     def self?.cached_lookup_set: (String hostname, ip_family family, Array[dns_result] addresses) -> void
 
-    def self?.lookup: (String hostname, Numeric ttl) -> Array[IPAddr]?
+    def self?.lookup: (String hostname, Hash[String, Array[dns_result]] lookups, Numeric ttl) -> Array[IPAddr]?
 
     def self?.generate_id: () -> Integer
 
     def self?.encode_dns_query: (String hostname, ?type: dns_resource, ?message_id: Integer) -> String
 
     def self?.decode_dns_answer: (String) -> dns_decoding_response
+
+    def self?.lookup_synchronize: [U] () { (Hash[String, Array[dns_result]] lookups) -> U } -> U
+
+    def self?.id_synchronize: () { () -> void } -> void
   end
 end
\ No newline at end of file