From 6697ef495521ffd80386b6ccf162db286b36375f Mon Sep 17 00:00:00 2001
From: Christian Ammer <chrie.ammer@gmail.com>
Date: Sun, 11 Nov 2018 23:25:38 +0100
Subject: [PATCH 02/10] Remaining queries counter fix, additional unit tests
 for `ares_getaddrinfo` (#233)

Remaining queries counter fix, added tests (ParallelLookups,
SearchDomains, SearchDomainsServFailOnAAAA).  Removed unnecessary
if and commented code in test.

Fix By: Christian Ammer (@ChristianAmmer)
---
 ares_getaddrinfo.c        |  39 ++++++++-------
 test/ares-test-ai.h       |   7 +++
 test/ares-test-mock-ai.cc | 125 +++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 153 insertions(+), 18 deletions(-)

diff --git a/ares_getaddrinfo.c b/ares_getaddrinfo.c
index be936ff..36f29b5 100644
--- a/ares_getaddrinfo.c
+++ b/ares_getaddrinfo.c
@@ -122,7 +122,7 @@ void ares_getaddrinfo(ares_channel channel,
   hquery->arg = arg;
   hquery->timeouts = 0;
   hquery->next_domain = 0;
-  hquery->remaining = ai_family == AF_UNSPEC ? 2 : 1;
+  hquery->remaining = 0;
 
   /* Host file lookup */
   if (file_lookup(hquery->name, ai_family, &hquery->ai) == ARES_SUCCESS) {
@@ -291,9 +291,11 @@ static void next_dns_lookup(struct host_query *hquery) {
   if (s) {
     if (hquery->ai_family == AF_INET || hquery->ai_family == AF_UNSPEC) {
       ares_query(hquery->channel, s, C_IN, T_A, host_callback, hquery);
+      hquery->remaining++;
     }
     if (hquery->ai_family == AF_INET6 || hquery->ai_family == AF_UNSPEC) {
       ares_query(hquery->channel, s, C_IN, T_AAAA, host_callback, hquery);
+      hquery->remaining++;
     }
     if (is_s_allocated) {
       ares_free(s);
@@ -306,9 +308,7 @@ static void next_dns_lookup(struct host_query *hquery) {
 }
 
 static void end_hquery(struct host_query *hquery, int status) {
-  if (hquery->ai) {
-    hquery->callback(hquery->arg, status, hquery->ai);
-  }
+  hquery->callback(hquery->arg, status, hquery->ai);
   ares_free(hquery->name);
   ares_free(hquery);
 }
@@ -319,11 +319,14 @@ static void host_callback(void *arg, int status, int timeouts,
   ares_channel channel = hquery->channel;
   struct hostent *host = NULL;
   int qtype;
+  int qtypestatus;
   hquery->timeouts += timeouts;
 
+  hquery->remaining--;
+
   if (status == ARES_SUCCESS) {
-    status = ares__parse_qtype_reply(abuf, alen, &qtype);
-    if (status == ARES_SUCCESS && qtype == T_A) {
+    qtypestatus = ares__parse_qtype_reply(abuf, alen, &qtype);
+    if (qtypestatus == ARES_SUCCESS && qtype == T_A) {
       /* Can ares_parse_a_reply be unsuccessful (after parse_qtype) */
       ares_parse_a_reply(abuf, alen, &host, NULL, NULL);
       if (host && channel->nsort) {
@@ -331,11 +334,8 @@ static void host_callback(void *arg, int status, int timeouts,
       }
       add_to_addrinfo(&hquery->ai, host);
       ares_free_hostent(host);
-      if (!--hquery->remaining) {
-	end_hquery(hquery, ARES_SUCCESS);
-      }
     }
-    else if (status == ARES_SUCCESS && qtype == T_AAAA) {
+    else if (qtypestatus == ARES_SUCCESS && qtype == T_AAAA) {
       /* Can ares_parse_a_reply be unsuccessful (after parse_qtype) */
       ares_parse_aaaa_reply(abuf, alen, &host, NULL, NULL);
       if (host && channel->nsort) {
@@ -343,18 +343,23 @@ static void host_callback(void *arg, int status, int timeouts,
       }
       add_to_addrinfo(&hquery->ai, host);
       ares_free_hostent(host);
-      if (!--hquery->remaining) {
-        end_hquery(hquery, ARES_SUCCESS);
-      }
+    }
+  }
+
+  if (!hquery->remaining) {
+    if (hquery->ai) {
+      // at least one query ended with ARES_SUCCESS
+      end_hquery(hquery, ARES_SUCCESS);
+    }
+    else if (status == ARES_ENOTFOUND) {
+      next_dns_lookup(hquery);
     }
     else {
-      assert(!hquery->ai);
       end_hquery(hquery, status);
     }
   }
-  else {
-    next_dns_lookup(hquery);
-  }
+
+  // at this point we keep on waiting for the next query to finish
 }
 
 static void sort_addresses(struct hostent *host,
diff --git a/test/ares-test-ai.h b/test/ares-test-ai.h
index e4c4403..a7a6a73 100644
--- a/test/ares-test-ai.h
+++ b/test/ares-test-ai.h
@@ -16,6 +16,13 @@ class MockChannelTestAI
   MockChannelTestAI() : MockChannelOptsTest(1, GetParam().first, GetParam().second, nullptr, 0) {}
 };
 
+class MockUDPChannelTestAI
+    : public MockChannelOptsTest,
+      public ::testing::WithParamInterface<int> {
+ public:
+  MockUDPChannelTestAI() : MockChannelOptsTest(1, GetParam(), false, nullptr, 0) {}
+};
+
 // Structure that describes the result of an ares_addr_callback invocation.
 struct AIResult {
   // Whether the callback has been invoked.
diff --git a/test/ares-test-mock-ai.cc b/test/ares-test-mock-ai.cc
index 8ba1611..28a01be 100644
--- a/test/ares-test-mock-ai.cc
+++ b/test/ares-test-mock-ai.cc
@@ -52,6 +52,50 @@ MATCHER_P(IncludesV6Address, address, "") {
   return false;
 }
 
+// UDP only so mock server doesn't get confused by concatenated requests
+TEST_P(MockUDPChannelTestAI, ParallelLookups) {
+  DNSPacket rsp1;
+  rsp1.set_response().set_aa()
+    .add_question(new DNSQuestion("www.google.com", ns_t_a))
+    .add_answer(new DNSARR("www.google.com", 100, {2, 3, 4, 5}));
+  ON_CALL(server_, OnRequest("www.google.com", ns_t_a))
+    .WillByDefault(SetReply(&server_, &rsp1));
+  DNSPacket rsp2;
+  rsp2.set_response().set_aa()
+    .add_question(new DNSQuestion("www.example.com", ns_t_a))
+    .add_answer(new DNSARR("www.example.com", 100, {1, 2, 3, 4}));
+  ON_CALL(server_, OnRequest("www.example.com", ns_t_a))
+    .WillByDefault(SetReply(&server_, &rsp2));
+
+  struct ares_addrinfo hints = {};
+  hints.ai_family = AF_INET;
+  AIResult result1;
+  ares_getaddrinfo(channel_, "www.google.com.", NULL, &hints, AICallback, &result1);
+  AIResult result2;
+  ares_getaddrinfo(channel_, "www.example.com.", NULL, &hints, AICallback, &result2);
+  AIResult result3;
+  ares_getaddrinfo(channel_, "www.google.com.", NULL, &hints, AICallback, &result3);
+  Process();
+
+  EXPECT_TRUE(result1.done);
+  EXPECT_EQ(result1.status, ARES_SUCCESS);
+  EXPECT_THAT(result1.airesult, IncludesNumAddresses(1));
+  EXPECT_THAT(result1.airesult, IncludesV4Address("2.3.4.5"));
+  ares_freeaddrinfo(result1.airesult);
+
+  EXPECT_TRUE(result2.done);
+  EXPECT_EQ(result2.status, ARES_SUCCESS);
+  EXPECT_THAT(result2.airesult, IncludesNumAddresses(1));
+  EXPECT_THAT(result2.airesult, IncludesV4Address("1.2.3.4"));
+  ares_freeaddrinfo(result2.airesult);
+
+  EXPECT_TRUE(result3.done);
+  EXPECT_EQ(result3.status, ARES_SUCCESS);
+  EXPECT_THAT(result3.airesult, IncludesNumAddresses(1));
+  EXPECT_THAT(result3.airesult, IncludesV4Address("2.3.4.5"));
+  ares_freeaddrinfo(result3.airesult);
+}
+
 TEST_P(MockChannelTestAI, FamilyV6) {
   DNSPacket rsp6;
   rsp6.set_response().set_aa()
@@ -146,9 +190,88 @@ TEST_P(MockChannelTestAI, FamilyUnspecified) {
   ares_freeaddrinfo(result.airesult);
 }
 
-INSTANTIATE_TEST_CASE_P(AddressFamilies, MockChannelTestAI,
+TEST_P(MockChannelTestAI, SearchDomains) {
+  DNSPacket nofirst;
+  nofirst.set_response().set_aa().set_rcode(ns_r_nxdomain)
+    .add_question(new DNSQuestion("www.first.com", ns_t_a));
+  ON_CALL(server_, OnRequest("www.first.com", ns_t_a))
+    .WillByDefault(SetReply(&server_, &nofirst));
+  DNSPacket nosecond;
+  nosecond.set_response().set_aa().set_rcode(ns_r_nxdomain)
+    .add_question(new DNSQuestion("www.second.org", ns_t_a));
+  ON_CALL(server_, OnRequest("www.second.org", ns_t_a))
+    .WillByDefault(SetReply(&server_, &nosecond));
+  DNSPacket yesthird;
+  yesthird.set_response().set_aa()
+    .add_question(new DNSQuestion("www.third.gov", ns_t_a))
+    .add_answer(new DNSARR("www.third.gov", 0x0200, {2, 3, 4, 5}));
+  ON_CALL(server_, OnRequest("www.third.gov", ns_t_a))
+    .WillByDefault(SetReply(&server_, &yesthird));
+
+  AIResult result;
+  struct ares_addrinfo hints = {};
+  hints.ai_family = AF_INET;
+  ares_getaddrinfo(channel_, "www", NULL, &hints, AICallback, &result);
+  Process();
+  EXPECT_TRUE(result.done);
+  EXPECT_EQ(result.status, ARES_SUCCESS);
+  EXPECT_THAT(result.airesult, IncludesNumAddresses(1));
+  EXPECT_THAT(result.airesult, IncludesV4Address("2.3.4.5"));
+  ares_freeaddrinfo(result.airesult);
+}
+
+TEST_P(MockChannelTestAI, SearchDomainsServFailOnAAAA) {
+  DNSPacket nofirst;
+  nofirst.set_response().set_aa().set_rcode(ns_r_nxdomain)
+    .add_question(new DNSQuestion("www.first.com", ns_t_aaaa));
+  ON_CALL(server_, OnRequest("www.first.com", ns_t_aaaa))
+    .WillByDefault(SetReply(&server_, &nofirst));
+  DNSPacket nofirst4;
+  nofirst4.set_response().set_aa().set_rcode(ns_r_nxdomain)
+    .add_question(new DNSQuestion("www.first.com", ns_t_a));
+  ON_CALL(server_, OnRequest("www.first.com", ns_t_a))
+    .WillByDefault(SetReply(&server_, &nofirst4));
+  
+  DNSPacket nosecond;
+  nosecond.set_response().set_aa().set_rcode(ns_r_nxdomain)
+    .add_question(new DNSQuestion("www.second.org", ns_t_aaaa));
+  ON_CALL(server_, OnRequest("www.second.org", ns_t_aaaa))
+    .WillByDefault(SetReply(&server_, &nosecond));
+  DNSPacket yessecond4;
+  yessecond4.set_response().set_aa()
+    .add_question(new DNSQuestion("www.second.org", ns_t_a))
+    .add_answer(new DNSARR("www.second.org", 0x0200, {2, 3, 4, 5}));
+  ON_CALL(server_, OnRequest("www.second.org", ns_t_a))
+    .WillByDefault(SetReply(&server_, &yessecond4));
+  
+  DNSPacket failthird;
+  failthird.set_response().set_aa().set_rcode(ns_r_servfail)
+    .add_question(new DNSQuestion("www.third.gov", ns_t_aaaa));
+  ON_CALL(server_, OnRequest("www.third.gov", ns_t_aaaa))
+    .WillByDefault(SetReply(&server_, &failthird));
+  DNSPacket failthird4;
+  failthird4.set_response().set_aa().set_rcode(ns_r_servfail)
+    .add_question(new DNSQuestion("www.third.gov", ns_t_a));
+  ON_CALL(server_, OnRequest("www.third.gov", ns_t_a))
+    .WillByDefault(SetReply(&server_, &failthird4));
+  
+  AIResult result;
+  struct ares_addrinfo hints = {};
+  hints.ai_family = AF_UNSPEC;
+  ares_getaddrinfo(channel_, "www", NULL, &hints, AICallback, &result);
+  Process();
+  EXPECT_TRUE(result.done);
+  EXPECT_EQ(result.status, ARES_SUCCESS);
+  EXPECT_THAT(result.airesult, IncludesNumAddresses(1));
+  EXPECT_THAT(result.airesult, IncludesV4Address("2.3.4.5"));
+  ares_freeaddrinfo(result.airesult);
+}
+
+INSTANTIATE_TEST_CASE_P(AddressFamiliesAI, MockChannelTestAI,
                         ::testing::Values(std::make_pair<int, bool>(AF_INET, false)));
 
+INSTANTIATE_TEST_CASE_P(AddressFamiliesAI, MockUDPChannelTestAI,
+                        ::testing::ValuesIn(ares::test::families));
 
 }  // namespace test
 }  // namespace ares
-- 
2.16.4

