Expand wildcard names in Zone lookups

Closes #169

src/main/java/org/xbill/DNS/RRset.java

@@ -8,6 +8,8 @@
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Objects;
+import lombok.EqualsAndHashCode;
 
 /**
  * A set of Records with the same name, type, and class. Also included are all RRSIG records signing
@@ -17,6 +19,7 @@
  * @see RRSIGRecord
  * @author Brian Wellington
  */
+@EqualsAndHashCode(of = {"rrs", "sigs"})
 public class RRset implements Serializable {
   private final ArrayList<Record> rrs;
   private final ArrayList<RRSIGRecord> sigs;
@@ -35,6 +38,19 @@ public RRset(Record record) {
     addRR(record);
   }
 
+  /**
+   * Creates an RRset and sets its contents to the specified record(s)
+   *
+   * @param records The records to add to the set. See {@link #addRR(Record)} for restrictions.
+   */
+  public RRset(Record... records) {
+    this();
+    Objects.requireNonNull(records);
+    for (Record r : records) {
+      addRR(r);
+    }
+  }
+
   /** Creates an RRset with the contents of an existing RRset */
   public RRset(RRset rrset) {
     rrs = new ArrayList<>(rrset.rrs);
@@ -47,6 +63,9 @@ public RRset(RRset rrset) {
    * Adds a signature to this RRset. If the TTL of the added signature is not the same as existing
    * records in the RRset, all records are set to the lowest TTL of either the added record or the
    * existing records.
+   *
+   * @throws IllegalArgumentException if the RRset already contains records and the signature to add
+   *     does not match.
    */
   public void addRR(RRSIGRecord r) {
     addRR(r, sigs);
@@ -56,6 +75,9 @@ public void addRR(RRSIGRecord r) {
    * Adds a Record to this RRset. If the TTL of the added record is not the same as existing records
    * in the RRset, all records are set to the lowest TTL of either the added record or the existing
    * records.
+   *
+   * @throws IllegalArgumentException if the RRset already contains records and the record to add
+   *     does not match.
    */
   public void addRR(Record r) {
     if (r instanceof RRSIGRecord) {

src/main/java/org/xbill/DNS/Zone.java

@@ -337,25 +337,18 @@ private synchronized void removeRRset(Name name, int type) {
   }
 
   private synchronized SetResponse lookup(Name name, int type) {
-    int labels;
-    int olabels;
-    int tlabels;
-    RRset rrset;
-    Name tname;
-    Object types;
-    SetResponse sr;
-
     if (!name.subdomain(origin)) {
       return SetResponse.ofType(SetResponse.NXDOMAIN);
     }
 
-    labels = name.labels();
-    olabels = origin.labels();
+    int labels = name.labels();
+    int olabels = origin.labels();
 
-    for (tlabels = olabels; tlabels <= labels; tlabels++) {
+    for (int tlabels = olabels; tlabels <= labels; tlabels++) {
       boolean isOrigin = tlabels == olabels;
       boolean isExact = tlabels == labels;
 
+      Name tname;
       if (isOrigin) {
         tname = origin;
       } else if (isExact) {
@@ -364,7 +357,7 @@ private synchronized SetResponse lookup(Name name, int type) {
         tname = new Name(name, labels - tlabels);
       }
 
-      types = exactName(tname);
+      Object types = exactName(tname);
       if (types == null) {
         continue;
       }
@@ -379,9 +372,8 @@ private synchronized SetResponse lookup(Name name, int type) {
 
       /* If this is an ANY lookup, return everything. */
       if (isExact && type == Type.ANY) {
-        sr = new SetResponse(SetResponse.SUCCESSFUL);
-        RRset[] sets = allRRsets(types);
-        for (RRset set : sets) {
+        SetResponse sr = new SetResponse(SetResponse.SUCCESSFUL);
+        for (RRset set : allRRsets(types)) {
           sr.addRRset(set);
         }
         return sr;
@@ -392,18 +384,16 @@ private synchronized SetResponse lookup(Name name, int type) {
        * Otherwise, look for a DNAME.
        */
       if (isExact) {
-        rrset = oneRRset(types, type);
+        RRset rrset = oneRRset(types, type);
         if (rrset != null) {
-          sr = new SetResponse(SetResponse.SUCCESSFUL);
-          sr.addRRset(rrset);
-          return sr;
+          return new SetResponse(SetResponse.SUCCESSFUL, rrset);
         }
         rrset = oneRRset(types, Type.CNAME);
         if (rrset != null) {
           return new SetResponse(SetResponse.CNAME, rrset);
         }
       } else {
-        rrset = oneRRset(types, Type.DNAME);
+        RRset rrset = oneRRset(types, Type.DNAME);
         if (rrset != null) {
           return new SetResponse(SetResponse.DNAME, rrset);
         }
@@ -417,25 +407,41 @@ private synchronized SetResponse lookup(Name name, int type) {
 
     if (hasWild) {
       for (int i = 0; i < labels - olabels; i++) {
-        tname = name.wild(i + 1);
-
-        types = exactName(tname);
+        Name tname = name.wild(i + 1);
+        Object types = exactName(tname);
         if (types == null) {
           continue;
         }
 
-        rrset = oneRRset(types, type);
-        if (rrset != null) {
-          sr = new SetResponse(SetResponse.SUCCESSFUL);
-          sr.addRRset(rrset);
+        if (type == Type.ANY) {
+          SetResponse sr = new SetResponse(SetResponse.SUCCESSFUL);
+          for (RRset set : allRRsets(types)) {
+            sr.addRRset(expandSet(set, name));
+          }
           return sr;
+        } else {
+          RRset rrset = oneRRset(types, type);
+          if (rrset != null) {
+            return new SetResponse(SetResponse.SUCCESSFUL, expandSet(rrset, name));
+          }
         }
       }
     }
 
     return SetResponse.ofType(SetResponse.NXDOMAIN);
   }
 
+  private RRset expandSet(RRset set, Name tname) {
+    RRset expandedSet = new RRset();
+    for (Record r : set.rrs()) {
+      expandedSet.addRR(r.withName(tname));
+    }
+    for (RRSIGRecord r : set.sigs()) {
+      expandedSet.addRR(r.withName(tname));
+    }
+    return expandedSet;
+  }
+
   /**
    * Looks up Records in the Zone. This follows CNAMEs and wildcards.
    *

src/test/java/org/xbill/DNS/RRsetTest.java

@@ -44,7 +44,10 @@
 import java.net.InetAddress;
 import java.net.UnknownHostException;
 import java.time.Instant;
+import java.util.Collections;
 import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
@@ -218,6 +221,31 @@ void ctor_1arg() {
     assertEquals(2, sigs.size());
   }
 
+  @Test
+  void ctor_vararg() {
+    RRset set = new RRset(m_a1, m_a2);
+    assertEquals(Stream.of(m_a1, m_a2).collect(Collectors.toList()), set.rrs());
+    assertEquals(Collections.emptyList(), set.sigs());
+  }
+
+  @Test
+  void ctor_vararg_sig() {
+    RRset set = new RRset(m_a1, m_a2, m_s1);
+    assertEquals(Stream.of(m_a1, m_a2).collect(Collectors.toList()), set.rrs());
+    assertEquals(Collections.singletonList(m_s1), set.sigs());
+  }
+
+  @Test
+  void ctor_vararg_mismatch() {
+    TXTRecord txt = new TXTRecord(m_name, DClass.IN, 3600, "test");
+    assertThrows(IllegalArgumentException.class, () -> new RRset(m_a1, m_a2, txt));
+  }
+
+  @Test
+  void ctor_vararg_null() {
+    assertThrows(NullPointerException.class, () -> new RRset((Record[]) null));
+  }
+
   @Test
   void test_toString() {
     m_rs.addRR(m_a1);

src/test/java/org/xbill/DNS/ZoneTest.java

@@ -0,0 +1,120 @@
+// SPDX-License-Identifier: BSD-2-Clause
+package org.xbill.DNS;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.io.IOException;
+import java.net.InetAddress;
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import org.junit.jupiter.api.Test;
+
+class ZoneTest {
+  private static final ARecord A_TEST;
+  private static final AAAARecord AAAA_1_TEST;
+  private static final AAAARecord AAAA_2_TEST;
+  private static final ARecord A_WILD;
+  private static final TXTRecord TXT_WILD;
+  private static final Zone ZONE;
+
+  static {
+    try {
+      Name nameZone = new Name("example.");
+      InetAddress localhost4 = InetAddress.getByName("127.0.0.1");
+      InetAddress localhost6a = InetAddress.getByName("::1");
+      InetAddress localhost6b = InetAddress.getByName("::2");
+      A_TEST = new ARecord(new Name("test", nameZone), DClass.IN, 3600, localhost4);
+      AAAA_1_TEST = new AAAARecord(new Name("test", nameZone), DClass.IN, 3600, localhost6a);
+      AAAA_2_TEST = new AAAARecord(new Name("test", nameZone), DClass.IN, 3600, localhost6b);
+      A_WILD = new ARecord(new Name("*", nameZone), DClass.IN, 3600, localhost4);
+      TXT_WILD = new TXTRecord(new Name("*", nameZone), DClass.IN, 3600, "sometext");
+
+      Record[] zoneRecords =
+          new Record[] {
+            new SOARecord(
+                nameZone,
+                DClass.IN,
+                3600L,
+                Name.fromConstantString("nameserver."),
+                new Name("hostmaster", nameZone),
+                1,
+                21600L,
+                7200L,
+                2160000L,
+                3600L),
+            new NSRecord(nameZone, DClass.IN, 300L, Name.fromConstantString("nameserver.")),
+            A_TEST,
+            AAAA_1_TEST,
+            AAAA_2_TEST,
+            A_WILD,
+            TXT_WILD,
+          };
+      ZONE = new Zone(nameZone, zoneRecords);
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  @Test
+  void exactNameExistingALookup() {
+    Name testName = Name.fromConstantString("test.example.");
+    SetResponse resp = ZONE.findRecords(testName, Type.A);
+    assertEquals(oneRRset(A_TEST), resp.answers());
+  }
+
+  @Test
+  void exactNameTwoAaaaLookup() {
+    Name testName = Name.fromConstantString("test.example.");
+    SetResponse resp = ZONE.findRecords(testName, Type.AAAA);
+    assertEquals(oneRRset(AAAA_1_TEST, AAAA_2_TEST), resp.answers());
+  }
+
+  @Test
+  void exactNameAnyLookup() {
+    Name testName = Name.fromConstantString("test.example.");
+    SetResponse resp = ZONE.findRecords(testName, Type.ANY);
+    assertTrue(resp.isSuccessful());
+    assertEquals(listOf(new RRset(A_TEST), new RRset(AAAA_1_TEST, AAAA_2_TEST)), resp.answers());
+  }
+
+  @Test
+  void wildNameExistingALookup() {
+    Name testName = Name.fromConstantString("undefined.example.");
+    SetResponse resp = ZONE.findRecords(testName, Type.A);
+    assertEquals(oneRRset(A_WILD.withName(testName)), resp.answers());
+  }
+
+  @Test
+  void wildNameExistingTxtLookup() {
+    Name testName = Name.fromConstantString("undefined.example.");
+    SetResponse resp = ZONE.findRecords(testName, Type.TXT);
+    assertEquals(oneRRset(TXT_WILD.withName(testName)), resp.answers());
+  }
+
+  @Test
+  void wildNameNonExistingMxLookup() {
+    SetResponse resp = ZONE.findRecords(Name.fromConstantString("undefined.example."), Type.MX);
+    assertTrue(resp.isNXDOMAIN());
+  }
+
+  @Test
+  void wildNameAnyLookup() {
+    Name testName = Name.fromConstantString("undefined.example.");
+    SetResponse resp = ZONE.findRecords(testName, Type.ANY);
+    assertTrue(resp.isSuccessful());
+    assertEquals(
+        listOf(new RRset(A_WILD.withName(testName)), new RRset(TXT_WILD.withName(testName))),
+        resp.answers());
+  }
+
+  private static List<RRset> listOf(RRset... rrsets) {
+    return Stream.of(rrsets).collect(Collectors.toList());
+  }
+
+  private static List<RRset> oneRRset(Record... r) {
+    return Collections.singletonList(new RRset(r));
+  }
+}