RDAP: Fixes for domain search by nameserver

Use ordered sets to make sure that we don't pick up the same domain more than once if we find more than one of its nameservers, while preserving the order in which the domains are found. Add tests to make sure it works.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=134821458

java/google/registry/rdap/RdapDomainSearchAction.java

@@ -43,6 +43,7 @@ import google.registry.util.Clock;
 import google.registry.util.Idn;
 import java.net.InetAddress;
 import java.util.ArrayList;
+import java.util.LinkedHashSet;
 import java.util.List;
 import javax.inject.Inject;
 import org.joda.time.DateTime;
@@ -199,32 +200,39 @@ public class RdapDomainSearchAction extends RdapActionBase {
   }
 
   /** Searches for domains by nameserver name, returning a JSON array of domain info maps. */
-  private ImmutableList<ImmutableMap<String, Object>>
+  private ImmutableList<ImmutableMap<String, Object>> searchByNameserverLdhName(
-      searchByNameserverLdhName(final RdapSearchPattern partialStringQuery, final DateTime now) {
+      final RdapSearchPattern partialStringQuery, final DateTime now) {
-    Iterable<Key<HostResource>> hostKeys;
+    Iterable<Key<HostResource>> hostKeys = getNameserverRefsByLdhName(partialStringQuery, now);
+    if (Iterables.isEmpty(hostKeys)) {
+      throw new NotFoundException("No matching nameservers found");
+    }
+    return searchByNameserverRefs(hostKeys, now);
+  }
+
+  /** Assembles a list of {@link HostResource} keys by name. */
+  private Iterable<Key<HostResource>>
+      getNameserverRefsByLdhName(final RdapSearchPattern partialStringQuery, final DateTime now) {
     // Handle queries without a wildcard; just load the host by foreign key in the usual way.
     if (!partialStringQuery.getHasWildcard()) {
       Key<HostResource> hostKey = loadAndGetKey(
           HostResource.class, partialStringQuery.getInitialString(), now);
       if (hostKey == null) {
         return ImmutableList.of();
+      } else {
+        return ImmutableList.of(hostKey);
       }
-      hostKeys = ImmutableList.of(hostKey);
     // Handle queries with a wildcard, but no suffix. Query the host resources themselves, rather
     // than the foreign key index, because then we have an index on fully qualified host name and
     // deletion time, so we can check the deletion status in the query itself. There are no pending
     // deletes for hosts, so we can call queryUndeleted.
     } else if (partialStringQuery.getSuffix() == null) {
       // TODO (b/24463238): figure out how to limit the size of these queries effectively
-      hostKeys =
+      return queryUndeleted(HostResource.class, "fullyQualifiedHostName", partialStringQuery, 1000)
-          queryUndeleted(HostResource.class, "fullyQualifiedHostName", partialStringQuery, 1000)
           .keys();
-      if (Iterables.isEmpty(hostKeys)) {
-        throw new NotFoundException("No matching nameservers found");
-      }
     // Handle queries with a wildcard and a suffix. In this case, it is more efficient to do things
     // differently. We use the suffix to look up the domain, then loop through the subordinate hosts
     // looking for matches.
+    // TODO(mountford): This might not be ok; it will only find nameservers on domains we control
     } else {
       DomainResource domainResource = loadByForeignKey(
           DomainResource.class, partialStringQuery.getSuffix(), now);
@@ -242,14 +250,9 @@ public class RdapDomainSearchAction extends RdapActionBase {
           }
         }
       }
-      hostKeys = builder.build();
+      return builder.build();
-      if (Iterables.isEmpty(hostKeys)) {
-        throw new NotFoundException("No matching nameservers found");
     }
   }
-    // Find all domains that link to any of these hosts, and return information about them.
-    return searchByNameserverRefs(hostKeys, now);
-  }
 
   /** Searches for domains by nameserver address, returning a JSON array of domain info maps. */
   private ImmutableList<ImmutableMap<String, Object>>
@@ -258,6 +261,7 @@ public class RdapDomainSearchAction extends RdapActionBase {
     // the query on nameserver name (because we're already using an inequality query), and it seems
     // dangerous and confusing to filter on deletion time differently between the two queries.
     // Find all domains that link to any of these hosts, and return information about them.
+    // TODO (b/24463238): figure out how to limit the size of these queries effectively
     return searchByNameserverRefs(
         ofy()
             .load()
@@ -277,16 +281,24 @@ public class RdapDomainSearchAction extends RdapActionBase {
   private ImmutableList<ImmutableMap<String, Object>>
       searchByNameserverRefs(final Iterable<Key<HostResource>> hostKeys, final DateTime now) {
     // We must break the query up into chunks, because the in operator is limited to 30 subqueries.
-    ImmutableList.Builder<DomainResource> domainListBuilder = new ImmutableList.Builder<>();
+    // Since it is possible for the same domain to show up more than once in our result list (if
+    // we do a wildcard nameserver search that returns multiple nameservers used by the same
+    // domain), we must create a set of resulting {@link DomainResource} objects. But we use a
+    // LinkedHashSet to preserve the order in which we found the domains.
+    LinkedHashSet<DomainResource> domainResources = new LinkedHashSet<>();
     for (List<Key<HostResource>> chunk : Iterables.partition(hostKeys, 30)) {
-      domainListBuilder.addAll(
+      domainResources.addAll(
           ofy().load()
               .type(DomainResource.class)
               .filter("nameservers.linked in", chunk)
               .filter("deletionTime >", now)
-              .limit(1000));
+              .limit(rdapResultSetMaxSize - domainResources.size())
+              .list());
+      if (domainResources.size() >= rdapResultSetMaxSize) {
+        break;
       }
-    return makeSearchResults(domainListBuilder.build(), now);
+    }
+    return makeSearchResults(ImmutableList.copyOf(domainResources), now);
   }
 
   /** Output JSON for a list of domains. */

javatests/google/registry/rdap/RdapDomainSearchActionTest.java

@@ -518,16 +518,32 @@ public class RdapDomainSearchActionTest {
     assertThat(response.getStatus()).isEqualTo(404);
   }
 
-  private void createManyDomains(int numActiveDomains, int numTotalDomainsPerActiveDomain) {
+  private void createManyDomainsAndHosts(
+      int numActiveDomains, int numTotalDomainsPerActiveDomain, int numHosts) {
+    ImmutableSet.Builder<Key<HostResource>> hostKeysBuilder = new ImmutableSet.Builder<>();
+    ImmutableSet.Builder<String> subordinateHostsBuilder = new ImmutableSet.Builder<>();
+    String mainDomainName = String.format("domain%d.lol", numTotalDomainsPerActiveDomain);
+    for (int i = 1; i <= numHosts; i++) {
+      String hostName = String.format("ns%d.%s", i, mainDomainName);
+      subordinateHostsBuilder.add(hostName);
+      HostResource host = makeAndPersistHostResource(
+          hostName, String.format("5.5.5.%d", i), clock.nowUtc().minusYears(1));
+      hostKeysBuilder.add(Key.create(host));
+    }
+    ImmutableSet<Key<HostResource>> hostKeys = hostKeysBuilder.build();
     // Create all the domains at once, then persist them in parallel, for increased efficiency.
     ImmutableList.Builder<DomainResource> domainsBuilder = new ImmutableList.Builder<>();
     for (int i = 1; i <= numActiveDomains * numTotalDomainsPerActiveDomain; i++) {
       String domainName = String.format("domain%d.lol", i);
       DomainResource.Builder builder =
           makeDomainResource(
-              domainName, contact1, contact2, contact3, hostNs1CatLol, hostNs2CatLol, registrar)
+              domainName, contact1, contact2, contact3, null, null, registrar)
           .asBuilder()
+          .setNameservers(hostKeys)
           .setCreationTimeForTest(clock.nowUtc().minusYears(3));
+      if (domainName.equals(mainDomainName)) {
+        builder.setSubordinateHosts(subordinateHostsBuilder.build());
+      }
       if (i % numTotalDomainsPerActiveDomain != 0) {
         builder = builder.setDeletionTime(clock.nowUtc().minusDays(1));
       }
@@ -551,7 +567,7 @@ public class RdapDomainSearchActionTest {
   @Test
   public void testDomainMatch_manyDeletedDomains_fullResultSet() throws Exception {
     // There are enough domains to fill a full result set; deleted domains are ignored.
-    createManyDomains(4, 4);
+    createManyDomainsAndHosts(4, 4, 2);
     Object obj = generateActualJson(RequestType.NAME, "domain*.lol");
     assertThat(response.getStatus()).isEqualTo(200);
     checkNumberOfDomainsInResult(obj, 4);
@@ -561,7 +577,7 @@ public class RdapDomainSearchActionTest {
   public void testDomainMatch_manyDeletedDomains_partialResultSetDueToInsufficientDomains()
       throws Exception {
     // There are not enough domains to fill a full result set.
-    createManyDomains(3, 100);
+    createManyDomainsAndHosts(3, 100, 2);
     Object obj = generateActualJson(RequestType.NAME, "domain*.lol");
     assertThat(response.getStatus()).isEqualTo(200);
     checkNumberOfDomainsInResult(obj, 3);
@@ -573,7 +589,7 @@ public class RdapDomainSearchActionTest {
     // This is not exactly desired behavior, but expected: There are enough domains to fill a full
     // result set, but there are so many deleted domains that we run out of patience before we work
     // our way through all of them.
-    createManyDomains(4, 150);
+    createManyDomainsAndHosts(4, 150, 2);
     Object obj = generateActualJson(RequestType.NAME, "domain*.lol");
     assertThat(response.getStatus()).isEqualTo(200);
     checkNumberOfDomainsInResult(obj, 3);
@@ -694,7 +710,8 @@ public class RdapDomainSearchActionTest {
     persistResource(
         hostNs1CatLol.asBuilder().setDeletionTime(clock.nowUtc().minusDays(1)).build());
     assertThat(generateActualJson(RequestType.NS_LDH_NAME, "ns1.cat.lol"))
-        .isEqualTo(generateExpectedJson("No domains found", null, null, "rdap_error_404.json"));
+        .isEqualTo(generateExpectedJson(
+            "No matching nameservers found", null, null, "rdap_error_404.json"));
     assertThat(response.getStatus()).isEqualTo(404);
   }
 
@@ -718,6 +735,37 @@ public class RdapDomainSearchActionTest {
     assertThat(response.getStatus()).isEqualTo(404);
   }
 
+  @Test
+  public void testNameserverMatchManyNameserversForTheSameDomains() throws Exception {
+    // 40 nameservers for each of 3 domains; we should get back all three undeleted domains, because
+    // each one references the nameserver.
+    createManyDomainsAndHosts(3, 1, 40);
+    Object obj = generateActualJson(RequestType.NS_LDH_NAME, "ns1.domain1.lol");
+    assertThat(response.getStatus()).isEqualTo(200);
+    checkNumberOfDomainsInResult(obj, 3);
+  }
+
+  @Test
+  public void testNameserverMatchManyNameserversForTheSameDomainsWithWildcard() throws Exception {
+    // Same as above, except with a wildcard (that still only finds one nameserver).
+    createManyDomainsAndHosts(3, 1, 40);
+    Object obj = generateActualJson(RequestType.NS_LDH_NAME, "ns1.domain1.l*");
+    assertThat(response.getStatus()).isEqualTo(200);
+    checkNumberOfDomainsInResult(obj, 3);
+  }
+
+  @Test
+  public void testNameserverMatchManyNameserversForTheSameDomainsWithSuffix() throws Exception {
+    // Same as above, except that we find all 40 nameservers because of the wildcard. But we
+    // should still only return 3 domains, because we merge duplicate domains together in a set.
+    // Since we fetch domains by nameserver in batches of 30 nameservers, we need to make sure to
+    // have more than that number of nameservers for an effective test.
+    createManyDomainsAndHosts(3, 1, 40);
+    Object obj = generateActualJson(RequestType.NS_LDH_NAME, "ns*.domain1.lol");
+    assertThat(response.getStatus()).isEqualTo(200);
+    checkNumberOfDomainsInResult(obj, 3);
+  }
+
   @Test
   public void testAddressMatchV4Address_foundMultiple() throws Exception {
     assertThat(generateActualJson(RequestType.NS_IP, "1.2.3.4"))