Support returning multiple results for dns parsing (#265)

Author: fantasy-peak

trantor/net/Resolver.h

@@ -22,6 +22,8 @@ class TRANTOR_EXPORT Resolver
 {
   public:
     using Callback = std::function<void(const trantor::InetAddress&)>;
+    using ResolverResultsCallback =
+        std::function<void(const std::vector<trantor::InetAddress>&)>;
 
     /**
      * @brief Create a new DNS resolver.
@@ -42,6 +44,15 @@ class TRANTOR_EXPORT Resolver
     virtual void resolve(const std::string& hostname,
                          const Callback& callback) = 0;
 
+    /**
+     * @brief Resolve an address array asynchronously.
+     *
+     * @param hostname
+     * @param callback
+     */
+    virtual void resolve(const std::string& hostname,
+                         const ResolverResultsCallback& callback) = 0;
+
     virtual ~Resolver()
     {
     }

trantor/net/inner/AresResolver.cc

@@ -110,13 +110,14 @@ AresResolver::~AresResolver()
 }
 
 void AresResolver::resolveInLoop(const std::string& hostname,
-                                 const Callback& cb)
+                                 const ResolverResultsCallback& cb)
 {
     loop_->assertInLoopThread();
 #ifdef _WIN32
     if (hostname == "localhost")
     {
-        const static trantor::InetAddress localhost_{"127.0.0.1", 0};
+        const static std::vector<trantor::InetAddress> localhost_{
+            trantor::InetAddress{"127.0.0.1", 0}};
         cb(localhost_);
         return;
     }
@@ -167,25 +168,39 @@ void AresResolver::onTimer()
 void AresResolver::onQueryResult(int status,
                                  struct hostent* result,
                                  const std::string& hostname,
-                                 const Callback& callback)
+                                 const ResolverResultsCallback& callback)
 {
     LOG_TRACE << "onQueryResult " << status;
-    struct sockaddr_in addr;
-    memset(&addr, 0, sizeof addr);
-    addr.sin_family = AF_INET;
-    addr.sin_port = 0;
+    std::vector<trantor::InetAddress> inets;
     if (result)
     {
-        addr.sin_addr = *reinterpret_cast<in_addr*>(result->h_addr);
+        auto pptr = (struct in_addr**)result->h_addr_list;
+        for (; *pptr != nullptr; pptr++)
+        {
+            struct sockaddr_in addr;
+            memset(&addr, 0, sizeof addr);
+            addr.sin_family = AF_INET;
+            addr.sin_port = 0;
+            addr.sin_addr = *reinterpret_cast<in_addr*>(*pptr);
+            inets.emplace_back(trantor::InetAddress{addr});
+        }
+    }
+    if (inets.empty())
+    {
+        struct sockaddr_in addr;
+        memset(&addr, 0, sizeof addr);
+        addr.sin_family = AF_INET;
+        addr.sin_port = 0;
+        InetAddress inet(addr);
+        inets.emplace_back(std::move(inet));
     }
-    InetAddress inet(addr);
     {
         std::lock_guard<std::mutex> lock(globalMutex());
         auto& addrItem = globalCache()[hostname];
-        addrItem.first = addr.sin_addr;
+        addrItem.first = inets;
         addrItem.second = trantor::Date::date();
     }
-    callback(inet);
+    callback(inets);
 }
 
 void AresResolver::onSockCreate(int sockfd, int type)

trantor/net/inner/AresResolver.h

@@ -43,12 +43,7 @@ class AresResolver : public Resolver,
                 if (timeout_ == 0 ||
                     cachedAddr.second.after(timeout_) > trantor::Date::date())
                 {
-                    struct sockaddr_in addr;
-                    memset(&addr, 0, sizeof addr);
-                    addr.sin_family = AF_INET;
-                    addr.sin_port = 0;
-                    addr.sin_addr = cachedAddr.first;
-                    inet = InetAddress(addr);
+                    inet = (cachedAddr.first)[0];
                     cached = true;
                 }
             }
@@ -59,6 +54,49 @@ class AresResolver : public Resolver,
             return;
         }
         if (loop_->isInLoopThread())
+        {
+            resolveInLoop(hostname,
+                          [cb](const std::vector<trantor::InetAddress>& inets) {
+                              cb(inets[0]);
+                          });
+        }
+        else
+        {
+            loop_->queueInLoop([thisPtr = shared_from_this(), hostname, cb]() {
+                thisPtr->resolveInLoop(
+                    hostname,
+                    [cb](const std::vector<trantor::InetAddress>& inets) {
+                        cb(inets[0]);
+                    });
+            });
+        }
+    }
+
+    virtual void resolve(const std::string& hostname,
+                         const ResolverResultsCallback& cb) override
+    {
+        bool cached = false;
+        std::vector<trantor::InetAddress> inets;
+        {
+            std::lock_guard<std::mutex> lock(globalMutex());
+            auto iter = globalCache().find(hostname);
+            if (iter != globalCache().end())
+            {
+                auto& cachedAddr = iter->second;
+                if (timeout_ == 0 ||
+                    cachedAddr.second.after(timeout_) > trantor::Date::date())
+                {
+                    inets = cachedAddr.first;
+                    cached = true;
+                }
+            }
+        }
+        if (cached)
+        {
+            cb(inets);
+            return;
+        }
+        if (loop_->isInLoopThread())
         {
             resolveInLoop(hostname, cb);
         }
@@ -74,29 +112,32 @@ class AresResolver : public Resolver,
     struct QueryData
     {
         AresResolver* owner_;
-        Callback callback_;
+        ResolverResultsCallback callback_;
         std::string hostname_;
         QueryData(AresResolver* o,
-                  const Callback& cb,
+                  const ResolverResultsCallback& cb,
                   const std::string& hostname)
             : owner_(o), callback_(cb), hostname_(hostname)
         {
         }
     };
-    void resolveInLoop(const std::string& hostname, const Callback& cb);
+    void resolveInLoop(const std::string& hostname,
+                       const ResolverResultsCallback& cb);
     void init();
     trantor::EventLoop* loop_;
     std::shared_ptr<bool> loopValid_;
     ares_channel ctx_{nullptr};
     bool timerActive_{false};
     using ChannelList = std::map<int, std::unique_ptr<trantor::Channel>>;
     ChannelList channels_;
-    static std::unordered_map<std::string,
-                              std::pair<struct in_addr, trantor::Date>>&
+    static std::unordered_map<
+        std::string,
+        std::pair<std::vector<InetAddress>, trantor::Date>>&
     globalCache()
     {
-        static std::unordered_map<std::string,
-                                  std::pair<struct in_addr, trantor::Date>>
+        static std::unordered_map<
+            std::string,
+            std::pair<std::vector<InetAddress>, trantor::Date>>
             dnsCache;
         return dnsCache;
     }
@@ -118,7 +159,7 @@ class AresResolver : public Resolver,
     void onQueryResult(int status,
                        struct hostent* result,
                        const std::string& hostname,
-                       const Callback& callback);
+                       const ResolverResultsCallback& callback);
     void onSockCreate(int sockfd, int type);
     void onSockStateChange(int sockfd, bool read, bool write);
 

trantor/net/inner/NormalResolver.h

@@ -23,6 +23,13 @@ class NormalResolver : public Resolver,
   public:
     virtual void resolve(const std::string& hostname,
                          const Callback& callback) override;
+    virtual void resolve(const std::string& hostname,
+                         const ResolverResultsCallback& callback) override
+    {
+        resolve(hostname, [callback](const trantor::InetAddress& inet) {
+            callback(std::vector<trantor::InetAddress>{inet});
+        });
+    }
     explicit NormalResolver(size_t timeout)
         : timeout_(timeout), resolveBuffer_(kResolveBufferLength)
     {

trantor/tests/DnsTest.cc

@@ -9,6 +9,15 @@ void dns(const std::shared_ptr<trantor::Resolver> &resolver)
         std::cout << "baidu:" << addr.toIp() << " " << interval / 1000 << "ms"
                   << std::endl;
     });
+    resolver->resolve("www.baidu.com",
+                      [now](const std::vector<trantor::InetAddress> &addrs) {
+                          auto interval =
+                              trantor::Date::now().microSecondsSinceEpoch() -
+                              now.microSecondsSinceEpoch();
+                          for (auto &addr : addrs)
+                              std::cout << "baidu:" << addr.toIp() << " "
+                                        << interval / 1000 << "ms" << std::endl;
+                      });
     resolver->resolve("www.google.com",
                       [now](const trantor::InetAddress &addr) {
                           auto interval =