Fix a bug where private DNS to v4-only servers don't validate.

Bug: 132670008
Test: Manual
Test: atest NetworkStackTests
Change-Id: I8ae6fa6f3800f5f6c9ee88ffe3b8ada34b96dee6
(cherry picked from commit aea0808d3b26dab24005afc39eb7ec4f4f333bae)

2 files changed, 20 insertions, 7 deletions

diff --git a/packages/NetworkStack/src/com/android/networkstack/util/DnsUtils.java b/packages/NetworkStack/src/com/android/networkstack/util/DnsUtils.java
index 85f94e17a088..4767d5574a00 100644
--- a/packages/NetworkStack/src/com/android/networkstack/util/DnsUtils.java
+++ b/packages/NetworkStack/src/com/android/networkstack/util/DnsUtils.java
@@ -55,12 +55,23 @@ public class DnsUtils {
             throws UnknownHostException {
         final List<InetAddress> result = new ArrayList<InetAddress>();
 
-        result.addAll(Arrays.asList(
-                getAllByName(dnsResolver, network, host, TYPE_AAAA, FLAG_NO_CACHE_LOOKUP,
-                timeout)));
-        result.addAll(Arrays.asList(
-                getAllByName(dnsResolver, network, host, TYPE_A, FLAG_NO_CACHE_LOOKUP,
-                timeout)));
+        try {
+            result.addAll(Arrays.asList(
+                    getAllByName(dnsResolver, network, host, TYPE_AAAA, FLAG_NO_CACHE_LOOKUP,
+                    timeout)));
+        } catch (UnknownHostException e) {
+            // Might happen if the host is v4-only, still need to query TYPE_A
+        }
+        try {
+            result.addAll(Arrays.asList(
+                    getAllByName(dnsResolver, network, host, TYPE_A, FLAG_NO_CACHE_LOOKUP,
+                    timeout)));
+        } catch (UnknownHostException e) {
+            // Might happen if the host is v6-only, still need to return AAAA answers
+        }
+        if (result.size() == 0) {
+            throw new UnknownHostException(host);
+        }
         return result.toArray(new InetAddress[0]);
     }
 
diff --git a/packages/NetworkStack/tests/src/com/android/server/connectivity/NetworkMonitorTest.java b/packages/NetworkStack/tests/src/com/android/server/connectivity/NetworkMonitorTest.java
index abfb9c8ae282..832b7124dc05 100644
--- a/packages/NetworkStack/tests/src/com/android/server/connectivity/NetworkMonitorTest.java
+++ b/packages/NetworkStack/tests/src/com/android/server/connectivity/NetworkMonitorTest.java
@@ -298,8 +298,10 @@ public class NetworkMonitorTest {
         setFallbackSpecs(null); // Test with no fallback spec by default
         when(mRandom.nextInt()).thenReturn(0);
 
+        // DNS probe timeout should not be defined more than half of HANDLER_TIMEOUT_MS. Otherwise,
+        // it will fail the test because of timeout expired for querying AAAA and A sequentially.
         when(mResources.getInteger(eq(R.integer.config_captive_portal_dns_probe_timeout)))
-                .thenReturn(500);
+                .thenReturn(200);
 
         doAnswer((invocation) -> {
             URL url = invocation.getArgument(0);