YARN-8410. Fixed a bug in A record lookup by CNAME record.

            Contributed by Shane Kumpf

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-registry/src/main/java/org/apache/hadoop/registry/server/dns/RegistryDNS.java

@@ -1126,19 +1126,38 @@ public class RegistryDNS extends AbstractService implements DNSOperations,
    */
   private byte remoteLookup(Message response, Name name, int type,
       int iterations) {
+    // If retrieving the root zone, query for NS record type
+    if (name.toString().equals(".")) {
+      type = Type.NS;
+    }
+
+    // Always add any CNAMEs to the response first
+    if (type != Type.CNAME) {
+      Record[] cnameAnswers = getRecords(name, Type.CNAME);
+      if (cnameAnswers != null) {
+        for (Record cnameR : cnameAnswers) {
+          if (!response.findRecord(cnameR)) {
+            response.addRecord(cnameR, Section.ANSWER);
+          }
+        }
+      }
+    }
+
     // Forward lookup to primary DNS servers
     Record[] answers = getRecords(name, type);
     try {
       for (Record r : answers) {
-        if (r.getType() == Type.SOA) {
-          response.addRecord(r, Section.AUTHORITY);
-        } else {
-          response.addRecord(r, Section.ANSWER);
+        if (!response.findRecord(r)) {
+          if (r.getType() == Type.SOA) {
+            response.addRecord(r, Section.AUTHORITY);
+          } else {
+            response.addRecord(r, Section.ANSWER);
+          }
         }
         if (r.getType() == Type.CNAME) {
           Name cname = ((CNAMERecord) r).getAlias();
           if (iterations < 6) {
-            remoteLookup(response, cname, Type.CNAME, iterations + 1);
+            remoteLookup(response, cname, type, iterations + 1);
           }
         }
       }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-registry/src/test/java/org/apache/hadoop/registry/server/dns/TestRegistryDNS.java

@@ -410,7 +410,7 @@ public class TestRegistryDNS extends Assert {
     return recs;
   }
 
-  Record[] assertDNSQueryNotNull(String lookup, int type)
+  Record[] assertDNSQueryNotNull(String lookup, int type, int answerCount)
       throws IOException {
     Name name = Name.fromString(lookup);
     Record question = Record.newRecord(name, type, DClass.IN);
@@ -424,7 +424,7 @@ public class TestRegistryDNS extends Assert {
     assertEquals("Questions do not match", query.getQuestion(),
         response.getQuestion());
     Record[] recs = response.getSectionArray(Section.ANSWER);
-    assertEquals(1, recs.length);
+    assertEquals(answerCount, recs.length);
     assertEquals(recs[0].getType(), type);
     return recs;
   }
@@ -656,7 +656,24 @@ public class TestRegistryDNS extends Assert {
 
     // start assessing whether correct records are available
     Record[] recs =
-        assertDNSQueryNotNull("mail.yahoo.com.", Type.CNAME);
+        assertDNSQueryNotNull("mail.yahoo.com.", Type.CNAME, 1);
+  }
+
+  @Test
+  public void testRootLookup() throws Exception {
+    setRegistryDNS(new RegistryDNS("TestRegistry"));
+    Configuration conf = new Configuration();
+    conf.set(RegistryConstants.KEY_DNS_DOMAIN, "dev.test");
+    conf.set(RegistryConstants.KEY_DNS_ZONE_SUBNET, "172.17.0");
+    conf.setTimeDuration(RegistryConstants.KEY_DNS_TTL, 30L, TimeUnit.SECONDS);
+    conf.set(RegistryConstants.KEY_DNS_ZONES_DIR,
+        getClass().getResource("/").getFile());
+    getRegistryDNS().setDomainName(conf);
+    getRegistryDNS().initializeZones(conf);
+
+    // start assessing whether correct records are available
+    Record[] recs =
+        assertDNSQueryNotNull(".", Type.NS, 13);
   }
 
   @Test