Merge "Fix querying of multiple subtypes" into main am: 742588457a

Original change: https://android-review.googlesource.com/c/platform/packages/modules/Connectivity/+/2939001

Change-Id: I01aaf966c5ea7eeb059371418d10ed4f7edcbe06
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
index 3a69d67..4cb88b4 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
@@ -39,6 +39,7 @@
 import java.net.Inet6Address;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
@@ -139,6 +140,7 @@
                     // before sending the query, it needs to be called just before sending it.
                     final List<MdnsResponse> servicesToResolve = makeResponsesForResolve(socketKey);
                     final QueryTask queryTask = new QueryTask(taskArgs, servicesToResolve,
+                            getAllDiscoverySubtypes(),
                             servicesToResolve.size() < listeners.size() /* sendDiscoveryQueries */);
                     executor.submit(queryTask);
                     break;
@@ -359,7 +361,6 @@
         // Keep tracking the ScheduledFuture for the task so we can cancel it if caller is not
         // interested anymore.
         final QueryTaskConfig taskConfig = new QueryTaskConfig(
-                searchOptions.getSubtypes(),
                 searchOptions.getQueryMode(),
                 searchOptions.onlyUseIpv6OnIpv6OnlyNetworks(),
                 searchOptions.numOfQueriesBeforeBackoff(),
@@ -387,6 +388,7 @@
             final QueryTask queryTask = new QueryTask(
                     mdnsQueryScheduler.scheduleFirstRun(taskConfig, now,
                             minRemainingTtl, currentSessionId), servicesToResolve,
+                    getAllDiscoverySubtypes(),
                     servicesToResolve.size() < listeners.size() /* sendDiscoveryQueries */);
             executor.submit(queryTask);
         }
@@ -394,6 +396,15 @@
         serviceCache.registerServiceExpiredCallback(cacheKey, serviceExpiredCallback);
     }
 
+    private Set<String> getAllDiscoverySubtypes() {
+        final Set<String> subtypes = MdnsUtils.newSet();
+        for (int i = 0; i < listeners.size(); i++) {
+            final MdnsSearchOptions listenerOptions = listeners.valueAt(i).searchOptions;
+            subtypes.addAll(listenerOptions.getSubtypes());
+        }
+        return subtypes;
+    }
+
     /**
      * Get the executor service.
      */
@@ -664,11 +675,15 @@
     private class QueryTask implements Runnable {
         private final MdnsQueryScheduler.ScheduledQueryTaskArgs taskArgs;
         private final List<MdnsResponse> servicesToResolve = new ArrayList<>();
+        private final List<String> subtypes = new ArrayList<>();
         private final boolean sendDiscoveryQueries;
         QueryTask(@NonNull MdnsQueryScheduler.ScheduledQueryTaskArgs taskArgs,
-                @NonNull List<MdnsResponse> servicesToResolve, boolean sendDiscoveryQueries) {
+                @NonNull Collection<MdnsResponse> servicesToResolve,
+                @NonNull Collection<String> subtypes,
+                boolean sendDiscoveryQueries) {
             this.taskArgs = taskArgs;
             this.servicesToResolve.addAll(servicesToResolve);
+            this.subtypes.addAll(subtypes);
             this.sendDiscoveryQueries = sendDiscoveryQueries;
         }
 
@@ -681,7 +696,7 @@
                                 socketClient,
                                 createMdnsPacketWriter(),
                                 serviceType,
-                                taskArgs.config.subtypes,
+                                subtypes,
                                 taskArgs.config.expectUnicastResponse,
                                 taskArgs.config.transactionId,
                                 taskArgs.config.socketKey,
@@ -693,7 +708,7 @@
                                 .call();
             } catch (RuntimeException e) {
                 sharedLog.e(String.format("Failed to run EnqueueMdnsQueryCallable for subtype: %s",
-                        TextUtils.join(",", taskArgs.config.subtypes)), e);
+                        TextUtils.join(",", subtypes)), e);
                 result = Pair.create(INVALID_TRANSACTION_ID, new ArrayList<>());
             }
             dependencies.sendMessage(
diff --git a/service-t/src/com/android/server/connectivity/mdns/QueryTaskConfig.java b/service-t/src/com/android/server/connectivity/mdns/QueryTaskConfig.java
index 63524a6..0894166 100644
--- a/service-t/src/com/android/server/connectivity/mdns/QueryTaskConfig.java
+++ b/service-t/src/com/android/server/connectivity/mdns/QueryTaskConfig.java
@@ -24,10 +24,6 @@
 
 import com.android.internal.annotations.VisibleForTesting;
 
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.List;
-
 /**
  * A configuration for the PeriodicalQueryTask that contains parameters to build a query packet.
  * Call to getConfigForNextRun returns a config that can be used to build the next query task.
@@ -53,9 +49,6 @@
     // they only listen once every [multiplier] DTIM intervals).
     static final int TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS = 100;
     static final int MAX_TIME_BETWEEN_AGGRESSIVE_BURSTS_MS = 60000;
-    // The following fields are used by QueryTask so we need to test them.
-    @VisibleForTesting
-    final List<String> subtypes;
     private final boolean alwaysAskForUnicastResponse =
             MdnsConfigs.alwaysAskForUnicastResponseInEachBurst();
     private final int queryMode;
@@ -78,7 +71,6 @@
             boolean expectUnicastResponse, boolean isFirstBurst, int burstCounter,
             int queriesPerBurst, int timeBetweenBurstsInMs,
             long delayUntilNextTaskWithoutBackoffMs) {
-        this.subtypes = new ArrayList<>(other.subtypes);
         this.queryMode = other.queryMode;
         this.onlyUseIpv6OnIpv6OnlyNetworks = other.onlyUseIpv6OnIpv6OnlyNetworks;
         this.numOfQueriesBeforeBackoff = other.numOfQueriesBeforeBackoff;
@@ -93,15 +85,13 @@
         this.socketKey = other.socketKey;
     }
 
-    QueryTaskConfig(@NonNull Collection<String> subtypes,
-            int queryMode,
+    QueryTaskConfig(int queryMode,
             boolean onlyUseIpv6OnIpv6OnlyNetworks,
             int numOfQueriesBeforeBackoff,
             @Nullable SocketKey socketKey) {
         this.queryMode = queryMode;
         this.onlyUseIpv6OnIpv6OnlyNetworks = onlyUseIpv6OnIpv6OnlyNetworks;
         this.numOfQueriesBeforeBackoff = numOfQueriesBeforeBackoff;
-        this.subtypes = new ArrayList<>(subtypes);
         this.queriesPerBurst = QUERIES_PER_BURST;
         this.burstCounter = 0;
         this.transactionId = 1;
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
index dab3b42..58124f3 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
@@ -241,15 +241,22 @@
             return true;
         }).when(mockDeps).sendMessage(any(Handler.class), any(Message.class));
 
-        client =
-                new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                        mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
-                        serviceCache) {
-                    @Override
-                    MdnsPacketWriter createMdnsPacketWriter() {
-                        return mockPacketWriter;
-                    }
-                };
+        client = makeMdnsServiceTypeClient(mockPacketWriter);
+    }
+
+    private MdnsServiceTypeClient makeMdnsServiceTypeClient(
+            @Nullable MdnsPacketWriter packetWriter) {
+        return new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
+                mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                serviceCache) {
+            @Override
+            MdnsPacketWriter createMdnsPacketWriter() {
+                if (packetWriter == null) {
+                    return super.createMdnsPacketWriter();
+                }
+                return packetWriter;
+            }
+        };
     }
 
     @After
@@ -562,13 +569,12 @@
         MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
                 .addSubtype(SUBTYPE).setQueryMode(ACTIVE_QUERY_MODE).build();
         QueryTaskConfig config = new QueryTaskConfig(
-                searchOptions.getSubtypes(), searchOptions.getQueryMode(),
+                searchOptions.getQueryMode(),
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */, 3 /* numOfQueriesBeforeBackoff */,
                 socketKey);
 
         // This is the first query. We will ask for unicast response.
         assertTrue(config.expectUnicastResponse);
-        assertEquals(config.subtypes, searchOptions.getSubtypes());
         assertEquals(config.transactionId, 1);
 
         // For the rest of queries in this burst, we will NOT ask for unicast response.
@@ -576,7 +582,6 @@
             int oldTransactionId = config.transactionId;
             config = config.getConfigForNextRun();
             assertFalse(config.expectUnicastResponse);
-            assertEquals(config.subtypes, searchOptions.getSubtypes());
             assertEquals(config.transactionId, oldTransactionId + 1);
         }
 
@@ -584,7 +589,6 @@
         int oldTransactionId = config.transactionId;
         config = config.getConfigForNextRun();
         assertTrue(config.expectUnicastResponse);
-        assertEquals(config.subtypes, searchOptions.getSubtypes());
         assertEquals(config.transactionId, oldTransactionId + 1);
     }
 
@@ -593,13 +597,12 @@
         MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
                 .addSubtype(SUBTYPE).setQueryMode(ACTIVE_QUERY_MODE).build();
         QueryTaskConfig config = new QueryTaskConfig(
-                searchOptions.getSubtypes(), searchOptions.getQueryMode(),
+                searchOptions.getQueryMode(),
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */, 3 /* numOfQueriesBeforeBackoff */,
                 socketKey);
 
         // This is the first query. We will ask for unicast response.
         assertTrue(config.expectUnicastResponse);
-        assertEquals(config.subtypes, searchOptions.getSubtypes());
         assertEquals(config.transactionId, 1);
 
         // For the rest of queries in this burst, we will NOT ask for unicast response.
@@ -607,7 +610,6 @@
             int oldTransactionId = config.transactionId;
             config = config.getConfigForNextRun();
             assertFalse(config.expectUnicastResponse);
-            assertEquals(config.subtypes, searchOptions.getSubtypes());
             assertEquals(config.transactionId, oldTransactionId + 1);
         }
 
@@ -615,7 +617,6 @@
         int oldTransactionId = config.transactionId;
         config = config.getConfigForNextRun();
         assertFalse(config.expectUnicastResponse);
-        assertEquals(config.subtypes, searchOptions.getSubtypes());
         assertEquals(config.transactionId, oldTransactionId + 1);
     }
 
@@ -693,6 +694,81 @@
                 any(), any(), eq(MdnsConfigs.timeBetweenQueriesInBurstMs()));
     }
 
+    @Test
+    public void testCombinedSubtypesQueriedWithMultipleListeners() throws Exception {
+        client = makeMdnsServiceTypeClient(/* packetWriter= */ null);
+        final MdnsSearchOptions searchOptions1 = MdnsSearchOptions.newBuilder()
+                .addSubtype("subtype1").build();
+        final MdnsSearchOptions searchOptions2 = MdnsSearchOptions.newBuilder()
+                .addSubtype("subtype2").build();
+        startSendAndReceive(mockListenerOne, searchOptions1);
+        currentThreadExecutor.getAndClearSubmittedRunnable().run();
+
+        InOrder inOrder = inOrder(mockListenerOne, mockSocketClient, mockDeps);
+
+        // Verify the query asks for subtype1
+        final ArgumentCaptor<DatagramPacket> subtype1QueryCaptor =
+                ArgumentCaptor.forClass(DatagramPacket.class);
+        currentThreadExecutor.getAndClearLastScheduledRunnable().run();
+        // Send twice for IPv4 and IPv6
+        inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingUnicastResponse(
+                subtype1QueryCaptor.capture(),
+                eq(socketKey), eq(false));
+
+        final MdnsPacket subtype1Query = MdnsPacket.parse(
+                new MdnsPacketReader(subtype1QueryCaptor.getValue()));
+
+        assertEquals(2, subtype1Query.questions.size());
+        assertTrue(hasQuestion(subtype1Query, MdnsRecord.TYPE_PTR, SERVICE_TYPE_LABELS));
+        assertTrue(hasQuestion(subtype1Query, MdnsRecord.TYPE_PTR,
+                getServiceTypeWithSubtype("_subtype1")));
+
+        // Add subtype2
+        startSendAndReceive(mockListenerTwo, searchOptions2);
+        inOrder.verify(mockDeps).removeMessages(any(), eq(EVENT_START_QUERYTASK));
+        currentThreadExecutor.getAndClearLastScheduledRunnable().run();
+
+        final ArgumentCaptor<DatagramPacket> combinedSubtypesQueryCaptor =
+                ArgumentCaptor.forClass(DatagramPacket.class);
+        inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingUnicastResponse(
+                combinedSubtypesQueryCaptor.capture(),
+                eq(socketKey), eq(false));
+        // The next query must have been scheduled
+        inOrder.verify(mockDeps).sendMessageDelayed(any(), any(), anyLong());
+
+        final MdnsPacket combinedSubtypesQuery = MdnsPacket.parse(
+                new MdnsPacketReader(combinedSubtypesQueryCaptor.getValue()));
+
+        assertEquals(3, combinedSubtypesQuery.questions.size());
+        assertTrue(hasQuestion(combinedSubtypesQuery, MdnsRecord.TYPE_PTR, SERVICE_TYPE_LABELS));
+        assertTrue(hasQuestion(combinedSubtypesQuery, MdnsRecord.TYPE_PTR,
+                getServiceTypeWithSubtype("_subtype1")));
+        assertTrue(hasQuestion(combinedSubtypesQuery, MdnsRecord.TYPE_PTR,
+                getServiceTypeWithSubtype("_subtype2")));
+
+        // Remove subtype1
+        stopSendAndReceive(mockListenerOne);
+
+        // Queries are not rescheduled, but the next query is affected
+        dispatchMessage();
+        currentThreadExecutor.getAndClearLastScheduledRunnable().run();
+
+        final ArgumentCaptor<DatagramPacket> subtype2QueryCaptor =
+                ArgumentCaptor.forClass(DatagramPacket.class);
+        // Send twice for IPv4 and IPv6
+        inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingMulticastResponse(
+                subtype2QueryCaptor.capture(),
+                eq(socketKey), eq(false));
+
+        final MdnsPacket subtype2Query = MdnsPacket.parse(
+                new MdnsPacketReader(subtype2QueryCaptor.getValue()));
+
+        assertEquals(2, subtype2Query.questions.size());
+        assertTrue(hasQuestion(subtype2Query, MdnsRecord.TYPE_PTR, SERVICE_TYPE_LABELS));
+        assertTrue(hasQuestion(subtype2Query, MdnsRecord.TYPE_PTR,
+                getServiceTypeWithSubtype("_subtype2")));
+    }
+
     private static void verifyServiceInfo(MdnsServiceInfo serviceInfo, String serviceName,
             String[] serviceType, List<String> ipv4Addresses, List<String> ipv6Addresses, int port,
             List<String> subTypes, Map<String, String> attributes, SocketKey socketKey) {
@@ -945,15 +1021,7 @@
     public void processResponse_searchOptionsEnableServiceRemoval_shouldRemove()
             throws Exception {
         final String serviceInstanceName = "service-instance-1";
-        client =
-                new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                        mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
-                        serviceCache) {
-                    @Override
-                    MdnsPacketWriter createMdnsPacketWriter() {
-                        return mockPacketWriter;
-                    }
-                };
+        client = makeMdnsServiceTypeClient(mockPacketWriter);
         MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
                 .setRemoveExpiredService(true)
                 .setNumOfQueriesBeforeBackoff(Integer.MAX_VALUE)
@@ -991,15 +1059,7 @@
     public void processResponse_searchOptionsNotEnableServiceRemoval_shouldNotRemove()
             throws Exception {
         final String serviceInstanceName = "service-instance-1";
-        client =
-                new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                        mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
-                        serviceCache) {
-                    @Override
-                    MdnsPacketWriter createMdnsPacketWriter() {
-                        return mockPacketWriter;
-                    }
-                };
+        client = makeMdnsServiceTypeClient(mockPacketWriter);
         startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
         Runnable firstMdnsTask = currentThreadExecutor.getAndClearSubmittedRunnable();
 
@@ -1025,15 +1085,7 @@
             throws Exception {
         //MdnsConfigsFlagsImpl.removeServiceAfterTtlExpires.override(true);
         final String serviceInstanceName = "service-instance-1";
-        client =
-                new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                        mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
-                        serviceCache) {
-                    @Override
-                    MdnsPacketWriter createMdnsPacketWriter() {
-                        return mockPacketWriter;
-                    }
-                };
+        client = makeMdnsServiceTypeClient(mockPacketWriter);
         startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
         Runnable firstMdnsTask = currentThreadExecutor.getAndClearSubmittedRunnable();
 
@@ -1148,9 +1200,7 @@
 
     @Test
     public void testProcessResponse_Resolve() throws Exception {
-        client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
-                serviceCache);
+        client = makeMdnsServiceTypeClient(/* packetWriter= */ null);
 
         final String instanceName = "service-instance";
         final String[] hostname = new String[] { "testhost "};
@@ -1243,9 +1293,7 @@
 
     @Test
     public void testRenewTxtSrvInResolve() throws Exception {
-        client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
-                serviceCache);
+        client = makeMdnsServiceTypeClient(/* packetWriter= */ null);
 
         final String instanceName = "service-instance";
         final String[] hostname = new String[] { "testhost "};
@@ -1359,9 +1407,7 @@
 
     @Test
     public void testProcessResponse_ResolveExcludesOtherServices() {
-        client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
-                serviceCache);
+        client = makeMdnsServiceTypeClient(/* packetWriter= */ null);
 
         final String requestedInstance = "instance1";
         final String otherInstance = "instance2";
@@ -1429,9 +1475,7 @@
 
     @Test
     public void testProcessResponse_SubtypeDiscoveryLimitedToSubtype() {
-        client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
-                serviceCache);
+        client = makeMdnsServiceTypeClient(/* packetWriter= */ null);
 
         final String matchingInstance = "instance1";
         final String subtype = "_subtype";
@@ -1519,9 +1563,7 @@
 
     @Test
     public void testProcessResponse_SubtypeChange() {
-        client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
-                serviceCache);
+        client = makeMdnsServiceTypeClient(/* packetWriter= */ null);
 
         final String matchingInstance = "instance1";
         final String subtype = "_subtype";
@@ -1604,9 +1646,7 @@
 
     @Test
     public void testNotifySocketDestroyed() throws Exception {
-        client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
-                serviceCache);
+        client = makeMdnsServiceTypeClient(/* packetWriter= */ null);
 
         final String requestedInstance = "instance1";
         final String otherInstance = "instance2";
@@ -1928,6 +1968,11 @@
                 Arrays.stream(SERVICE_TYPE_LABELS)).toArray(String[]::new);
     }
 
+    private static String[] getServiceTypeWithSubtype(String subtype) {
+        return Stream.concat(Stream.of(subtype, "_sub"),
+                Arrays.stream(SERVICE_TYPE_LABELS)).toArray(String[]::new);
+    }
+
     private static boolean hasQuestion(MdnsPacket packet, int type) {
         return hasQuestion(packet, type, null);
     }