Invoke separate IPv4 and IPv6 DNS queries in parallel am: 5404598e33

Original change: https://googleplex-android-review.googlesource.com/c/platform/packages/services/Iwlan/+/23237099

Change-Id: Id9099b9f4850b197c76a250136ad8050250cf88d
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

4 files changed, 183 insertions, 62 deletions

diff --git a/src/com/google/android/iwlan/IwlanHelper.java b/src/com/google/android/iwlan/IwlanHelper.java
index efda5b7..ade7189 100644
--- a/src/com/google/android/iwlan/IwlanHelper.java
+++ b/src/com/google/android/iwlan/IwlanHelper.java
@@ -110,14 +110,15 @@ public class IwlanHelper {
         return info;
     }
 
-    public static List<InetAddress> getAddressesForNetwork(Network network, Context context) {
+    // Retrieves all IP addresses for this Network, including stacked IPv4 link addresses.
+    public static List<InetAddress> getAllAddressesForNetwork(Network network, Context context) {
         ConnectivityManager connectivityManager =
                 context.getSystemService(ConnectivityManager.class);
         List<InetAddress> gatewayList = new ArrayList<>();
         if (network != null) {
             LinkProperties linkProperties = connectivityManager.getLinkProperties(network);
             if (linkProperties != null) {
-                for (LinkAddress linkAddr : linkProperties.getLinkAddresses()) {
+                for (LinkAddress linkAddr : linkProperties.getAllLinkAddresses()) {
                     InetAddress inetAddr = linkAddr.getAddress();
                     // skip linklocal and loopback addresses
                     if (!inetAddr.isLoopbackAddress() && !inetAddr.isLinkLocalAddress()) {
@@ -132,25 +133,6 @@ public class IwlanHelper {
         return gatewayList;
     }
 
-    public static List<InetAddress> getStackedAddressesForNetwork(
-            Network network, Context context) {
-        ConnectivityManager connectivityManager =
-                context.getSystemService(ConnectivityManager.class);
-        List<InetAddress> gatewayList = new ArrayList<>();
-        if (network != null) {
-            LinkProperties linkProperties = connectivityManager.getLinkProperties(network);
-            if (linkProperties != null) {
-                for (LinkAddress linkAddr : linkProperties.getAllLinkAddresses()) {
-                    InetAddress inetAddr = linkAddr.getAddress();
-                    if ((inetAddr instanceof Inet4Address)) {
-                        gatewayList.add(inetAddr);
-                    }
-                }
-            }
-        }
-        return gatewayList;
-    }
-
     /**
      * The method is to check if this IP address is an IPv4-embedded IPv6 address(Pref64::/n).
      *
@@ -162,22 +144,24 @@ public class IwlanHelper {
     }
 
     public static boolean hasIpv6Address(List<InetAddress> localAddresses) {
-        for (InetAddress address : localAddresses) {
-            if (address instanceof Inet6Address) {
-                return true;
+        if (localAddresses != null) {
+            for (InetAddress address : localAddresses) {
+                if (address instanceof Inet6Address) {
+                    return true;
+                }
             }
         }
-
         return false;
     }
 
     public static boolean hasIpv4Address(List<InetAddress> localAddresses) {
-        for (InetAddress address : localAddresses) {
-            if (address instanceof Inet4Address) {
-                return true;
+        if (localAddresses != null) {
+            for (InetAddress address : localAddresses) {
+                if (address instanceof Inet4Address) {
+                    return true;
+                }
             }
         }
-
         return false;
     }
 
diff --git a/src/com/google/android/iwlan/epdg/EpdgSelector.java b/src/com/google/android/iwlan/epdg/EpdgSelector.java
index 52392fd..e56196c 100644
--- a/src/com/google/android/iwlan/epdg/EpdgSelector.java
+++ b/src/com/google/android/iwlan/epdg/EpdgSelector.java
@@ -84,7 +84,8 @@ public class EpdgSelector {
     // IWLAN applies an internal timeout of 6 seconds, slightly longer than the default timeout
     private static final long DNS_RESOLVER_TIMEOUT_DURATION_SEC = 6L;
 
-    private static final long PARALLEL_DNS_RESOLVER_TIMEOUT_DURATION_SEC = 20L;
+    private static final long PARALLEL_STATIC_RESOLUTION_TIMEOUT_DURATION_SEC = 6L;
+    private static final long PARALLEL_PLMN_RESOLUTION_TIMEOUT_DURATION_SEC = 20L;
     private static final int NUM_EPDG_SELECTION_EXECUTORS = 2; // 1 each for normal selection, SOS.
     private static final int MAX_EPDG_SELECTION_THREADS = 2; // 1 each for prefetch, tunnel bringup.
     private static final int MAX_DNS_RESOLVER_THREADS = 25; // Do not expect > 25 FQDNs per carrier.
@@ -206,7 +207,7 @@ public class EpdgSelector {
     }
 
     private CompletableFuture<Map.Entry<String, List<InetAddress>>> submitDnsResolverQuery(
-            String domainName, Network network, Executor executor) {
+            String domainName, Network network, int queryType, Executor executor) {
         CompletableFuture<Map.Entry<String, List<InetAddress>>> result = new CompletableFuture();
 
         final DnsResolver.Callback<List<InetAddress>> cb =
@@ -234,7 +235,7 @@ public class EpdgSelector {
                     }
                 };
         DnsResolver.getInstance()
-                .query(network, domainName, DnsResolver.FLAG_EMPTY, executor, null, cb);
+                .query(network, domainName, queryType, DnsResolver.FLAG_EMPTY, executor, null, cb);
         return result;
     }
 
@@ -280,6 +281,16 @@ public class EpdgSelector {
                                 .collect(Collectors.<T>toList()));
     }
 
+    @VisibleForTesting
+    protected boolean hasIpv4Address(Network network) {
+        return IwlanHelper.hasIpv4Address(IwlanHelper.getAllAddressesForNetwork(network, mContext));
+    }
+
+    @VisibleForTesting
+    protected boolean hasIpv6Address(Network network) {
+        return IwlanHelper.hasIpv6Address(IwlanHelper.getAllAddressesForNetwork(network, mContext));
+    }
+
     /**
      * Returns a list of unique IP addresses corresponding to the given domain names, in the same
      * order of the input. Runs DNS resolution across parallel threads.
@@ -287,10 +298,11 @@ public class EpdgSelector {
      * @param domainNames Domain names for which DNS resolution needs to be performed.
      * @param filter Selects for IPv4, IPv6 (or both) addresses from the resulting DNS records
      * @param network {@link Network} Network on which to run the DNS query.
+     * @param timeout timeout in seconds.
      * @return List of unique IP addresses corresponding to the domainNames.
      */
     private LinkedHashMap<String, List<InetAddress>> getIP(
-            List<String> domainNames, int filter, Network network) {
+            List<String> domainNames, int filter, Network network, long timeout) {
         // LinkedHashMap preserves insertion order (and hence priority) of domain names passed in.
         LinkedHashMap<String, List<InetAddress>> domainNameToIpAddr = new LinkedHashMap<>();
 
@@ -298,16 +310,27 @@ public class EpdgSelector {
                 new ArrayList<>();
         for (String domainName : domainNames) {
             domainNameToIpAddr.put(domainName, new ArrayList<>());
-            futuresList.add(submitDnsResolverQuery(domainName, network, mDnsResolutionExecutor));
+            // Dispatches separate IPv4 and IPv6 queries to avoid being blocked on either result.
+            if (hasIpv4Address(network)) {
+                futuresList.add(
+                        submitDnsResolverQuery(
+                                domainName, network, DnsResolver.TYPE_A, mDnsResolutionExecutor));
+            }
+            if (hasIpv6Address(network)) {
+                futuresList.add(
+                        submitDnsResolverQuery(
+                                domainName,
+                                network,
+                                DnsResolver.TYPE_AAAA,
+                                mDnsResolutionExecutor));
+            }
         }
         CompletableFuture<List<Map.Entry<String, List<InetAddress>>>> allFuturesResult =
                 allOf(futuresList);
 
         List<Map.Entry<String, List<InetAddress>>> resultList = null;
         try {
-            resultList =
-                    allFuturesResult.get(
-                            PARALLEL_DNS_RESOLVER_TIMEOUT_DURATION_SEC, TimeUnit.SECONDS);
+            resultList = allFuturesResult.get(timeout, TimeUnit.SECONDS);
         } catch (ExecutionException e) {
             Log.e(TAG, "Cause of ExecutionException: ", e.getCause());
         } catch (InterruptedException e) {
@@ -320,8 +343,17 @@ public class EpdgSelector {
                 Log.w(TAG, "No IP addresses in parallel DNS query!");
             } else {
                 for (Map.Entry<String, List<InetAddress>> entry : resultList) {
-                    domainNameToIpAddr.put(
-                            entry.getKey(), v4v6ProtocolFilter(entry.getValue(), filter));
+                    String resultDomainName = entry.getKey();
+                    List<InetAddress> resultIpAddr = v4v6ProtocolFilter(entry.getValue(), filter);
+
+                    if (!domainNameToIpAddr.containsKey(resultDomainName)) {
+                        Log.w(
+                                TAG,
+                                "Unexpected domain name in DnsResolver result: "
+                                        + resultDomainName);
+                        continue;
+                    }
+                    domainNameToIpAddr.get(resultDomainName).addAll(resultIpAddr);
                 }
             }
         }
@@ -570,9 +602,13 @@ public class EpdgSelector {
         }
 
         Log.d(TAG, "Static Domain Names: " + Arrays.toString(domainNames));
-        for (String domainName : domainNames) {
-            getIP(domainName, filter, validIpList, network);
-        }
+        LinkedHashMap<String, List<InetAddress>> domainNameToIpAddr =
+                getIP(
+                        Arrays.asList(domainNames),
+                        filter,
+                        network,
+                        PARALLEL_STATIC_RESOLUTION_TIMEOUT_DURATION_SEC);
+        domainNameToIpAddr.values().forEach(validIpList::addAll);
     }
 
     private String[] getDomainNames(String key) {
@@ -647,7 +683,7 @@ public class EpdgSelector {
         }
 
         LinkedHashMap<String, List<InetAddress>> domainNameToIpAddr =
-                getIP(domainNames, filter, network);
+                getIP(domainNames, filter, network, PARALLEL_PLMN_RESOLUTION_TIMEOUT_DURATION_SEC);
         domainNameToIpAddr.values().forEach(validIpList::addAll);
         return domainNameToIpAddr;
     }
diff --git a/src/com/google/android/iwlan/epdg/EpdgTunnelManager.java b/src/com/google/android/iwlan/epdg/EpdgTunnelManager.java
index bfe4f95..b8b70e4 100644
--- a/src/com/google/android/iwlan/epdg/EpdgTunnelManager.java
+++ b/src/com/google/android/iwlan/epdg/EpdgTunnelManager.java
@@ -2283,7 +2283,7 @@ public class EpdgTunnelManager {
 
     @VisibleForTesting
     List<InetAddress> getAddressForNetwork(Network network, Context context) {
-        return IwlanHelper.getAddressesForNetwork(network, context);
+        return IwlanHelper.getAllAddressesForNetwork(network, context);
     }
 
     @VisibleForTesting
diff --git a/test/com/google/android/iwlan/epdg/EpdgSelectorTest.java b/test/com/google/android/iwlan/epdg/EpdgSelectorTest.java
index b2d514c..97526d1 100644
--- a/test/com/google/android/iwlan/epdg/EpdgSelectorTest.java
+++ b/test/com/google/android/iwlan/epdg/EpdgSelectorTest.java
@@ -124,7 +124,7 @@ public class EpdgSelectorTest {
 
         when(ErrorPolicyManager.getInstance(mMockContext, DEFAULT_SLOT_INDEX))
                 .thenReturn(mMockErrorPolicyManager);
-        mEpdgSelector = new EpdgSelector(mMockContext, DEFAULT_SLOT_INDEX);
+        mEpdgSelector = spy(new EpdgSelector(mMockContext, DEFAULT_SLOT_INDEX));
 
         when(mMockContext.getSystemService(eq(SubscriptionManager.class)))
                 .thenReturn(mMockSubscriptionManager);
@@ -163,8 +163,6 @@ public class EpdgSelectorTest {
                 .thenReturn(mMockCarrierConfigManager);
         when(mMockCarrierConfigManager.getConfigForSubId(anyInt())).thenReturn(mTestBundle);
 
-        lenient().when(DnsResolver.getInstance()).thenReturn(mMockDnsResolver);
-
         mFakeDns = new FakeDns();
         mFakeDns.startMocking();
     }
@@ -177,6 +175,10 @@ public class EpdgSelectorTest {
 
     @Test
     public void testStaticMethodPass() throws Exception {
+        when(DnsResolver.getInstance()).thenReturn(mMockDnsResolver);
+        doReturn(true).when(mEpdgSelector).hasIpv4Address(mMockNetwork);
+        doReturn(true).when(mEpdgSelector).hasIpv6Address(mMockNetwork);
+
         // Set DnsResolver query mock
         final String testStaticAddress = "epdg.epc.mnc088.mcc888.pub.3gppnetwork.org";
         mFakeDns.setAnswer(testStaticAddress, new String[] {TEST_IP_ADDRESS}, TYPE_A);
@@ -199,6 +201,10 @@ public class EpdgSelectorTest {
 
     @Test
     public void testRoamStaticMethodPass() throws Exception {
+        when(DnsResolver.getInstance()).thenReturn(mMockDnsResolver);
+        doReturn(true).when(mEpdgSelector).hasIpv4Address(mMockNetwork);
+        doReturn(true).when(mEpdgSelector).hasIpv6Address(mMockNetwork);
+
         // Set DnsResolver query mock
         final String testRoamStaticAddress = "epdg.epc.mnc088.mcc888.pub.3gppnetwork.org";
         mFakeDns.setAnswer(testRoamStaticAddress, new String[] {TEST_IP_ADDRESS}, TYPE_A);
@@ -232,6 +238,10 @@ public class EpdgSelectorTest {
 
     @Test
     public void testPlmnResolutionMethodWithNoPlmnInCarrierConfig() throws Exception {
+        when(DnsResolver.getInstance()).thenReturn(mMockDnsResolver);
+        doReturn(true).when(mEpdgSelector).hasIpv4Address(mMockNetwork);
+        doReturn(true).when(mEpdgSelector).hasIpv6Address(mMockNetwork);
+
         // setUp() fills default values for mcc-mnc
         String expectedFqdnFromImsi = "epdg.epc.mnc120.mcc311.pub.3gppnetwork.org";
         String expectedFqdnFromEhplmn = "epdg.epc.mnc120.mcc300.pub.3gppnetwork.org";
@@ -248,6 +258,10 @@ public class EpdgSelectorTest {
     }
 
     private void testPlmnResolutionMethod(boolean isEmergency) throws Exception {
+        when(DnsResolver.getInstance()).thenReturn(mMockDnsResolver);
+        doReturn(true).when(mEpdgSelector).hasIpv4Address(mMockNetwork);
+        doReturn(true).when(mEpdgSelector).hasIpv6Address(mMockNetwork);
+
         String expectedFqdnFromImsi = "epdg.epc.mnc120.mcc311.pub.3gppnetwork.org";
         String expectedFqdnFromRplmn = "epdg.epc.mnc121.mcc311.pub.3gppnetwork.org";
         String expectedFqdnFromEhplmn = "epdg.epc.mnc120.mcc300.pub.3gppnetwork.org";
@@ -295,6 +309,10 @@ public class EpdgSelectorTest {
 
     @Test
     public void testPlmnResolutionMethodWithDuplicatedImsiAndEhplmn() throws Exception {
+        when(DnsResolver.getInstance()).thenReturn(mMockDnsResolver);
+        doReturn(true).when(mEpdgSelector).hasIpv4Address(mMockNetwork);
+        doReturn(true).when(mEpdgSelector).hasIpv6Address(mMockNetwork);
+
         String fqdnFromEhplmn1 = "epdg.epc.mnc120.mcc300.pub.3gppnetwork.org";
         String fqdnFromEhplmn2AndImsi = "epdg.epc.mnc120.mcc311.pub.3gppnetwork.org";
         String fqdnFromEhplmn3 = "epdg.epc.mnc122.mcc300.pub.3gppnetwork.org";
@@ -330,6 +348,10 @@ public class EpdgSelectorTest {
 
     @Test
     public void testPlmnResolutionMethodWithInvalidLengthPlmns() throws Exception {
+        when(DnsResolver.getInstance()).thenReturn(mMockDnsResolver);
+        doReturn(true).when(mEpdgSelector).hasIpv4Address(mMockNetwork);
+        doReturn(true).when(mEpdgSelector).hasIpv6Address(mMockNetwork);
+
         when(mMockSubscriptionInfo.getMccString()).thenReturn("31");
         when(mMockSubscriptionInfo.getMncString()).thenReturn("12");
 
@@ -355,6 +377,10 @@ public class EpdgSelectorTest {
 
     @Test
     public void testPlmnResolutionMethodWithInvalidCharacterPlmns() throws Exception {
+        when(DnsResolver.getInstance()).thenReturn(mMockDnsResolver);
+        doReturn(true).when(mEpdgSelector).hasIpv4Address(mMockNetwork);
+        doReturn(true).when(mEpdgSelector).hasIpv6Address(mMockNetwork);
+
         when(mMockSubscriptionInfo.getMccString()).thenReturn("a b");
         when(mMockSubscriptionInfo.getMncString()).thenReturn("!@#");
 
@@ -381,6 +407,10 @@ public class EpdgSelectorTest {
 
     @Test
     public void testPlmnResolutionMethodWithEmptyPlmns() throws Exception {
+        when(DnsResolver.getInstance()).thenReturn(mMockDnsResolver);
+        doReturn(true).when(mEpdgSelector).hasIpv4Address(mMockNetwork);
+        doReturn(true).when(mEpdgSelector).hasIpv6Address(mMockNetwork);
+
         when(mMockSubscriptionInfo.getMccString()).thenReturn(null);
         when(mMockSubscriptionInfo.getMncString()).thenReturn(null);
 
@@ -405,6 +435,10 @@ public class EpdgSelectorTest {
 
     @Test
     public void testPlmnResolutionMethodWithFirstEhplmn() throws Exception {
+        when(DnsResolver.getInstance()).thenReturn(mMockDnsResolver);
+        doReturn(true).when(mEpdgSelector).hasIpv4Address(mMockNetwork);
+        doReturn(true).when(mEpdgSelector).hasIpv6Address(mMockNetwork);
+
         String fqdnFromEhplmn1 = "epdg.epc.mnc120.mcc300.pub.3gppnetwork.org";
         String fqdnFromEhplmn2 = "epdg.epc.mnc121.mcc300.pub.3gppnetwork.org";
         String fqdnFromEhplmn3 = "epdg.epc.mnc122.mcc300.pub.3gppnetwork.org";
@@ -434,6 +468,10 @@ public class EpdgSelectorTest {
 
     @Test
     public void testPlmnResolutionMethodWithRplmn() throws Exception {
+        when(DnsResolver.getInstance()).thenReturn(mMockDnsResolver);
+        doReturn(true).when(mEpdgSelector).hasIpv4Address(mMockNetwork);
+        doReturn(true).when(mEpdgSelector).hasIpv6Address(mMockNetwork);
+
         String fqdnFromRplmn = "epdg.epc.mnc122.mcc300.pub.3gppnetwork.org";
         String fqdnFromEhplmn1 = "epdg.epc.mnc120.mcc300.pub.3gppnetwork.org";
         String fqdnFromEhplmn2 = "epdg.epc.mnc121.mcc300.pub.3gppnetwork.org";
@@ -464,6 +502,10 @@ public class EpdgSelectorTest {
 
     @Test
     public void testCarrierConfigStaticAddressList() throws Exception {
+        when(DnsResolver.getInstance()).thenReturn(mMockDnsResolver);
+        doReturn(true).when(mEpdgSelector).hasIpv4Address(mMockNetwork);
+        doReturn(true).when(mEpdgSelector).hasIpv6Address(mMockNetwork);
+
         // Set Network.getAllByName mock
         final String addr1 = "epdg.epc.mnc480.mcc310.pub.3gppnetwork.org";
         final String addr2 = "epdg.epc.mnc120.mcc300.pub.3gppnetwork.org";
@@ -588,6 +630,8 @@ public class EpdgSelectorTest {
     }
 
     private void testCellularResolutionMethod(boolean isEmergency) throws Exception {
+        when(DnsResolver.getInstance()).thenReturn(mMockDnsResolver);
+
         int testMcc = 311;
         int testMnc = 120;
         String testMccString = "311";
@@ -683,6 +727,10 @@ public class EpdgSelectorTest {
 
     @Test
     public void testGetValidatedServerListIpv4Preferred() throws Exception {
+        when(DnsResolver.getInstance()).thenReturn(mMockDnsResolver);
+        doReturn(true).when(mEpdgSelector).hasIpv4Address(mMockNetwork);
+        doReturn(true).when(mEpdgSelector).hasIpv6Address(mMockNetwork);
+
         final String addr1 = "epdg.epc.mnc120.mcc300.pub.3gppnetwork.org";
         final String addr2 = "epdg.epc.mnc120.mcc311.pub.3gppnetwork.org";
         final String testStaticAddress = addr1 + "," + addr2;
@@ -710,6 +758,10 @@ public class EpdgSelectorTest {
 
     @Test
     public void testGetValidatedServerListIpv6Preferred() throws Exception {
+        when(DnsResolver.getInstance()).thenReturn(mMockDnsResolver);
+        doReturn(true).when(mEpdgSelector).hasIpv4Address(mMockNetwork);
+        doReturn(true).when(mEpdgSelector).hasIpv6Address(mMockNetwork);
+
         final String addr1 = "epdg.epc.mnc120.mcc300.pub.3gppnetwork.org";
         final String addr2 = "epdg.epc.mnc120.mcc311.pub.3gppnetwork.org";
         final String testStaticAddress = addr1 + "," + addr2;
@@ -737,6 +789,10 @@ public class EpdgSelectorTest {
 
     @Test
     public void testGetValidatedServerListIpv4Only() throws Exception {
+        when(DnsResolver.getInstance()).thenReturn(mMockDnsResolver);
+        doReturn(true).when(mEpdgSelector).hasIpv4Address(mMockNetwork);
+        doReturn(true).when(mEpdgSelector).hasIpv6Address(mMockNetwork);
+
         final String addr1 = "epdg.epc.mnc120.mcc300.pub.3gppnetwork.org";
         final String addr2 = "epdg.epc.mnc120.mcc311.pub.3gppnetwork.org";
         final String testStaticAddress = addr1 + "," + addr2;
@@ -763,6 +819,10 @@ public class EpdgSelectorTest {
 
     @Test
     public void testGetValidatedServerListIpv4OnlyCongestion() throws Exception {
+        when(DnsResolver.getInstance()).thenReturn(mMockDnsResolver);
+        doReturn(true).when(mEpdgSelector).hasIpv4Address(mMockNetwork);
+        doReturn(true).when(mEpdgSelector).hasIpv6Address(mMockNetwork);
+
         when(mMockErrorPolicyManager.getMostRecentDataFailCause())
                 .thenReturn(DataFailCause.IWLAN_CONGESTION);
         when(mMockErrorPolicyManager.getCurrentFqdnIndex(anyInt())).thenReturn(0);
@@ -794,6 +854,10 @@ public class EpdgSelectorTest {
 
     @Test
     public void testGetValidatedServerListIpv6Only() throws Exception {
+        when(DnsResolver.getInstance()).thenReturn(mMockDnsResolver);
+        doReturn(true).when(mEpdgSelector).hasIpv4Address(mMockNetwork);
+        doReturn(true).when(mEpdgSelector).hasIpv6Address(mMockNetwork);
+
         final String addr1 = "epdg.epc.mnc120.mcc300.pub.3gppnetwork.org";
         final String addr2 = "epdg.epc.mnc120.mcc311.pub.3gppnetwork.org";
         final String testStaticAddress = addr1 + "," + addr2;
@@ -820,6 +884,10 @@ public class EpdgSelectorTest {
 
     @Test
     public void testGetValidatedServerListSystemPreferred() throws Exception {
+        when(DnsResolver.getInstance()).thenReturn(mMockDnsResolver);
+        doReturn(true).when(mEpdgSelector).hasIpv4Address(mMockNetwork);
+        doReturn(true).when(mEpdgSelector).hasIpv6Address(mMockNetwork);
+
         final String addr1 = "epdg.epc.mnc120.mcc300.pub.3gppnetwork.org";
         final String addr2 = "epdg.epc.mnc120.mcc311.pub.3gppnetwork.org";
         final String addr3 = "epdg.epc.mnc120.mcc312.pub.3gppnetwork.org";
@@ -873,20 +941,20 @@ public class EpdgSelectorTest {
             }
         }
 
-        private final ArrayList<DnsEntry> mAnswers = new ArrayList<DnsEntry>();
+        private final List<DnsEntry> mAnswers = new ArrayList<>();
 
         /** Clears all DNS entries. */
         private synchronized void clearAll() {
             mAnswers.clear();
         }
 
-        /** Returns the answer for a given name and type on the given mock network. */
-        private synchronized List<InetAddress> getAnswer(Object mock, String hostname, int type) {
+        /** Returns the answer for a given name and type. */
+        private synchronized List<InetAddress> getAnswer(String hostname, int type) {
             return mAnswers.stream()
                     .filter(e -> e.matches(hostname, type))
                     .map(answer -> answer.mAddresses)
                     .findFirst()
-                    .orElse(null);
+                    .orElse(List.of());
         }
 
         /** Sets the answer for a given name and type. */
@@ -907,10 +975,20 @@ public class EpdgSelectorTest {
         }
 
         // Regardless of the type, depends on what the responses contained in the network.
-        private List<InetAddress> queryAllTypes(Object mock, String hostname) {
+        private List<InetAddress> queryIpv4(String hostname) {
+            return getAnswer(hostname, TYPE_A);
+        }
+
+        // Regardless of the type, depends on what the responses contained in the network.
+        private List<InetAddress> queryIpv6(String hostname) {
+            return getAnswer(hostname, TYPE_AAAA);
+        }
+
+        // Regardless of the type, depends on what the responses contained in the network.
+        private List<InetAddress> queryAllTypes(String hostname) {
             List<InetAddress> answer = new ArrayList<>();
-            addAllIfNotNull(answer, getAnswer(mock, hostname, TYPE_A));
-            addAllIfNotNull(answer, getAnswer(mock, hostname, TYPE_AAAA));
+            answer.addAll(queryIpv4(hostname));
+            answer.addAll(queryIpv6(hostname));
             return answer;
         }
 
@@ -922,32 +1000,55 @@ public class EpdgSelectorTest {
 
         /** Starts mocking DNS queries. */
         private void startMocking() throws UnknownHostException {
+            // 5-arg DnsResolver.query()
             doAnswer(
                             invocation -> {
                                 return mockQuery(
                                         invocation,
                                         1 /* posHostname */,
+                                        -1 /* posType */,
                                         3 /* posExecutor */,
-                                        5 /* posCallback */,
-                                        -1 /* posType */);
+                                        5 /* posCallback */);
                             })
                     .when(mMockDnsResolver)
-                    .query(any(), any(), anyInt(), any(), any(), any());
+                    .query(any(), anyString(), anyInt(), any(), any(), any());
+
+            // 6-arg DnsResolver.query() with explicit query type (IPv4 or v6).
+            doAnswer(
+                            invocation -> {
+                                return mockQuery(
+                                        invocation,
+                                        1 /* posHostname */,
+                                        2 /* posType */,
+                                        4 /* posExecutor */,
+                                        6 /* posCallback */);
+                            })
+                    .when(mMockDnsResolver)
+                    .query(any(), anyString(), anyInt(), anyInt(), any(), any(), any());
         }
 
         // Mocking queries on DnsResolver#query.
         private Answer mockQuery(
                 InvocationOnMock invocation,
                 int posHostname,
+                int posType,
                 int posExecutor,
-                int posCallback,
-                int posType) {
-            String hostname = (String) invocation.getArgument(posHostname);
-            Executor executor = (Executor) invocation.getArgument(posExecutor);
+                int posCallback) {
+            String hostname = invocation.getArgument(posHostname);
+            Executor executor = invocation.getArgument(posExecutor);
             DnsResolver.Callback<List<InetAddress>> callback = invocation.getArgument(posCallback);
             List<InetAddress> answer;
 
-            answer = queryAllTypes(invocation.getMock(), hostname);
+            switch (posType) {
+                case TYPE_A:
+                    answer = queryIpv4(hostname);
+                    break;
+                case TYPE_AAAA:
+                    answer = queryIpv6(hostname);
+                    break;
+                default:
+                    answer = queryAllTypes(hostname);
+            }
 
             if (answer != null && answer.size() > 0) {
                 new Handler(Looper.getMainLooper())