From 833617c8fee2ba6a7c910a1589177cc779160cd7 Mon Sep 17 00:00:00 2001
From: Ali Mohammad Pur <ali.mpfard@gmail.com>
Date: Wed, 18 Jun 2025 08:37:54 +0200
Subject: [PATCH] LibDNS+dns: Avoid multi-question queries

This is not supported by much of anything (correctly), so send the
requests in parallel instead.
---
 Libraries/LibDNS/Resolver.h | 35 ++++++++++++++++++++++++-
 Utilities/dns.cpp           | 51 +++++++++++++++++++------------------
 2 files changed, 60 insertions(+), 26 deletions(-)

diff --git a/Libraries/LibDNS/Resolver.h b/Libraries/LibDNS/Resolver.h
index 82b43cda1d3..dbb323e44e8 100644
--- a/Libraries/LibDNS/Resolver.h
+++ b/Libraries/LibDNS/Resolver.h
@@ -301,6 +301,37 @@ public:
         });
     }
 
+    NonnullRefPtr<Core::Promise<NonnullRefPtr<LookupResult const>>> lookup(ByteString name, Messages::Class class_, Vector<Vector<Messages::ResourceType>> desired_types, LookupOptions options = LookupOptions::default_())
+    {
+        using ResultPromise = Core::Promise<NonnullRefPtr<LookupResult const>>;
+        Vector<NonnullRefPtr<ResultPromise>> promises;
+        promises.ensure_capacity(desired_types.size());
+
+        for (auto& types : desired_types)
+            promises.unchecked_append(lookup(name, class_, types, options));
+
+        auto result_promise = Core::Promise<NonnullRefPtr<LookupResult const>>::construct();
+        result_promise->add_child(Core::Promise<Empty>::after(promises)
+                ->when_resolved([promises, result_promise = result_promise->make_weak_ptr<ResultPromise>()](auto&&) {
+                    if (!result_promise.ptr())
+                        return;
+                    VERIFY(promises.first()->is_resolved());
+                    result_promise->resolve(MUST(promises.first()->await()));
+                })
+                .when_rejected([promises, result_promise = result_promise->make_weak_ptr<ResultPromise>()](auto&& error) {
+                    if (!result_promise.ptr())
+                        return;
+                    for (auto& promise : promises) {
+                        if (promise->is_resolved()) {
+                            result_promise->resolve(MUST(promise->await()));
+                            return;
+                        }
+                    }
+                    result_promise->reject(move(error));
+                }));
+        return result_promise;
+    }
+
     NonnullRefPtr<Core::Promise<NonnullRefPtr<LookupResult const>>> lookup(ByteString name, Messages::Class class_ = Messages::Class::IN, LookupOptions options = LookupOptions::default_())
     {
         return lookup(move(name), class_, { Messages::ResourceType::A, Messages::ResourceType::AAAA }, options);
@@ -449,7 +480,9 @@ public:
         }
 
         Messages::Message query;
-        if (options.repeating_lookup) {
+        if (cached_result_id.has_value()) {
+            query.header.id = cached_result_id.value();
+        } else if (options.repeating_lookup) {
             query.header.id = options.repeating_lookup->id;
             options.repeating_lookup->times_repeated++;
         } else {
diff --git a/Utilities/dns.cpp b/Utilities/dns.cpp
index db274613479..74ce2314282 100644
--- a/Utilities/dns.cpp
+++ b/Utilities/dns.cpp
@@ -14,7 +14,7 @@
 ErrorOr<int> serenity_main(Main::Arguments arguments)
 {
     struct Request {
-        Vector<DNS::Messages::ResourceType> types;
+        Vector<Vector<DNS::Messages::ResourceType>> types;
         ByteString name;
     };
     Vector<Request> requests;
@@ -40,7 +40,7 @@ ErrorOr<int> serenity_main(Main::Arguments arguments)
                 return Error::from_string_literal("Invalid record/name format");
 
             if (parts.size() == 1) {
-                current_request.types.append(DNS::Messages::ResourceType::ANY);
+                current_request.types.append({ DNS::Messages::ResourceType::ANY });
                 current_request.name = parts[0];
             } else {
                 auto rr_parts = parts[0].split_view(',');
@@ -49,7 +49,7 @@ ErrorOr<int> serenity_main(Main::Arguments arguments)
                     auto type = DNS::Messages::resource_type_from_string(rr_name.to_uppercase());
                     if (!type.has_value())
                         return Error::from_string_literal("Invalid resource type");
-                    current_request.types.append(type.value());
+                    current_request.types.append({ type.value() });
                 }
                 current_request.name = parts[1];
             }
@@ -105,33 +105,34 @@ ErrorOr<int> serenity_main(Main::Arguments arguments)
     TRY(resolver.when_socket_ready()->await());
 
     size_t pending_requests = requests.size();
+    Vector<NonnullRefPtr<Core::Promise<NonnullRefPtr<DNS::LookupResult const>>>> promises;
     for (auto& request : requests) {
-        resolver.lookup(request.name, DNS::Messages::Class::IN, request.types, { .validate_dnssec_locally = dnssec })
-            ->when_resolved([&](auto& result) {
-                outln("Resolved {}:", request.name);
-                HashTable<DNS::Messages::ResourceType> types;
-                auto recs = result->records();
-                for (auto& record : recs)
-                    types.set(record.type);
+        promises.append(resolver.lookup(request.name, DNS::Messages::Class::IN, request.types, { .validate_dnssec_locally = dnssec })
+                ->when_resolved([&](auto& result) {
+                    outln("Resolved {}:", request.name);
+                    HashTable<DNS::Messages::ResourceType> types;
+                    auto recs = result->records();
+                    for (auto& record : recs)
+                        types.set(record.type);
 
-                for (auto& type : types) {
-                    outln("  - {} IN {}:", request.name, DNS::Messages::to_string(type));
-                    for (auto& record : recs) {
-                        if (type != record.type)
-                            continue;
+                    for (auto& type : types) {
+                        outln("  - {} IN {}:", request.name, DNS::Messages::to_string(type));
+                        for (auto& record : recs) {
+                            if (type != record.type)
+                                continue;
 
-                        outln("    - {}", record.to_string());
+                            outln("    - {}", record.to_string());
+                        }
                     }
-                }
 
-                if (--pending_requests == 0)
-                    loop.quit(0);
-            })
-            .when_rejected([&](auto& error) {
-                outln("Failed to resolve {} IN {}: {}", request.name, DNS::Messages::to_string(request.types.first()), error);
-                if (--pending_requests == 0)
-                    loop.quit(1);
-            });
+                    if (--pending_requests == 0)
+                        loop.quit(0);
+                })
+                .when_rejected([&](auto& error) {
+                    outln("Failed to resolve {} IN {}: {}", request.name, DNS::Messages::to_string(request.types.first().first()), error);
+                    if (--pending_requests == 0)
+                        loop.quit(1);
+                }));
     }
 
     return loop.exec();