• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /*
2  * Copyright (C) 2021 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License
15  */
16 
17 package com.android.server.connectivity.mdns;
18 
19 import static com.android.server.connectivity.mdns.MdnsQueryScheduler.INITIAL_AGGRESSIVE_TIME_BETWEEN_BURSTS_MS;
20 import static com.android.server.connectivity.mdns.MdnsQueryScheduler.MAX_TIME_BETWEEN_AGGRESSIVE_BURSTS_MS;
21 import static com.android.server.connectivity.mdns.MdnsQueryScheduler.TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS;
22 import static com.android.server.connectivity.mdns.MdnsSearchOptions.ACTIVE_QUERY_MODE;
23 import static com.android.server.connectivity.mdns.MdnsSearchOptions.AGGRESSIVE_QUERY_MODE;
24 import static com.android.server.connectivity.mdns.MdnsSearchOptions.PASSIVE_QUERY_MODE;
25 import static com.android.server.connectivity.mdns.MdnsServiceTypeClient.EVENT_START_QUERYTASK;
26 import static com.android.testutils.DevSdkIgnoreRuleKt.SC_V2;
27 
28 import static org.junit.Assert.assertArrayEquals;
29 import static org.junit.Assert.assertEquals;
30 import static org.junit.Assert.assertFalse;
31 import static org.junit.Assert.assertNotNull;
32 import static org.junit.Assert.assertNull;
33 import static org.junit.Assert.assertTrue;
34 import static org.mockito.ArgumentMatchers.any;
35 import static org.mockito.ArgumentMatchers.anyBoolean;
36 import static org.mockito.ArgumentMatchers.anyInt;
37 import static org.mockito.ArgumentMatchers.anyLong;
38 import static org.mockito.ArgumentMatchers.argThat;
39 import static org.mockito.ArgumentMatchers.eq;
40 import static org.mockito.Mockito.doAnswer;
41 import static org.mockito.Mockito.doCallRealMethod;
42 import static org.mockito.Mockito.doReturn;
43 import static org.mockito.Mockito.inOrder;
44 import static org.mockito.Mockito.never;
45 import static org.mockito.Mockito.times;
46 import static org.mockito.Mockito.verify;
47 import static org.mockito.Mockito.verifyNoMoreInteractions;
48 import static org.mockito.Mockito.when;
49 
50 import static java.nio.charset.StandardCharsets.UTF_8;
51 
52 import android.annotation.NonNull;
53 import android.annotation.Nullable;
54 import android.net.InetAddresses;
55 import android.net.Network;
56 import android.os.Handler;
57 import android.os.HandlerThread;
58 import android.os.Message;
59 import android.text.TextUtils;
60 
61 import com.android.net.module.util.CollectionUtils;
62 import com.android.net.module.util.SharedLog;
63 import com.android.server.connectivity.mdns.MdnsServiceInfo.TextEntry;
64 import com.android.server.connectivity.mdns.util.MdnsUtils;
65 import com.android.testutils.DevSdkIgnoreRule;
66 import com.android.testutils.DevSdkIgnoreRunner;
67 import com.android.testutils.HandlerUtils;
68 
69 import org.junit.After;
70 import org.junit.Before;
71 import org.junit.Ignore;
72 import org.junit.Test;
73 import org.junit.runner.RunWith;
74 import org.mockito.ArgumentCaptor;
75 import org.mockito.ArgumentMatcher;
76 import org.mockito.Captor;
77 import org.mockito.InOrder;
78 import org.mockito.Mock;
79 import org.mockito.Mockito;
80 import org.mockito.MockitoAnnotations;
81 
82 import java.io.IOException;
83 import java.net.DatagramPacket;
84 import java.net.InetAddress;
85 import java.net.InetSocketAddress;
86 import java.util.ArrayList;
87 import java.util.Arrays;
88 import java.util.Collections;
89 import java.util.List;
90 import java.util.Map;
91 import java.util.concurrent.Future;
92 import java.util.concurrent.ScheduledFuture;
93 import java.util.concurrent.ScheduledThreadPoolExecutor;
94 import java.util.concurrent.TimeUnit;
95 import java.util.stream.Stream;
96 
97 /** Tests for {@link MdnsServiceTypeClient}. */
98 @DevSdkIgnoreRunner.MonitorThreadLeak
99 @RunWith(DevSdkIgnoreRunner.class)
100 @DevSdkIgnoreRule.IgnoreUpTo(SC_V2)
101 public class MdnsServiceTypeClientTests {
102     private static final int INTERFACE_INDEX = 999;
103     private static final long DEFAULT_TIMEOUT = 2000L;
104     private static final String SERVICE_TYPE = "_googlecast._tcp.local";
105     private static final String SUBTYPE = "_subtype";
106     private static final String[] SERVICE_TYPE_LABELS = TextUtils.split(SERVICE_TYPE, "\\.");
107     private static final InetSocketAddress IPV4_ADDRESS = new InetSocketAddress(
108             MdnsConstants.getMdnsIPv4Address(), MdnsConstants.MDNS_PORT);
109     private static final InetSocketAddress IPV6_ADDRESS = new InetSocketAddress(
110             MdnsConstants.getMdnsIPv6Address(), MdnsConstants.MDNS_PORT);
111 
112     private static final long TEST_TTL = 120000L;
113     private static final long TEST_ELAPSED_REALTIME = 123L;
114     private static final long TEST_TIMEOUT_MS = 10_000L;
115 
116     @Mock
117     private MdnsServiceBrowserListener mockListenerOne;
118     @Mock
119     private MdnsServiceBrowserListener mockListenerTwo;
120     @Mock
121     private MdnsMultinetworkSocketClient mockSocketClient;
122     @Mock
123     private Network mockNetwork;
124     @Mock
125     private MdnsUtils.Clock mockDecoderClock;
126     @Mock
127     private SharedLog mockSharedLog;
128     @Mock
129     private MdnsServiceTypeClient.Dependencies mockDeps;
130     @Mock
131     private Scheduler mockScheduler;
132     @Captor
133     private ArgumentCaptor<MdnsServiceInfo> serviceInfoCaptor;
134 
135     private final byte[] buf = new byte[10];
136 
137     private DatagramPacket[] expectedIPv4Packets;
138     private DatagramPacket[] expectedIPv6Packets;
139     private FakeExecutor currentThreadExecutor = new FakeExecutor();
140 
141     private MdnsServiceTypeClient client;
142     private SocketKey socketKey;
143     private HandlerThread thread;
144     private Handler handler;
145     private MdnsServiceCache serviceCache;
146     private long latestDelayMs = 0;
147     private Message delayMessage = null;
148     private Handler realHandler = null;
149     private MdnsFeatureFlags featureFlags = MdnsFeatureFlags.newBuilder().build();
150     private Message message = null;
151 
152     @Before
153     @SuppressWarnings("DoNotMock")
setUp()154     public void setUp() throws IOException {
155         MockitoAnnotations.initMocks(this);
156         doReturn(TEST_ELAPSED_REALTIME).when(mockDecoderClock).elapsedRealtime();
157 
158         expectedIPv4Packets = new DatagramPacket[24];
159         expectedIPv6Packets = new DatagramPacket[24];
160         socketKey = new SocketKey(mockNetwork, INTERFACE_INDEX);
161 
162         for (int i = 0; i < expectedIPv4Packets.length; ++i) {
163             expectedIPv4Packets[i] = new DatagramPacket(buf, 0 /* offset */, 5 /* length */,
164                     MdnsConstants.getMdnsIPv4Address(), MdnsConstants.MDNS_PORT);
165             expectedIPv6Packets[i] = new DatagramPacket(buf, 0 /* offset */, 5 /* length */,
166                     MdnsConstants.getMdnsIPv6Address(), MdnsConstants.MDNS_PORT);
167         }
168         when(mockDeps.getDatagramPacketsFromMdnsPacket(
169                 any(), any(MdnsPacket.class), eq(IPV4_ADDRESS), anyBoolean()))
170                 .thenReturn(List.of(expectedIPv4Packets[0]))
171                 .thenReturn(List.of(expectedIPv4Packets[1]))
172                 .thenReturn(List.of(expectedIPv4Packets[2]))
173                 .thenReturn(List.of(expectedIPv4Packets[3]))
174                 .thenReturn(List.of(expectedIPv4Packets[4]))
175                 .thenReturn(List.of(expectedIPv4Packets[5]))
176                 .thenReturn(List.of(expectedIPv4Packets[6]))
177                 .thenReturn(List.of(expectedIPv4Packets[7]))
178                 .thenReturn(List.of(expectedIPv4Packets[8]))
179                 .thenReturn(List.of(expectedIPv4Packets[9]))
180                 .thenReturn(List.of(expectedIPv4Packets[10]))
181                 .thenReturn(List.of(expectedIPv4Packets[11]))
182                 .thenReturn(List.of(expectedIPv4Packets[12]))
183                 .thenReturn(List.of(expectedIPv4Packets[13]))
184                 .thenReturn(List.of(expectedIPv4Packets[14]))
185                 .thenReturn(List.of(expectedIPv4Packets[15]))
186                 .thenReturn(List.of(expectedIPv4Packets[16]))
187                 .thenReturn(List.of(expectedIPv4Packets[17]))
188                 .thenReturn(List.of(expectedIPv4Packets[18]))
189                 .thenReturn(List.of(expectedIPv4Packets[19]))
190                 .thenReturn(List.of(expectedIPv4Packets[20]))
191                 .thenReturn(List.of(expectedIPv4Packets[21]))
192                 .thenReturn(List.of(expectedIPv4Packets[22]))
193                 .thenReturn(List.of(expectedIPv4Packets[23]));
194 
195         when(mockDeps.getDatagramPacketsFromMdnsPacket(
196                 any(), any(MdnsPacket.class), eq(IPV6_ADDRESS), anyBoolean()))
197                 .thenReturn(List.of(expectedIPv6Packets[0]))
198                 .thenReturn(List.of(expectedIPv6Packets[1]))
199                 .thenReturn(List.of(expectedIPv6Packets[2]))
200                 .thenReturn(List.of(expectedIPv6Packets[3]))
201                 .thenReturn(List.of(expectedIPv6Packets[4]))
202                 .thenReturn(List.of(expectedIPv6Packets[5]))
203                 .thenReturn(List.of(expectedIPv6Packets[6]))
204                 .thenReturn(List.of(expectedIPv6Packets[7]))
205                 .thenReturn(List.of(expectedIPv6Packets[8]))
206                 .thenReturn(List.of(expectedIPv6Packets[9]))
207                 .thenReturn(List.of(expectedIPv6Packets[10]))
208                 .thenReturn(List.of(expectedIPv6Packets[11]))
209                 .thenReturn(List.of(expectedIPv6Packets[12]))
210                 .thenReturn(List.of(expectedIPv6Packets[13]))
211                 .thenReturn(List.of(expectedIPv6Packets[14]))
212                 .thenReturn(List.of(expectedIPv6Packets[15]))
213                 .thenReturn(List.of(expectedIPv6Packets[16]))
214                 .thenReturn(List.of(expectedIPv6Packets[17]))
215                 .thenReturn(List.of(expectedIPv6Packets[18]))
216                 .thenReturn(List.of(expectedIPv6Packets[19]))
217                 .thenReturn(List.of(expectedIPv6Packets[20]))
218                 .thenReturn(List.of(expectedIPv6Packets[21]))
219                 .thenReturn(List.of(expectedIPv6Packets[22]))
220                 .thenReturn(List.of(expectedIPv6Packets[23]));
221 
222         thread = new HandlerThread("MdnsServiceTypeClientTests");
223         thread.start();
224         handler = new Handler(thread.getLooper());
225         serviceCache = new MdnsServiceCache(
226                 thread.getLooper(),
227                 MdnsFeatureFlags.newBuilder().setIsExpiredServicesRemovalEnabled(false).build(),
228                 mockDecoderClock);
229 
230         doAnswer(inv -> {
231             latestDelayMs = 0;
232             delayMessage = null;
233             return true;
234         }).when(mockDeps).removeMessages(any(Handler.class), eq(EVENT_START_QUERYTASK));
235 
236         doAnswer(inv -> {
237             realHandler = (Handler) inv.getArguments()[0];
238             delayMessage = (Message) inv.getArguments()[1];
239             latestDelayMs = (long) inv.getArguments()[2];
240             return true;
241         }).when(mockDeps).sendMessageDelayed(any(Handler.class), any(Message.class), anyLong());
242 
243         doAnswer(inv -> {
244             final Handler handler = (Handler) inv.getArguments()[0];
245             final Message message = (Message) inv.getArguments()[1];
246             runOnHandler(() -> handler.dispatchMessage(message));
247             return true;
248         }).when(mockDeps).sendMessage(any(Handler.class), any(Message.class));
249 
250         doAnswer(inv -> {
251             realHandler = (Handler) inv.getArguments()[0];
252             return mockScheduler;
253         }).when(mockDeps).createScheduler(any(Handler.class));
254 
255         doAnswer(inv -> {
256             message = (Message) inv.getArguments()[0];
257             latestDelayMs = (long) inv.getArguments()[1];
258             return null;
259         }).when(mockScheduler).sendDelayedMessage(any(), anyLong());
260 
261         client = makeMdnsServiceTypeClient(featureFlags);
262     }
263 
makeMdnsServiceTypeClient(MdnsFeatureFlags featureFlags)264     private MdnsServiceTypeClient makeMdnsServiceTypeClient(MdnsFeatureFlags featureFlags) {
265         return new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
266                 mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
267                 serviceCache, featureFlags);
268     }
269 
270     @After
tearDown()271     public void tearDown() throws Exception {
272         if (thread != null) {
273             thread.quitSafely();
274             thread.join();
275         }
276     }
277 
runOnHandler(Runnable r)278     private void runOnHandler(Runnable r) {
279         handler.post(r);
280         HandlerUtils.waitForIdle(handler, DEFAULT_TIMEOUT);
281     }
282 
startSendAndReceive(MdnsServiceBrowserListener listener, MdnsSearchOptions searchOptions)283     private void startSendAndReceive(MdnsServiceBrowserListener listener,
284             MdnsSearchOptions searchOptions) {
285         runOnHandler(() -> client.startSendAndReceive(listener, searchOptions));
286     }
287 
processResponse(MdnsPacket packet, SocketKey socketKey)288     private void processResponse(MdnsPacket packet, SocketKey socketKey) {
289         runOnHandler(() -> client.processResponse(packet, socketKey));
290     }
291 
stopSendAndReceive(MdnsServiceBrowserListener listener)292     private void stopSendAndReceive(MdnsServiceBrowserListener listener) {
293         runOnHandler(() -> client.stopSendAndReceive(listener));
294     }
295 
notifySocketDestroyed()296     private void notifySocketDestroyed() {
297         runOnHandler(() -> client.notifySocketDestroyed());
298     }
299 
dispatchMessage()300     private void dispatchMessage() {
301         runOnHandler(() -> realHandler.dispatchMessage(delayMessage));
302         delayMessage = null;
303     }
304 
305     @Test
sendQueries_activeScanMode()306     public void sendQueries_activeScanMode() {
307         MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
308                 .addSubtype(SUBTYPE).setQueryMode(ACTIVE_QUERY_MODE).build();
309         startSendAndReceive(mockListenerOne, searchOptions);
310         // Always try to remove the task.
311         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
312 
313         // First burst, 3 queries.
314         verifyAndSendQuery(0, 0, /* expectsUnicastResponse= */ true);
315         verifyAndSendQuery(
316                 1, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
317         verifyAndSendQuery(
318                 2, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
319         // Second burst will be sent after initialTimeBetweenBurstsMs, 3 queries.
320         verifyAndSendQuery(
321                 3, MdnsConfigs.initialTimeBetweenBurstsMs(), /* expectsUnicastResponse= */ false);
322         verifyAndSendQuery(
323                 4, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
324         verifyAndSendQuery(
325                 5, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
326         // Third burst will be sent after initialTimeBetweenBurstsMs * 2, 3 queries.
327         verifyAndSendQuery(
328                 6, MdnsConfigs.initialTimeBetweenBurstsMs() * 2, /* expectsUnicastResponse= */
329                 false);
330         verifyAndSendQuery(
331                 7, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
332         verifyAndSendQuery(
333                 8, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
334         // Forth burst will be sent after initialTimeBetweenBurstsMs * 4, 3 queries.
335         verifyAndSendQuery(
336                 9, MdnsConfigs.initialTimeBetweenBurstsMs() * 4, /* expectsUnicastResponse= */
337                 false);
338         verifyAndSendQuery(
339                 10, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
340         verifyAndSendQuery(
341                 11, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
342         // Fifth burst will be sent after timeBetweenBurstsMs, 3 queries.
343         verifyAndSendQuery(12, MdnsConfigs.timeBetweenBurstsMs(), /* expectsUnicastResponse= */
344                 false);
345         verifyAndSendQuery(
346                 13, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
347         verifyAndSendQuery(
348                 14, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
349         // Verify that Task is not removed before stopSendAndReceive was called.
350         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
351 
352         // Stop sending packets.
353         stopSendAndReceive(mockListenerOne);
354         verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
355     }
356 
357     @Test
sendQueries_reentry_activeScanMode()358     public void sendQueries_reentry_activeScanMode() {
359         MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
360                 .addSubtype(SUBTYPE).setQueryMode(ACTIVE_QUERY_MODE).build();
361         startSendAndReceive(mockListenerOne, searchOptions);
362         // Always try to remove the task.
363         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
364 
365         // First burst, first query is sent.
366         verifyAndSendQuery(0, 0, /* expectsUnicastResponse= */ true);
367 
368         // After the first query is sent, change the subtypes, and restart.
369         searchOptions =
370                 MdnsSearchOptions.newBuilder()
371                         .addSubtype(SUBTYPE)
372                         .addSubtype("_subtype2")
373                         .setQueryMode(ACTIVE_QUERY_MODE)
374                         .build();
375         startSendAndReceive(mockListenerOne, searchOptions);
376         // The previous scheduled task should be canceled.
377         verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
378 
379         // Queries should continue to be sent.
380         verifyAndSendQuery(1, 0, /* expectsUnicastResponse= */ true);
381         verifyAndSendQuery(
382                 2, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
383         verifyAndSendQuery(
384                 3, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
385 
386         // Stop sending packets.
387         stopSendAndReceive(mockListenerOne);
388         verify(mockDeps, times(3)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
389     }
390 
391     @Test
sendQueries_passiveScanMode()392     public void sendQueries_passiveScanMode() {
393         MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
394                 .addSubtype(SUBTYPE).setQueryMode(PASSIVE_QUERY_MODE).build();
395         startSendAndReceive(mockListenerOne, searchOptions);
396         // Always try to remove the task.
397         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
398 
399         // First burst, 3 query.
400         verifyAndSendQuery(0, 0, /* expectsUnicastResponse= */ true);
401         verifyAndSendQuery(
402                 1, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
403         verifyAndSendQuery(
404                 2, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
405         // Second burst will be sent after timeBetweenBurstsMs, 1 query.
406         verifyAndSendQuery(3, MdnsConfigs.timeBetweenBurstsMs(), /* expectsUnicastResponse= */
407                 false);
408         // Third burst will be sent after timeBetweenBurstsMs, 1 query.
409         verifyAndSendQuery(4, MdnsConfigs.timeBetweenBurstsMs(), /* expectsUnicastResponse= */
410                 false);
411 
412         // Stop sending packets.
413         stopSendAndReceive(mockListenerOne);
414         verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
415     }
416 
417     @Test
sendQueries_activeScanWithQueryBackoff()418     public void sendQueries_activeScanWithQueryBackoff() {
419         MdnsSearchOptions searchOptions =
420                 MdnsSearchOptions.newBuilder()
421                         .addSubtype(SUBTYPE)
422                         .setQueryMode(ACTIVE_QUERY_MODE)
423                         .setNumOfQueriesBeforeBackoff(11).build();
424         startSendAndReceive(mockListenerOne, searchOptions);
425         // Always try to remove the task.
426         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
427 
428         // First burst, 3 queries.
429         verifyAndSendQuery(0, 0, /* expectsUnicastResponse= */ true);
430         verifyAndSendQuery(
431                 1, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
432         verifyAndSendQuery(
433                 2, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
434         // Second burst will be sent after initialTimeBetweenBurstsMs, 3 queries.
435         verifyAndSendQuery(
436                 3, MdnsConfigs.initialTimeBetweenBurstsMs(), /* expectsUnicastResponse= */ false);
437         verifyAndSendQuery(
438                 4, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
439         verifyAndSendQuery(
440                 5, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
441         // Third burst will be sent after initialTimeBetweenBurstsMs * 2, 3 queries.
442         verifyAndSendQuery(
443                 6, MdnsConfigs.initialTimeBetweenBurstsMs() * 2, /* expectsUnicastResponse= */
444                 false);
445         verifyAndSendQuery(
446                 7, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
447         verifyAndSendQuery(
448                 8, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
449         // Forth burst will be sent after initialTimeBetweenBurstsMs * 4, 3 queries.
450         verifyAndSendQuery(
451                 9, MdnsConfigs.initialTimeBetweenBurstsMs() * 4, /* expectsUnicastResponse= */
452                 false);
453         verifyAndSendQuery(
454                 10, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
455         verifyAndSendQuery(
456                 11, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
457         // In backoff mode, the current scheduled task will be canceled and reschedule if the
458         // 0.8 * smallestRemainingTtl is larger than time to next run.
459         long currentTime = TEST_TTL / 2 + TEST_ELAPSED_REALTIME;
460         doReturn(currentTime).when(mockDecoderClock).elapsedRealtime();
461         doReturn(true).when(mockDeps).hasMessages(any(), eq(EVENT_START_QUERYTASK));
462         processResponse(createResponse(
463                 "service-instance-1", "192.0.2.123", 5353,
464                 SERVICE_TYPE_LABELS,
465                 Collections.emptyMap(), TEST_TTL), socketKey);
466         verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
467         assertNotNull(delayMessage);
468         verifyAndSendQuery(12 /* index */, (long) (TEST_TTL / 2 * 0.8) /* timeInMs */,
469                 false /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
470                 14 /* scheduledCount */);
471         currentTime += (long) (TEST_TTL / 2 * 0.8);
472         doReturn(currentTime).when(mockDecoderClock).elapsedRealtime();
473         verifyAndSendQuery(13 /* index */, MdnsConfigs.timeBetweenQueriesInBurstMs(),
474                 false /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
475                 15 /* scheduledCount */);
476     }
477 
478     @Test
sendQueries_passiveScanWithQueryBackoff()479     public void sendQueries_passiveScanWithQueryBackoff() {
480         MdnsSearchOptions searchOptions =
481                 MdnsSearchOptions.newBuilder()
482                         .addSubtype(SUBTYPE)
483                         .setQueryMode(PASSIVE_QUERY_MODE)
484                         .setNumOfQueriesBeforeBackoff(3).build();
485         startSendAndReceive(mockListenerOne, searchOptions);
486         // Always try to remove the task.
487         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
488 
489         verifyAndSendQuery(0 /* index */, 0 /* timeInMs */, true /* expectsUnicastResponse */,
490                 true /* multipleSocketDiscovery */, 1 /* scheduledCount */);
491         verifyAndSendQuery(1 /* index */, MdnsConfigs.timeBetweenQueriesInBurstMs(),
492                 false /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
493                 2 /* scheduledCount */);
494         verifyAndSendQuery(2 /* index */, MdnsConfigs.timeBetweenQueriesInBurstMs(),
495                 false /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
496                 3 /* scheduledCount */);
497         verifyAndSendQuery(3 /* index */, MdnsConfigs.timeBetweenBurstsMs(),
498                 false /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
499                 4 /* scheduledCount */);
500 
501         // In backoff mode, the current scheduled task will be canceled and reschedule if the
502         // 0.8 * smallestRemainingTtl is larger than time to next run.
503         doReturn(TEST_ELAPSED_REALTIME + 20000).when(mockDecoderClock).elapsedRealtime();
504         doReturn(true).when(mockDeps).hasMessages(any(), eq(EVENT_START_QUERYTASK));
505         processResponse(createResponse(
506                 "service-instance-1", "192.0.2.123", 5353,
507                 SERVICE_TYPE_LABELS,
508                 Collections.emptyMap(), TEST_TTL), socketKey);
509         verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
510         assertNotNull(delayMessage);
511         verifyAndSendQuery(4 /* index */, 80000 /* timeInMs */, false /* expectsUnicastResponse */,
512                 true /* multipleSocketDiscovery */, 6 /* scheduledCount */);
513         // Next run should also be scheduled in 0.8 * smallestRemainingTtl
514         verifyAndSendQuery(5 /* index */, 80000 /* timeInMs */, false /* expectsUnicastResponse */,
515                 true /* multipleSocketDiscovery */, 7 /* scheduledCount */);
516 
517         // If the records is not refreshed, the current scheduled task will not be canceled.
518         doReturn(TEST_ELAPSED_REALTIME + 20001).when(mockDecoderClock).elapsedRealtime();
519         processResponse(createResponse(
520                 "service-instance-1", "192.0.2.123", 5353,
521                 SERVICE_TYPE_LABELS,
522                 Collections.emptyMap(), TEST_TTL,
523                 TEST_ELAPSED_REALTIME - 1), socketKey);
524         verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
525 
526         // In backoff mode, the current scheduled task will not be canceled if the
527         // 0.8 * smallestRemainingTtl is smaller than time to next run.
528         doReturn(TEST_ELAPSED_REALTIME).when(mockDecoderClock).elapsedRealtime();
529         processResponse(createResponse(
530                 "service-instance-1", "192.0.2.123", 5353,
531                 SERVICE_TYPE_LABELS,
532                 Collections.emptyMap(), TEST_TTL), socketKey);
533         verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
534 
535         stopSendAndReceive(mockListenerOne);
536         verify(mockDeps, times(3)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
537     }
538 
539     @Test
sendQueries_reentry_passiveScanMode()540     public void sendQueries_reentry_passiveScanMode() {
541         MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
542                 .addSubtype(SUBTYPE).setQueryMode(PASSIVE_QUERY_MODE).build();
543         startSendAndReceive(mockListenerOne, searchOptions);
544         // Always try to remove the task.
545         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
546 
547         // First burst, first query is sent.
548         verifyAndSendQuery(0, 0, /* expectsUnicastResponse= */ true);
549 
550         // After the first query is sent, change the subtypes, and restart.
551         searchOptions =
552                 MdnsSearchOptions.newBuilder()
553                         .addSubtype(SUBTYPE)
554                         .addSubtype("_subtype2")
555                         .setQueryMode(PASSIVE_QUERY_MODE)
556                         .build();
557         startSendAndReceive(mockListenerOne, searchOptions);
558         // The previous scheduled task should be canceled.
559         verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
560 
561         // Queries should continue to be sent.
562         verifyAndSendQuery(1, 0, /* expectsUnicastResponse= */ true);
563         verifyAndSendQuery(
564                 2, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
565         verifyAndSendQuery(
566                 3, MdnsConfigs.timeBetweenQueriesInBurstMs(), /* expectsUnicastResponse= */ false);
567 
568         // Stop sending packets.
569         stopSendAndReceive(mockListenerOne);
570         verify(mockDeps, times(3)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
571     }
572 
573     @Test
574     @Ignore("MdnsConfigs is not configurable currently.")
testQueryTaskConfig_alwaysAskForUnicastResponse()575     public void testQueryTaskConfig_alwaysAskForUnicastResponse() {
576         //MdnsConfigsFlagsImpl.alwaysAskForUnicastResponseInEachBurst.override(true);
577         MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
578                 .addSubtype(SUBTYPE).setQueryMode(ACTIVE_QUERY_MODE).build();
579         QueryTaskConfig config = new QueryTaskConfig(searchOptions.getQueryMode());
580 
581         // This is the first query. We will ask for unicast response.
582         assertTrue(config.expectUnicastResponse);
583         assertEquals(config.getTransactionId(), 1);
584 
585         // For the rest of queries in this burst, we will NOT ask for unicast response.
586         for (int i = 1; i < MdnsConfigs.queriesPerBurst(); i++) {
587             int oldTransactionId = config.getTransactionId();
588             config = config.getConfigForNextRun(ACTIVE_QUERY_MODE);
589             assertFalse(config.expectUnicastResponse);
590             assertEquals(config.getTransactionId(), oldTransactionId + 1);
591         }
592 
593         // This is the first query of a new burst. We will ask for unicast response.
594         int oldTransactionId = config.getTransactionId();
595         config = config.getConfigForNextRun(ACTIVE_QUERY_MODE);
596         assertTrue(config.expectUnicastResponse);
597         assertEquals(config.getTransactionId(), oldTransactionId + 1);
598     }
599 
600     @Test
testQueryTaskConfig_askForUnicastInFirstQuery()601     public void testQueryTaskConfig_askForUnicastInFirstQuery() {
602         MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
603                 .addSubtype(SUBTYPE).setQueryMode(ACTIVE_QUERY_MODE).build();
604         QueryTaskConfig config = new QueryTaskConfig(searchOptions.getQueryMode());
605 
606         // This is the first query. We will ask for unicast response.
607         assertTrue(config.expectUnicastResponse);
608         assertEquals(config.getTransactionId(), 1);
609 
610         // For the rest of queries in this burst, we will NOT ask for unicast response.
611         for (int i = 1; i < MdnsConfigs.queriesPerBurst(); i++) {
612             int oldTransactionId = config.getTransactionId();
613             config = config.getConfigForNextRun(ACTIVE_QUERY_MODE);
614             assertFalse(config.expectUnicastResponse);
615             assertEquals(config.getTransactionId(), oldTransactionId + 1);
616         }
617 
618         // This is the first query of a new burst. We will NOT ask for unicast response.
619         int oldTransactionId = config.getTransactionId();
620         config = config.getConfigForNextRun(ACTIVE_QUERY_MODE);
621         assertFalse(config.expectUnicastResponse);
622         assertEquals(config.getTransactionId(), oldTransactionId + 1);
623     }
624 
625     @Test
testIfPreviousTaskIsCanceledWhenNewSessionStarts()626     public void testIfPreviousTaskIsCanceledWhenNewSessionStarts() {
627         MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
628                 .addSubtype(SUBTYPE).setQueryMode(PASSIVE_QUERY_MODE).build();
629         startSendAndReceive(mockListenerOne, searchOptions);
630         Runnable firstMdnsTask = currentThreadExecutor.getAndClearSubmittedRunnable();
631 
632         // Change the sutypes and start a new session.
633         searchOptions =
634                 MdnsSearchOptions.newBuilder()
635                         .addSubtype(SUBTYPE)
636                         .addSubtype("_subtype2")
637                         .setQueryMode(PASSIVE_QUERY_MODE)
638                         .build();
639         startSendAndReceive(mockListenerOne, searchOptions);
640 
641         // Clear the scheduled runnable.
642         currentThreadExecutor.getAndClearLastScheduledRunnable();
643 
644         // Simulate the case where the first mdns task is not successful canceled and it gets
645         // executed anyway.
646         firstMdnsTask.run();
647 
648         // Although it gets executes, no more task gets scheduled.
649         assertNull(currentThreadExecutor.getAndClearLastScheduledRunnable());
650     }
651 
652     @Test
653     @Ignore("MdnsConfigs is not configurable currently.")
testIfPreviousTaskIsCanceledWhenSessionStops()654     public void testIfPreviousTaskIsCanceledWhenSessionStops() {
655         //MdnsConfigsFlagsImpl.shouldCancelScanTaskWhenFutureIsNull.override(true);
656         MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
657                 .addSubtype(SUBTYPE).setQueryMode(PASSIVE_QUERY_MODE).build();
658         startSendAndReceive(mockListenerOne, searchOptions);
659         // Change the sutypes and start a new session.
660         stopSendAndReceive(mockListenerOne);
661         // Clear the scheduled runnable.
662         currentThreadExecutor.getAndClearLastScheduledRunnable();
663 
664         // Simulate the case where the first mdns task is not successful canceled and it gets
665         // executed anyway.
666         currentThreadExecutor.getAndClearSubmittedRunnable().run();
667 
668         // Although it gets executes, no more task gets scheduled.
669         assertNull(currentThreadExecutor.getAndClearLastScheduledRunnable());
670     }
671 
672     @Test
testQueryScheduledWhenAnsweredFromCache()673     public void testQueryScheduledWhenAnsweredFromCache() {
674         final MdnsSearchOptions searchOptions = MdnsSearchOptions.getDefaultOptions();
675         startSendAndReceive(mockListenerOne, searchOptions);
676         assertNotNull(currentThreadExecutor.getAndClearSubmittedRunnable());
677 
678         processResponse(createResponse(
679                 "service-instance-1", "192.0.2.123", 5353,
680                 SERVICE_TYPE_LABELS,
681                 Collections.emptyMap(), TEST_TTL), socketKey);
682 
683         verify(mockListenerOne).onServiceNameDiscovered(any(), eq(false) /* isServiceFromCache */);
684         verify(mockListenerOne).onServiceFound(any(), eq(false) /* isServiceFromCache */);
685 
686         // File another identical query
687         startSendAndReceive(mockListenerTwo, searchOptions);
688 
689         verify(mockListenerTwo).onServiceNameDiscovered(any(), eq(true) /* isServiceFromCache */);
690         verify(mockListenerTwo).onServiceFound(any(), eq(true) /* isServiceFromCache */);
691 
692         // This time no query is submitted, only scheduled
693         assertNull(currentThreadExecutor.getAndClearSubmittedRunnable());
694         // This just skips the first query of the first burst
695         verify(mockDeps).sendMessageDelayed(
696                 any(), any(), eq(MdnsConfigs.timeBetweenQueriesInBurstMs()));
697     }
698 
699     @Test
testCombinedSubtypesQueriedWithMultipleListeners()700     public void testCombinedSubtypesQueriedWithMultipleListeners() throws Exception {
701         final MdnsSearchOptions searchOptions1 = MdnsSearchOptions.newBuilder()
702                 .addSubtype("subtype1").build();
703         final MdnsSearchOptions searchOptions2 = MdnsSearchOptions.newBuilder()
704                 .addSubtype("subtype2").build();
705         doCallRealMethod().when(mockDeps).getDatagramPacketsFromMdnsPacket(
706                 any(), any(MdnsPacket.class), any(InetSocketAddress.class), anyBoolean());
707         startSendAndReceive(mockListenerOne, searchOptions1);
708         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
709 
710         InOrder inOrder = inOrder(mockListenerOne, mockSocketClient, mockDeps);
711 
712         // Verify the query asks for subtype1
713         final ArgumentCaptor<List<DatagramPacket>> subtype1QueryCaptor =
714                 ArgumentCaptor.forClass(List.class);
715         // Send twice for IPv4 and IPv6
716         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingUnicastResponse(
717                 subtype1QueryCaptor.capture(),
718                 eq(socketKey), eq(false));
719 
720         final MdnsPacket subtype1Query = MdnsPacket.parse(
721                 new MdnsPacketReader(subtype1QueryCaptor.getValue().get(0)));
722 
723         assertEquals(2, subtype1Query.questions.size());
724         assertTrue(hasQuestion(subtype1Query, MdnsRecord.TYPE_PTR, SERVICE_TYPE_LABELS));
725         assertTrue(hasQuestion(subtype1Query, MdnsRecord.TYPE_PTR,
726                 getServiceTypeWithSubtype("_subtype1")));
727 
728         // Add subtype2
729         startSendAndReceive(mockListenerTwo, searchOptions2);
730         inOrder.verify(mockDeps).removeMessages(any(), eq(EVENT_START_QUERYTASK));
731         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
732 
733         final ArgumentCaptor<List<DatagramPacket>> combinedSubtypesQueryCaptor =
734                 ArgumentCaptor.forClass(List.class);
735         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingUnicastResponse(
736                 combinedSubtypesQueryCaptor.capture(),
737                 eq(socketKey), eq(false));
738         // The next query must have been scheduled
739         inOrder.verify(mockDeps).sendMessageDelayed(any(), any(), anyLong());
740 
741         final MdnsPacket combinedSubtypesQuery = MdnsPacket.parse(
742                 new MdnsPacketReader(combinedSubtypesQueryCaptor.getValue().get(0)));
743 
744         assertEquals(3, combinedSubtypesQuery.questions.size());
745         assertTrue(hasQuestion(combinedSubtypesQuery, MdnsRecord.TYPE_PTR, SERVICE_TYPE_LABELS));
746         assertTrue(hasQuestion(combinedSubtypesQuery, MdnsRecord.TYPE_PTR,
747                 getServiceTypeWithSubtype("_subtype1")));
748         assertTrue(hasQuestion(combinedSubtypesQuery, MdnsRecord.TYPE_PTR,
749                 getServiceTypeWithSubtype("_subtype2")));
750 
751         // Remove subtype1
752         stopSendAndReceive(mockListenerOne);
753 
754         // Queries are not rescheduled, but the next query is affected
755         dispatchMessage();
756         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
757 
758         final ArgumentCaptor<List<DatagramPacket>> subtype2QueryCaptor =
759                 ArgumentCaptor.forClass(List.class);
760         // Send twice for IPv4 and IPv6
761         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingMulticastResponse(
762                 subtype2QueryCaptor.capture(),
763                 eq(socketKey), eq(false));
764 
765         final MdnsPacket subtype2Query = MdnsPacket.parse(
766                 new MdnsPacketReader(subtype2QueryCaptor.getValue().get(0)));
767 
768         assertEquals(2, subtype2Query.questions.size());
769         assertTrue(hasQuestion(subtype2Query, MdnsRecord.TYPE_PTR, SERVICE_TYPE_LABELS));
770         assertTrue(hasQuestion(subtype2Query, MdnsRecord.TYPE_PTR,
771                 getServiceTypeWithSubtype("_subtype2")));
772     }
773 
verifyServiceInfo(MdnsServiceInfo serviceInfo, String serviceName, String[] serviceType, List<String> ipv4Addresses, List<String> ipv6Addresses, int port, List<String> subTypes, Map<String, String> attributes, SocketKey socketKey)774     private static void verifyServiceInfo(MdnsServiceInfo serviceInfo, String serviceName,
775             String[] serviceType, List<String> ipv4Addresses, List<String> ipv6Addresses, int port,
776             List<String> subTypes, Map<String, String> attributes, SocketKey socketKey) {
777         assertEquals(serviceName, serviceInfo.getServiceInstanceName());
778         assertArrayEquals(serviceType, serviceInfo.getServiceType());
779         assertEquals(ipv4Addresses, serviceInfo.getIpv4Addresses());
780         assertEquals(ipv6Addresses, serviceInfo.getIpv6Addresses());
781         assertEquals(port, serviceInfo.getPort());
782         assertEquals(subTypes, serviceInfo.getSubtypes());
783         for (String key : attributes.keySet()) {
784             assertTrue(attributes.containsKey(key));
785             assertEquals(attributes.get(key), serviceInfo.getAttributeByKey(key));
786         }
787         assertEquals(socketKey.getInterfaceIndex(), serviceInfo.getInterfaceIndex());
788         assertEquals(socketKey.getNetwork(), serviceInfo.getNetwork());
789     }
790 
791     @Test
processResponse_incompleteResponse()792     public void processResponse_incompleteResponse() {
793         startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
794 
795         processResponse(createResponse(
796                 "service-instance-1", null /* host */, 0 /* port */,
797                 SERVICE_TYPE_LABELS,
798                 Collections.emptyMap(), TEST_TTL), socketKey);
799         verify(mockListenerOne).onServiceNameDiscovered(
800                 serviceInfoCaptor.capture(), eq(false) /* isServiceFromCache */);
801         verifyServiceInfo(serviceInfoCaptor.getAllValues().get(0),
802                 "service-instance-1",
803                 SERVICE_TYPE_LABELS,
804                 /* ipv4Addresses= */ List.of(),
805                 /* ipv6Addresses= */ List.of(),
806                 /* port= */ 0,
807                 /* subTypes= */ List.of(),
808                 Collections.emptyMap(),
809                 socketKey);
810 
811         verify(mockListenerOne, never()).onServiceFound(any(MdnsServiceInfo.class), anyBoolean());
812         verify(mockListenerOne, never()).onServiceUpdated(any(MdnsServiceInfo.class));
813     }
814 
815     @Test
processIPv4Response_completeResponseForNewServiceInstance()816     public void processIPv4Response_completeResponseForNewServiceInstance() throws Exception {
817         final String ipV4Address = "192.168.1.1";
818         startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
819 
820         // Process the initial response.
821         processResponse(createResponse(
822                 "service-instance-1", ipV4Address, 5353, SUBTYPE,
823                 Collections.emptyMap(), TEST_TTL), socketKey);
824 
825         // Process a second response with a different port and updated text attributes.
826         processResponse(createResponse(
827                         "service-instance-1", ipV4Address, 5354, SUBTYPE,
828                         Collections.singletonMap("key", "value"), TEST_TTL),
829                 socketKey);
830 
831         // Verify onServiceNameDiscovered was called once for the initial response.
832         verify(mockListenerOne).onServiceNameDiscovered(
833                 serviceInfoCaptor.capture(), eq(false) /* isServiceFromCache */);
834         verifyServiceInfo(serviceInfoCaptor.getAllValues().get(0),
835                 "service-instance-1",
836                 SERVICE_TYPE_LABELS,
837                 List.of(ipV4Address) /* ipv4Address */,
838                 List.of() /* ipv6Address */,
839                 5353 /* port */,
840                 Collections.singletonList(SUBTYPE) /* subTypes */,
841                 Collections.singletonMap("key", null) /* attributes */,
842                 socketKey);
843 
844         // Verify onServiceFound was called once for the initial response.
845         verify(mockListenerOne).onServiceFound(
846                 serviceInfoCaptor.capture(), eq(false) /* isServiceFromCache */);
847         MdnsServiceInfo initialServiceInfo = serviceInfoCaptor.getAllValues().get(1);
848         assertEquals(initialServiceInfo.getServiceInstanceName(), "service-instance-1");
849         assertEquals(initialServiceInfo.getIpv4Address(), ipV4Address);
850         assertEquals(initialServiceInfo.getPort(), 5353);
851         assertEquals(initialServiceInfo.getSubtypes(), Collections.singletonList(SUBTYPE));
852         assertNull(initialServiceInfo.getAttributeByKey("key"));
853         assertEquals(socketKey.getInterfaceIndex(), initialServiceInfo.getInterfaceIndex());
854         assertEquals(socketKey.getNetwork(), initialServiceInfo.getNetwork());
855 
856         // Verify onServiceUpdated was called once for the second response.
857         verify(mockListenerOne).onServiceUpdated(serviceInfoCaptor.capture());
858         MdnsServiceInfo updatedServiceInfo = serviceInfoCaptor.getAllValues().get(2);
859         assertEquals(updatedServiceInfo.getServiceInstanceName(), "service-instance-1");
860         assertEquals(updatedServiceInfo.getIpv4Address(), ipV4Address);
861         assertEquals(updatedServiceInfo.getPort(), 5354);
862         assertTrue(updatedServiceInfo.hasSubtypes());
863         assertEquals(updatedServiceInfo.getSubtypes(), Collections.singletonList(SUBTYPE));
864         assertEquals(updatedServiceInfo.getAttributeByKey("key"), "value");
865         assertEquals(socketKey.getInterfaceIndex(), updatedServiceInfo.getInterfaceIndex());
866         assertEquals(socketKey.getNetwork(), updatedServiceInfo.getNetwork());
867     }
868 
869     @Test
processIPv6Response_getCorrectServiceInfo()870     public void processIPv6Response_getCorrectServiceInfo() throws Exception {
871         final String ipV6Address = "2000:3333::da6c:63ff:fe7c:7483";
872         startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
873 
874         // Process the initial response.
875         processResponse(createResponse(
876                 "service-instance-1", ipV6Address, 5353, SUBTYPE,
877                 Collections.emptyMap(), TEST_TTL), socketKey);
878 
879         // Process a second response with a different port and updated text attributes.
880         processResponse(createResponse(
881                         "service-instance-1", ipV6Address, 5354, SUBTYPE,
882                         Collections.singletonMap("key", "value"), TEST_TTL),
883                 socketKey);
884 
885         // Verify onServiceNameDiscovered was called once for the initial response.
886         verify(mockListenerOne).onServiceNameDiscovered(
887                 serviceInfoCaptor.capture(), eq(false) /* isServiceFromCache */);
888         verifyServiceInfo(serviceInfoCaptor.getAllValues().get(0),
889                 "service-instance-1",
890                 SERVICE_TYPE_LABELS,
891                 List.of() /* ipv4Address */,
892                 List.of(ipV6Address) /* ipv6Address */,
893                 5353 /* port */,
894                 Collections.singletonList(SUBTYPE) /* subTypes */,
895                 Collections.singletonMap("key", null) /* attributes */,
896                 socketKey);
897 
898         // Verify onServiceFound was called once for the initial response.
899         verify(mockListenerOne).onServiceFound(
900                 serviceInfoCaptor.capture(), eq(false) /* isServiceFromCache */);
901         MdnsServiceInfo initialServiceInfo = serviceInfoCaptor.getAllValues().get(1);
902         assertEquals(initialServiceInfo.getServiceInstanceName(), "service-instance-1");
903         assertEquals(initialServiceInfo.getIpv6Address(), ipV6Address);
904         assertEquals(initialServiceInfo.getPort(), 5353);
905         assertEquals(initialServiceInfo.getSubtypes(), Collections.singletonList(SUBTYPE));
906         assertNull(initialServiceInfo.getAttributeByKey("key"));
907         assertEquals(socketKey.getInterfaceIndex(), initialServiceInfo.getInterfaceIndex());
908         assertEquals(socketKey.getNetwork(), initialServiceInfo.getNetwork());
909 
910         // Verify onServiceUpdated was called once for the second response.
911         verify(mockListenerOne).onServiceUpdated(serviceInfoCaptor.capture());
912         MdnsServiceInfo updatedServiceInfo = serviceInfoCaptor.getAllValues().get(2);
913         assertEquals(updatedServiceInfo.getServiceInstanceName(), "service-instance-1");
914         assertEquals(updatedServiceInfo.getIpv6Address(), ipV6Address);
915         assertEquals(updatedServiceInfo.getPort(), 5354);
916         assertTrue(updatedServiceInfo.hasSubtypes());
917         assertEquals(updatedServiceInfo.getSubtypes(), Collections.singletonList(SUBTYPE));
918         assertEquals(updatedServiceInfo.getAttributeByKey("key"), "value");
919         assertEquals(socketKey.getInterfaceIndex(), updatedServiceInfo.getInterfaceIndex());
920         assertEquals(socketKey.getNetwork(), updatedServiceInfo.getNetwork());
921     }
922 
verifyServiceRemovedNoCallback(MdnsServiceBrowserListener listener)923     private void verifyServiceRemovedNoCallback(MdnsServiceBrowserListener listener) {
924         verify(listener, never()).onServiceRemoved(any());
925         verify(listener, never()).onServiceNameRemoved(any());
926     }
927 
verifyServiceRemovedCallback(MdnsServiceBrowserListener listener, String serviceName, String[] serviceType, SocketKey socketKey)928     private void verifyServiceRemovedCallback(MdnsServiceBrowserListener listener,
929             String serviceName, String[] serviceType, SocketKey socketKey) {
930         verify(listener).onServiceRemoved(argThat(
931                 info -> serviceName.equals(info.getServiceInstanceName())
932                         && Arrays.equals(serviceType, info.getServiceType())
933                         && info.getInterfaceIndex() == socketKey.getInterfaceIndex()
934                         && socketKey.getNetwork().equals(info.getNetwork())));
935         verify(listener).onServiceNameRemoved(argThat(
936                 info -> serviceName.equals(info.getServiceInstanceName())
937                         && Arrays.equals(serviceType, info.getServiceType())
938                         && info.getInterfaceIndex() == socketKey.getInterfaceIndex()
939                         && socketKey.getNetwork().equals(info.getNetwork())));
940     }
941 
942     @Test
processResponse_goodBye()943     public void processResponse_goodBye() throws Exception {
944         startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
945         startSendAndReceive(mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
946 
947         final String serviceName = "service-instance-1";
948         final String ipV6Address = "2000:3333::da6c:63ff:fe7c:7483";
949         // Process the initial response.
950         processResponse(createResponse(
951                 serviceName, ipV6Address, 5353,
952                 SERVICE_TYPE_LABELS,
953                 Collections.emptyMap(), TEST_TTL), socketKey);
954 
955         processResponse(createResponse(
956                 "goodbye-service", ipV6Address, 5353,
957                 SERVICE_TYPE_LABELS,
958                 Collections.emptyMap(), /* ptrTtlMillis= */ 0L), socketKey);
959 
960         // Verify removed callback won't be called if the service is not existed.
961         verifyServiceRemovedNoCallback(mockListenerOne);
962         verifyServiceRemovedNoCallback(mockListenerTwo);
963 
964         // Verify removed callback would be called.
965         processResponse(createResponse(
966                 serviceName, ipV6Address, 5353,
967                 SERVICE_TYPE_LABELS,
968                 Collections.emptyMap(), 0L), socketKey);
969         verifyServiceRemovedCallback(
970                 mockListenerOne, serviceName, SERVICE_TYPE_LABELS, socketKey);
971         verifyServiceRemovedCallback(
972                 mockListenerTwo, serviceName, SERVICE_TYPE_LABELS, socketKey);
973     }
974 
975     @Test
reportExistingServiceToNewlyRegisteredListeners()976     public void reportExistingServiceToNewlyRegisteredListeners() throws Exception {
977         // Process the initial response.
978         processResponse(createResponse(
979                 "service-instance-1", "192.168.1.1", 5353, SUBTYPE,
980                 Collections.emptyMap(), TEST_TTL), socketKey);
981 
982         startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
983 
984         // Verify onServiceNameDiscovered was called once for the existing response.
985         verify(mockListenerOne).onServiceNameDiscovered(
986                 serviceInfoCaptor.capture(), eq(true) /* isServiceFromCache */);
987         verifyServiceInfo(serviceInfoCaptor.getAllValues().get(0),
988                 "service-instance-1",
989                 SERVICE_TYPE_LABELS,
990                 List.of("192.168.1.1") /* ipv4Address */,
991                 List.of() /* ipv6Address */,
992                 5353 /* port */,
993                 Collections.singletonList(SUBTYPE) /* subTypes */,
994                 Collections.singletonMap("key", null) /* attributes */,
995                 socketKey);
996 
997         // Verify onServiceFound was called once for the existing response.
998         verify(mockListenerOne).onServiceFound(
999                 serviceInfoCaptor.capture(), eq(true) /* isServiceFromCache */);
1000         MdnsServiceInfo existingServiceInfo = serviceInfoCaptor.getAllValues().get(1);
1001         assertEquals(existingServiceInfo.getServiceInstanceName(), "service-instance-1");
1002         assertEquals(existingServiceInfo.getIpv4Address(), "192.168.1.1");
1003         assertEquals(existingServiceInfo.getPort(), 5353);
1004         assertEquals(existingServiceInfo.getSubtypes(), Collections.singletonList(SUBTYPE));
1005         assertNull(existingServiceInfo.getAttributeByKey("key"));
1006 
1007         // Process a goodbye message for the existing response.
1008         processResponse(createResponse(
1009                 "service-instance-1", "192.168.1.1", 5353,
1010                 SERVICE_TYPE_LABELS,
1011                 Collections.emptyMap(), /* ptrTtlMillis= */ 0L), socketKey);
1012 
1013         startSendAndReceive(mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
1014 
1015         // Verify onServiceFound was not called on the newly registered listener after the existing
1016         // response is gone.
1017         verify(mockListenerTwo, never()).onServiceNameDiscovered(
1018                 any(MdnsServiceInfo.class), eq(false));
1019         verify(mockListenerTwo, never()).onServiceFound(any(MdnsServiceInfo.class), anyBoolean());
1020     }
1021 
1022     @Test
processResponse_searchOptionsEnableServiceRemoval_shouldRemove()1023     public void processResponse_searchOptionsEnableServiceRemoval_shouldRemove()
1024             throws Exception {
1025         final String serviceInstanceName = "service-instance-1";
1026         MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
1027                 .setRemoveExpiredService(true)
1028                 .setNumOfQueriesBeforeBackoff(Integer.MAX_VALUE)
1029                 .build();
1030         startSendAndReceive(mockListenerOne, searchOptions);
1031         Runnable firstMdnsTask = currentThreadExecutor.getAndClearSubmittedRunnable();
1032 
1033         // Process the initial response.
1034         processResponse(createResponse(
1035                 serviceInstanceName, "192.168.1.1", 5353, SUBTYPE,
1036                 Collections.emptyMap(), TEST_TTL), socketKey);
1037 
1038         // Clear the scheduled runnable.
1039         currentThreadExecutor.getAndClearLastScheduledRunnable();
1040 
1041         // Simulate the case where the response is under TTL.
1042         doReturn(TEST_ELAPSED_REALTIME + TEST_TTL - 1L).when(mockDecoderClock).elapsedRealtime();
1043         firstMdnsTask.run();
1044         verify(mockDeps, times(1)).sendMessage(any(), any(Message.class));
1045 
1046         // Verify removed callback was not called.
1047         verifyServiceRemovedNoCallback(mockListenerOne);
1048 
1049         // Simulate the case where the response is after TTL.
1050         doReturn(TEST_ELAPSED_REALTIME + TEST_TTL + 1L).when(mockDecoderClock).elapsedRealtime();
1051         firstMdnsTask.run();
1052         verify(mockDeps, times(2)).sendMessage(any(), any(Message.class));
1053 
1054         // Verify removed callback was called.
1055         verifyServiceRemovedCallback(
1056                 mockListenerOne, serviceInstanceName, SERVICE_TYPE_LABELS, socketKey);
1057     }
1058 
1059     @Test
processResponse_searchOptionsNotEnableServiceRemoval_shouldNotRemove()1060     public void processResponse_searchOptionsNotEnableServiceRemoval_shouldNotRemove()
1061             throws Exception {
1062         final String serviceInstanceName = "service-instance-1";
1063         startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
1064         Runnable firstMdnsTask = currentThreadExecutor.getAndClearSubmittedRunnable();
1065 
1066         // Process the initial response.
1067         processResponse(createResponse(
1068                 serviceInstanceName, "192.168.1.1", 5353, SUBTYPE,
1069                 Collections.emptyMap(), TEST_TTL), socketKey);
1070 
1071         // Clear the scheduled runnable.
1072         currentThreadExecutor.getAndClearLastScheduledRunnable();
1073 
1074         // Simulate the case where the response is after TTL.
1075         doReturn(TEST_ELAPSED_REALTIME + TEST_TTL + 1L).when(mockDecoderClock).elapsedRealtime();
1076         firstMdnsTask.run();
1077 
1078         // Verify removed callback was not called.
1079         verifyServiceRemovedNoCallback(mockListenerOne);
1080     }
1081 
1082     @Test
1083     @Ignore("MdnsConfigs is not configurable currently.")
processResponse_removeServiceAfterTtlExpiresEnabled_shouldRemove()1084     public void processResponse_removeServiceAfterTtlExpiresEnabled_shouldRemove()
1085             throws Exception {
1086         //MdnsConfigsFlagsImpl.removeServiceAfterTtlExpires.override(true);
1087         final String serviceInstanceName = "service-instance-1";
1088         startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
1089         Runnable firstMdnsTask = currentThreadExecutor.getAndClearSubmittedRunnable();
1090 
1091         // Process the initial response.
1092         processResponse(createResponse(
1093                 serviceInstanceName, "192.168.1.1", 5353, SUBTYPE,
1094                 Collections.emptyMap(), TEST_TTL), socketKey);
1095 
1096         // Clear the scheduled runnable.
1097         currentThreadExecutor.getAndClearLastScheduledRunnable();
1098 
1099         // Simulate the case where the response is after TTL.
1100         doReturn(TEST_ELAPSED_REALTIME + TEST_TTL + 1L).when(mockDecoderClock).elapsedRealtime();
1101         firstMdnsTask.run();
1102 
1103         // Verify removed callback was called.
1104         verifyServiceRemovedCallback(
1105                 mockListenerOne, serviceInstanceName, SERVICE_TYPE_LABELS, socketKey);
1106     }
1107 
1108     @Test
testProcessResponse_InOrder()1109     public void testProcessResponse_InOrder() throws Exception {
1110         final String serviceName = "service-instance";
1111         final String ipV4Address = "192.0.2.0";
1112         final String ipV6Address = "2001:db8::";
1113         startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
1114         InOrder inOrder = inOrder(mockListenerOne);
1115 
1116         // Process the initial response which is incomplete.
1117         processResponse(createResponse(
1118                 serviceName, null, 5353, SUBTYPE,
1119                 Collections.emptyMap(), TEST_TTL), socketKey);
1120 
1121         // Process a second response which has ip address to make response become complete.
1122         processResponse(createResponse(
1123                 serviceName, ipV4Address, 5353, SUBTYPE,
1124                 Collections.emptyMap(), TEST_TTL), socketKey);
1125 
1126         // Process a third response with a different ip address, port and updated text attributes.
1127         processResponse(createResponse(
1128                 serviceName, ipV6Address, 5354, SUBTYPE,
1129                 Collections.singletonMap("key", "value"), TEST_TTL), socketKey);
1130 
1131         // Process the last response which is goodbye message (with the main type, not subtype).
1132         processResponse(createResponse(
1133                         serviceName, ipV6Address, 5354, SERVICE_TYPE_LABELS,
1134                         Collections.singletonMap("key", "value"), /* ptrTtlMillis= */ 0L),
1135                 socketKey);
1136 
1137         // Verify onServiceNameDiscovered was first called for the initial response.
1138         inOrder.verify(mockListenerOne).onServiceNameDiscovered(
1139                 serviceInfoCaptor.capture(), eq(false) /* isServiceFromCache */);
1140         verifyServiceInfo(serviceInfoCaptor.getAllValues().get(0),
1141                 serviceName,
1142                 SERVICE_TYPE_LABELS,
1143                 List.of() /* ipv4Address */,
1144                 List.of() /* ipv6Address */,
1145                 5353 /* port */,
1146                 Collections.singletonList(SUBTYPE) /* subTypes */,
1147                 Collections.singletonMap("key", null) /* attributes */,
1148                 socketKey);
1149 
1150         // Verify onServiceFound was second called for the second response.
1151         inOrder.verify(mockListenerOne).onServiceFound(
1152                 serviceInfoCaptor.capture(), eq(false) /* isServiceFromCache */);
1153         verifyServiceInfo(serviceInfoCaptor.getAllValues().get(1),
1154                 serviceName,
1155                 SERVICE_TYPE_LABELS,
1156                 List.of(ipV4Address) /* ipv4Address */,
1157                 List.of() /* ipv6Address */,
1158                 5353 /* port */,
1159                 Collections.singletonList(SUBTYPE) /* subTypes */,
1160                 Collections.singletonMap("key", null) /* attributes */,
1161                 socketKey);
1162 
1163         // Verify onServiceUpdated was third called for the third response.
1164         inOrder.verify(mockListenerOne).onServiceUpdated(serviceInfoCaptor.capture());
1165         verifyServiceInfo(serviceInfoCaptor.getAllValues().get(2),
1166                 serviceName,
1167                 SERVICE_TYPE_LABELS,
1168                 List.of(ipV4Address) /* ipv4Address */,
1169                 List.of(ipV6Address) /* ipv6Address */,
1170                 5354 /* port */,
1171                 Collections.singletonList(SUBTYPE) /* subTypes */,
1172                 Collections.singletonMap("key", "value") /* attributes */,
1173                 socketKey);
1174 
1175         // Verify onServiceRemoved was called for the last response.
1176         inOrder.verify(mockListenerOne).onServiceRemoved(serviceInfoCaptor.capture());
1177         verifyServiceInfo(serviceInfoCaptor.getAllValues().get(3),
1178                 serviceName,
1179                 SERVICE_TYPE_LABELS,
1180                 List.of(ipV4Address) /* ipv4Address */,
1181                 List.of(ipV6Address) /* ipv6Address */,
1182                 5354 /* port */,
1183                 Collections.singletonList(SUBTYPE) /* subTypes */,
1184                 Collections.singletonMap("key", "value") /* attributes */,
1185                 socketKey);
1186 
1187         // Verify onServiceNameRemoved was called for the last response.
1188         inOrder.verify(mockListenerOne).onServiceNameRemoved(serviceInfoCaptor.capture());
1189         verifyServiceInfo(serviceInfoCaptor.getAllValues().get(4),
1190                 serviceName,
1191                 SERVICE_TYPE_LABELS,
1192                 List.of(ipV4Address) /* ipv4Address */,
1193                 List.of(ipV6Address) /* ipv6Address */,
1194                 5354 /* port */,
1195                 Collections.singletonList(SUBTYPE) /* subTypes */,
1196                 Collections.singletonMap("key", "value") /* attributes */,
1197                 socketKey);
1198     }
1199 
1200     @Test
testProcessResponse_Resolve()1201     public void testProcessResponse_Resolve() throws Exception {
1202         final String instanceName = "service-instance";
1203         final String[] hostname = new String[] { "testhost "};
1204         final String ipV4Address = "192.0.2.0";
1205         final String ipV6Address = "2001:db8::";
1206 
1207         final MdnsSearchOptions resolveOptions1 = MdnsSearchOptions.newBuilder()
1208                 .setResolveInstanceName(instanceName).build();
1209         final MdnsSearchOptions resolveOptions2 = MdnsSearchOptions.newBuilder()
1210                 .setResolveInstanceName(instanceName).build();
1211 
1212         doCallRealMethod().when(mockDeps).getDatagramPacketsFromMdnsPacket(
1213                 any(), any(MdnsPacket.class), any(InetSocketAddress.class), anyBoolean());
1214 
1215         startSendAndReceive(mockListenerOne, resolveOptions1);
1216         startSendAndReceive(mockListenerTwo, resolveOptions2);
1217         // No need to verify order for both listeners; and order is not guaranteed between them
1218         InOrder inOrder = inOrder(mockListenerOne, mockSocketClient);
1219 
1220         // Verify a query for SRV/TXT was sent, but no PTR query
1221         final ArgumentCaptor<List<DatagramPacket>> srvTxtQueryCaptor =
1222                 ArgumentCaptor.forClass(List.class);
1223         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
1224         // Send twice for IPv4 and IPv6
1225         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingUnicastResponse(
1226                 srvTxtQueryCaptor.capture(),
1227                 eq(socketKey), eq(false));
1228         verify(mockDeps, times(1)).sendMessage(any(), any(Message.class));
1229         assertNotNull(delayMessage);
1230         inOrder.verify(mockListenerOne).onDiscoveryQuerySent(any(), anyInt());
1231         verify(mockListenerTwo).onDiscoveryQuerySent(any(), anyInt());
1232 
1233         final MdnsPacket srvTxtQueryPacket = MdnsPacket.parse(
1234                 new MdnsPacketReader(srvTxtQueryCaptor.getValue().get(0)));
1235 
1236         final String[] serviceName = getTestServiceName(instanceName);
1237         assertEquals(1, srvTxtQueryPacket.questions.size());
1238         assertFalse(hasQuestion(srvTxtQueryPacket, MdnsRecord.TYPE_PTR));
1239         assertTrue(hasQuestion(srvTxtQueryPacket, MdnsRecord.TYPE_ANY, serviceName));
1240         assertEquals(0, srvTxtQueryPacket.answers.size());
1241         assertEquals(0, srvTxtQueryPacket.authorityRecords.size());
1242         assertEquals(0, srvTxtQueryPacket.additionalRecords.size());
1243 
1244         // Process a response with SRV+TXT
1245         final MdnsPacket srvTxtResponse = new MdnsPacket(
1246                 0 /* flags */,
1247                 Collections.emptyList() /* questions */,
1248                 List.of(
1249                         new MdnsServiceRecord(serviceName, 0L /* receiptTimeMillis */,
1250                                 true /* cacheFlush */, TEST_TTL, 0 /* servicePriority */,
1251                                 0 /* serviceWeight */, 1234 /* servicePort */, hostname),
1252                         new MdnsTextRecord(serviceName, 0L /* receiptTimeMillis */,
1253                                 true /* cacheFlush */, TEST_TTL,
1254                                 Collections.emptyList() /* entries */)),
1255                 Collections.emptyList() /* authorityRecords */,
1256                 Collections.emptyList() /* additionalRecords */);
1257 
1258         processResponse(srvTxtResponse, socketKey);
1259         inOrder.verify(mockListenerOne).onServiceNameDiscovered(
1260                 matchServiceName(instanceName), eq(false) /* isServiceFromCache */);
1261         verify(mockListenerTwo).onServiceNameDiscovered(
1262                 matchServiceName(instanceName), eq(false) /* isServiceFromCache */);
1263 
1264         // Expect a query for A/AAAA
1265         dispatchMessage();
1266         final ArgumentCaptor<List<DatagramPacket>> addressQueryCaptor =
1267                 ArgumentCaptor.forClass(List.class);
1268         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
1269         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingMulticastResponse(
1270                 addressQueryCaptor.capture(),
1271                 eq(socketKey), eq(false));
1272         inOrder.verify(mockListenerOne).onDiscoveryQuerySent(any(), anyInt());
1273         // onDiscoveryQuerySent was called 2 times in total
1274         verify(mockListenerTwo, times(2)).onDiscoveryQuerySent(any(), anyInt());
1275 
1276         final MdnsPacket addressQueryPacket = MdnsPacket.parse(
1277                 new MdnsPacketReader(addressQueryCaptor.getValue().get(0)));
1278         assertEquals(2, addressQueryPacket.questions.size());
1279         assertTrue(hasQuestion(addressQueryPacket, MdnsRecord.TYPE_A, hostname));
1280         assertTrue(hasQuestion(addressQueryPacket, MdnsRecord.TYPE_AAAA, hostname));
1281         assertEquals(0, addressQueryPacket.answers.size());
1282         assertEquals(0, addressQueryPacket.authorityRecords.size());
1283         assertEquals(0, addressQueryPacket.additionalRecords.size());
1284 
1285         // Process a response with address records
1286         final MdnsPacket addressResponse = new MdnsPacket(
1287                 0 /* flags */,
1288                 Collections.emptyList() /* questions */,
1289                 List.of(
1290                         new MdnsInetAddressRecord(hostname, 0L /* receiptTimeMillis */,
1291                                 true /* cacheFlush */, TEST_TTL,
1292                                 InetAddresses.parseNumericAddress(ipV4Address)),
1293                         new MdnsInetAddressRecord(hostname, 0L /* receiptTimeMillis */,
1294                                 true /* cacheFlush */, TEST_TTL,
1295                                 InetAddresses.parseNumericAddress(ipV6Address))),
1296                 Collections.emptyList() /* authorityRecords */,
1297                 Collections.emptyList() /* additionalRecords */);
1298 
1299         inOrder.verify(mockListenerOne, never()).onServiceNameDiscovered(any(), anyBoolean());
1300         verifyNoMoreInteractions(mockListenerTwo);
1301         processResponse(addressResponse, socketKey);
1302 
1303         inOrder.verify(mockListenerOne).onServiceFound(
1304                 serviceInfoCaptor.capture(), eq(false) /* isServiceFromCache */);
1305         verify(mockListenerTwo).onServiceFound(any(), anyBoolean());
1306         verifyServiceInfo(serviceInfoCaptor.getValue(),
1307                 instanceName,
1308                 SERVICE_TYPE_LABELS,
1309                 List.of(ipV4Address),
1310                 List.of(ipV6Address),
1311                 1234 /* port */,
1312                 Collections.emptyList() /* subTypes */,
1313                 Collections.emptyMap() /* attributes */,
1314                 socketKey);
1315     }
1316 
1317     @Test
testRenewTxtSrvInResolve()1318     public void testRenewTxtSrvInResolve() throws Exception {
1319         final String instanceName = "service-instance";
1320         final String[] hostname = new String[] { "testhost "};
1321         final String ipV4Address = "192.0.2.0";
1322         final String ipV6Address = "2001:db8::";
1323 
1324         final MdnsSearchOptions resolveOptions = MdnsSearchOptions.newBuilder()
1325                 .setResolveInstanceName(instanceName).build();
1326 
1327         doCallRealMethod().when(mockDeps).getDatagramPacketsFromMdnsPacket(
1328                 any(), any(MdnsPacket.class), any(InetSocketAddress.class), anyBoolean());
1329 
1330         startSendAndReceive(mockListenerOne, resolveOptions);
1331         InOrder inOrder = inOrder(mockListenerOne, mockSocketClient);
1332 
1333         // Get the query for SRV/TXT
1334         final ArgumentCaptor<List<DatagramPacket>> srvTxtQueryCaptor =
1335                 ArgumentCaptor.forClass(List.class);
1336         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
1337         // Send twice for IPv4 and IPv6
1338         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingUnicastResponse(
1339                 srvTxtQueryCaptor.capture(),
1340                 eq(socketKey), eq(false));
1341         verify(mockDeps, times(1)).sendMessage(any(), any(Message.class));
1342         assertNotNull(delayMessage);
1343 
1344         final MdnsPacket srvTxtQueryPacket = MdnsPacket.parse(
1345                 new MdnsPacketReader(srvTxtQueryCaptor.getValue().get(0)));
1346 
1347         final String[] serviceName = getTestServiceName(instanceName);
1348         assertTrue(hasQuestion(srvTxtQueryPacket, MdnsRecord.TYPE_ANY, serviceName));
1349 
1350         // Process a response with all records
1351         final MdnsPacket srvTxtResponse = new MdnsPacket(
1352                 0 /* flags */,
1353                 Collections.emptyList() /* questions */,
1354                 List.of(
1355                         new MdnsServiceRecord(serviceName, TEST_ELAPSED_REALTIME,
1356                                 true /* cacheFlush */, TEST_TTL, 0 /* servicePriority */,
1357                                 0 /* serviceWeight */, 1234 /* servicePort */, hostname),
1358                         new MdnsTextRecord(serviceName, TEST_ELAPSED_REALTIME,
1359                                 true /* cacheFlush */, TEST_TTL,
1360                                 Collections.emptyList() /* entries */),
1361                         new MdnsInetAddressRecord(hostname, TEST_ELAPSED_REALTIME,
1362                                 true /* cacheFlush */, TEST_TTL,
1363                                 InetAddresses.parseNumericAddress(ipV4Address)),
1364                         new MdnsInetAddressRecord(hostname, TEST_ELAPSED_REALTIME,
1365                                 true /* cacheFlush */, TEST_TTL,
1366                                 InetAddresses.parseNumericAddress(ipV6Address))),
1367                 Collections.emptyList() /* authorityRecords */,
1368                 Collections.emptyList() /* additionalRecords */);
1369         processResponse(srvTxtResponse, socketKey);
1370         dispatchMessage();
1371         inOrder.verify(mockListenerOne).onServiceNameDiscovered(
1372                 any(), eq(false) /* isServiceFromCache */);
1373         inOrder.verify(mockListenerOne).onServiceFound(
1374                 any(), eq(false) /* isServiceFromCache */);
1375 
1376         // Expect no query on the next run
1377         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
1378         inOrder.verifyNoMoreInteractions();
1379 
1380         // Advance time so 75% of TTL passes and re-execute
1381         doReturn(TEST_ELAPSED_REALTIME + (long) (TEST_TTL * 0.75))
1382                 .when(mockDecoderClock).elapsedRealtime();
1383         verify(mockDeps, times(2)).sendMessage(any(), any(Message.class));
1384         assertNotNull(delayMessage);
1385         dispatchMessage();
1386         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
1387 
1388         // Expect a renewal query
1389         final ArgumentCaptor<List<DatagramPacket>> renewalQueryCaptor =
1390                 ArgumentCaptor.forClass(List.class);
1391         // Second and later sends are sent as "expect multicast response" queries
1392         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingMulticastResponse(
1393                 renewalQueryCaptor.capture(),
1394                 eq(socketKey), eq(false));
1395         verify(mockDeps, times(3)).sendMessage(any(), any(Message.class));
1396         assertNotNull(delayMessage);
1397         inOrder.verify(mockListenerOne).onDiscoveryQuerySent(any(), anyInt());
1398         final MdnsPacket renewalPacket = MdnsPacket.parse(
1399                 new MdnsPacketReader(renewalQueryCaptor.getValue().get(0)));
1400         assertTrue(hasQuestion(renewalPacket, MdnsRecord.TYPE_ANY, serviceName));
1401         inOrder.verifyNoMoreInteractions();
1402 
1403         long updatedReceiptTime =  TEST_ELAPSED_REALTIME + TEST_TTL;
1404         final MdnsPacket refreshedSrvTxtResponse = new MdnsPacket(
1405                 0 /* flags */,
1406                 Collections.emptyList() /* questions */,
1407                 List.of(
1408                         new MdnsServiceRecord(serviceName, updatedReceiptTime,
1409                                 true /* cacheFlush */, TEST_TTL, 0 /* servicePriority */,
1410                                 0 /* serviceWeight */, 1234 /* servicePort */, hostname),
1411                         new MdnsTextRecord(serviceName, updatedReceiptTime,
1412                                 true /* cacheFlush */, TEST_TTL,
1413                                 Collections.emptyList() /* entries */),
1414                         new MdnsInetAddressRecord(hostname, updatedReceiptTime,
1415                                 true /* cacheFlush */, TEST_TTL,
1416                                 InetAddresses.parseNumericAddress(ipV4Address)),
1417                         new MdnsInetAddressRecord(hostname, updatedReceiptTime,
1418                                 true /* cacheFlush */, TEST_TTL,
1419                                 InetAddresses.parseNumericAddress(ipV6Address))),
1420                 Collections.emptyList() /* authorityRecords */,
1421                 Collections.emptyList() /* additionalRecords */);
1422         processResponse(refreshedSrvTxtResponse, socketKey);
1423         dispatchMessage();
1424 
1425         // Advance time to updatedReceiptTime + 1, expected no refresh query because the cache
1426         // should contain the record that have update last receipt time.
1427         doReturn(updatedReceiptTime + 1).when(mockDecoderClock).elapsedRealtime();
1428         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
1429         inOrder.verifyNoMoreInteractions();
1430     }
1431 
1432     @Test
testProcessResponse_ResolveExcludesOtherServices()1433     public void testProcessResponse_ResolveExcludesOtherServices() {
1434         final String requestedInstance = "instance1";
1435         final String otherInstance = "instance2";
1436         final String ipV4Address = "192.0.2.0";
1437         final String ipV6Address = "2001:db8::";
1438         final String capitalizedRequestInstance = "Instance1";
1439 
1440         final MdnsSearchOptions resolveOptions = MdnsSearchOptions.newBuilder()
1441                 // Use different case in the options
1442                 .setResolveInstanceName(capitalizedRequestInstance).build();
1443 
1444         startSendAndReceive(mockListenerOne, resolveOptions);
1445         startSendAndReceive(mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
1446 
1447         // Complete response from instanceName
1448         processResponse(createResponse(
1449                         requestedInstance, ipV4Address, 5353, SERVICE_TYPE_LABELS,
1450                         Collections.emptyMap() /* textAttributes */, TEST_TTL),
1451                 socketKey);
1452 
1453         // Complete response from otherInstanceName
1454         processResponse(createResponse(
1455                         otherInstance, ipV4Address, 5353, SERVICE_TYPE_LABELS,
1456                         Collections.emptyMap() /* textAttributes */, TEST_TTL),
1457                 socketKey);
1458 
1459         // Address update from otherInstanceName
1460         processResponse(createResponse(
1461                 otherInstance, ipV6Address, 5353, SERVICE_TYPE_LABELS,
1462                 Collections.emptyMap(), TEST_TTL), socketKey);
1463 
1464         // Goodbye from otherInstanceName
1465         processResponse(createResponse(
1466                 otherInstance, ipV6Address, 5353, SERVICE_TYPE_LABELS,
1467                 Collections.emptyMap(), 0L /* ttl */), socketKey);
1468 
1469         // mockListenerOne gets notified for the requested instance
1470         verify(mockListenerOne).onServiceNameDiscovered(
1471                 matchServiceName(capitalizedRequestInstance), eq(false) /* isServiceFromCache */);
1472         verify(mockListenerOne).onServiceFound(
1473                 matchServiceName(capitalizedRequestInstance), eq(false) /* isServiceFromCache */);
1474 
1475         // ...but does not get any callback for the other instance
1476         verify(mockListenerOne, never()).onServiceFound(
1477                 matchServiceName(otherInstance), anyBoolean());
1478         verify(mockListenerOne, never()).onServiceNameDiscovered(
1479                 matchServiceName(otherInstance), anyBoolean());
1480         verify(mockListenerOne, never()).onServiceUpdated(matchServiceName(otherInstance));
1481         verify(mockListenerOne, never()).onServiceRemoved(matchServiceName(otherInstance));
1482 
1483         // mockListenerTwo gets notified for both though
1484         final InOrder inOrder = inOrder(mockListenerTwo);
1485         inOrder.verify(mockListenerTwo).onServiceNameDiscovered(
1486                 matchServiceName(capitalizedRequestInstance), eq(false) /* isServiceFromCache */);
1487         inOrder.verify(mockListenerTwo).onServiceFound(
1488                 matchServiceName(capitalizedRequestInstance), eq(false) /* isServiceFromCache */);
1489 
1490         inOrder.verify(mockListenerTwo).onServiceNameDiscovered(
1491                 matchServiceName(otherInstance), eq(false) /* isServiceFromCache */);
1492         inOrder.verify(mockListenerTwo).onServiceFound(
1493                 matchServiceName(otherInstance), eq(false) /* isServiceFromCache */);
1494         inOrder.verify(mockListenerTwo).onServiceUpdated(matchServiceName(otherInstance));
1495         inOrder.verify(mockListenerTwo).onServiceRemoved(matchServiceName(otherInstance));
1496     }
1497 
1498     @Test
testProcessResponse_SubtypeDiscoveryLimitedToSubtype()1499     public void testProcessResponse_SubtypeDiscoveryLimitedToSubtype() {
1500         final String matchingInstance = "instance1";
1501         final String subtype = "_subtype";
1502         final String otherInstance = "instance2";
1503         final String ipV4Address = "192.0.2.0";
1504         final String ipV6Address = "2001:db8::";
1505 
1506         final MdnsSearchOptions options = MdnsSearchOptions.newBuilder()
1507                 // Search with different case. Note MdnsSearchOptions subtype doesn't start with "_"
1508                 .addSubtype("Subtype").build();
1509 
1510         startSendAndReceive(mockListenerOne, options);
1511         startSendAndReceive(mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
1512 
1513         // Complete response from instanceName
1514         final MdnsPacket packetWithoutSubtype = createResponse(
1515                 matchingInstance, ipV4Address, 5353, SERVICE_TYPE_LABELS,
1516                 Collections.emptyMap() /* textAttributes */, TEST_TTL);
1517         final MdnsPointerRecord originalPtr = (MdnsPointerRecord) CollectionUtils.findFirst(
1518                 packetWithoutSubtype.answers, r -> r instanceof MdnsPointerRecord);
1519 
1520         // Add a subtype PTR record
1521         final ArrayList<MdnsRecord> newAnswers = new ArrayList<>(packetWithoutSubtype.answers);
1522         newAnswers.add(new MdnsPointerRecord(
1523                 // PTR should be _subtype._sub._type._tcp.local -> instance1._type._tcp.local
1524                 Stream.concat(Stream.of(subtype, "_sub"), Arrays.stream(SERVICE_TYPE_LABELS))
1525                         .toArray(String[]::new),
1526                 originalPtr.getReceiptTime(), originalPtr.getCacheFlush(), originalPtr.getTtl(),
1527                 originalPtr.getPointer()));
1528         final MdnsPacket packetWithSubtype = new MdnsPacket(
1529                 packetWithoutSubtype.flags,
1530                 packetWithoutSubtype.questions,
1531                 newAnswers,
1532                 packetWithoutSubtype.authorityRecords,
1533                 packetWithoutSubtype.additionalRecords);
1534         processResponse(packetWithSubtype, socketKey);
1535 
1536         // Complete response from otherInstanceName, without subtype
1537         processResponse(createResponse(
1538                         otherInstance, ipV4Address, 5353, SERVICE_TYPE_LABELS,
1539                         Collections.emptyMap() /* textAttributes */, TEST_TTL),
1540                 socketKey);
1541 
1542         // Address update from otherInstanceName
1543         processResponse(createResponse(
1544                 otherInstance, ipV6Address, 5353, SERVICE_TYPE_LABELS,
1545                 Collections.emptyMap(), TEST_TTL), socketKey);
1546 
1547         // Goodbye from otherInstanceName
1548         processResponse(createResponse(
1549                 otherInstance, ipV6Address, 5353, SERVICE_TYPE_LABELS,
1550                 Collections.emptyMap(), 0L /* ttl */), socketKey);
1551 
1552         // mockListenerOne gets notified for the requested instance
1553         final ArgumentMatcher<MdnsServiceInfo> subtypeInstanceMatcher = info ->
1554                 info.getServiceInstanceName().equals(matchingInstance)
1555                         && info.getSubtypes().equals(Collections.singletonList(subtype));
1556         verify(mockListenerOne).onServiceNameDiscovered(
1557                 argThat(subtypeInstanceMatcher), eq(false) /* isServiceFromCache */);
1558         verify(mockListenerOne).onServiceFound(
1559                 argThat(subtypeInstanceMatcher), eq(false) /* isServiceFromCache */);
1560 
1561         // ...but does not get any callback for the other instance
1562         verify(mockListenerOne, never()).onServiceFound(
1563                 matchServiceName(otherInstance), anyBoolean());
1564         verify(mockListenerOne, never()).onServiceNameDiscovered(
1565                 matchServiceName(otherInstance), anyBoolean());
1566         verify(mockListenerOne, never()).onServiceUpdated(matchServiceName(otherInstance));
1567         verify(mockListenerOne, never()).onServiceRemoved(matchServiceName(otherInstance));
1568 
1569         // mockListenerTwo gets notified for both though
1570         final InOrder inOrder = inOrder(mockListenerTwo);
1571         inOrder.verify(mockListenerTwo).onServiceNameDiscovered(
1572                 argThat(subtypeInstanceMatcher), eq(false) /* isServiceFromCache */);
1573         inOrder.verify(mockListenerTwo).onServiceFound(
1574                 argThat(subtypeInstanceMatcher), eq(false) /* isServiceFromCache */);
1575 
1576         inOrder.verify(mockListenerTwo).onServiceNameDiscovered(
1577                 matchServiceName(otherInstance), eq(false) /* isServiceFromCache */);
1578         inOrder.verify(mockListenerTwo).onServiceFound(
1579                 matchServiceName(otherInstance), eq(false) /* isServiceFromCache */);
1580         inOrder.verify(mockListenerTwo).onServiceUpdated(matchServiceName(otherInstance));
1581         inOrder.verify(mockListenerTwo).onServiceRemoved(matchServiceName(otherInstance));
1582     }
1583 
1584     @Test
testProcessResponse_SubtypeChange()1585     public void testProcessResponse_SubtypeChange() {
1586         final String matchingInstance = "instance1";
1587         final String subtype = "_subtype";
1588         final String ipV4Address = "192.0.2.0";
1589         final String ipV6Address = "2001:db8::";
1590 
1591         final MdnsSearchOptions options = MdnsSearchOptions.newBuilder()
1592                 .addSubtype("othersub").build();
1593 
1594         startSendAndReceive(mockListenerOne, options);
1595 
1596         // Complete response from instanceName
1597         final MdnsPacket packetWithoutSubtype = createResponse(
1598                 matchingInstance, ipV4Address, 5353, SERVICE_TYPE_LABELS,
1599                 Collections.emptyMap() /* textAttributes */, TEST_TTL);
1600         final MdnsPointerRecord originalPtr = (MdnsPointerRecord) CollectionUtils.findFirst(
1601                 packetWithoutSubtype.answers, r -> r instanceof MdnsPointerRecord);
1602 
1603         // Add a subtype PTR record
1604         final ArrayList<MdnsRecord> newAnswers = new ArrayList<>(packetWithoutSubtype.answers);
1605         newAnswers.add(new MdnsPointerRecord(
1606                 // PTR should be _subtype._sub._type._tcp.local -> instance1._type._tcp.local
1607                 Stream.concat(Stream.of(subtype, "_sub"), Arrays.stream(SERVICE_TYPE_LABELS))
1608                         .toArray(String[]::new),
1609                 originalPtr.getReceiptTime(), originalPtr.getCacheFlush(), originalPtr.getTtl(),
1610                 originalPtr.getPointer()));
1611         processResponse(new MdnsPacket(
1612                 packetWithoutSubtype.flags,
1613                 packetWithoutSubtype.questions,
1614                 newAnswers,
1615                 packetWithoutSubtype.authorityRecords,
1616                 packetWithoutSubtype.additionalRecords), socketKey);
1617 
1618         // The subtype does not match
1619         final InOrder inOrder = inOrder(mockListenerOne);
1620         inOrder.verify(mockListenerOne, never()).onServiceNameDiscovered(any(), anyBoolean());
1621 
1622         // Add another matching subtype
1623         newAnswers.add(new MdnsPointerRecord(
1624                 // PTR should be _subtype._sub._type._tcp.local -> instance1._type._tcp.local
1625                 Stream.concat(Stream.of("_othersub", "_sub"), Arrays.stream(SERVICE_TYPE_LABELS))
1626                         .toArray(String[]::new),
1627                 originalPtr.getReceiptTime(), originalPtr.getCacheFlush(), originalPtr.getTtl(),
1628                 originalPtr.getPointer()));
1629         processResponse(new MdnsPacket(
1630                 packetWithoutSubtype.flags,
1631                 packetWithoutSubtype.questions,
1632                 newAnswers,
1633                 packetWithoutSubtype.authorityRecords,
1634                 packetWithoutSubtype.additionalRecords), socketKey);
1635 
1636         final ArgumentMatcher<MdnsServiceInfo> subtypeInstanceMatcher = info ->
1637                 info.getServiceInstanceName().equals(matchingInstance)
1638                         && info.getSubtypes().equals(List.of("_subtype", "_othersub"));
1639 
1640         // Service found callbacks are sent now
1641         inOrder.verify(mockListenerOne).onServiceNameDiscovered(
1642                 argThat(subtypeInstanceMatcher), eq(false) /* isServiceFromCache */);
1643         inOrder.verify(mockListenerOne).onServiceFound(
1644                 argThat(subtypeInstanceMatcher), eq(false) /* isServiceFromCache */);
1645 
1646         // Address update: update callbacks are sent
1647         processResponse(createResponse(
1648                 matchingInstance, ipV6Address, 5353, SERVICE_TYPE_LABELS,
1649                 Collections.emptyMap(), TEST_TTL), socketKey);
1650 
1651         inOrder.verify(mockListenerOne).onServiceUpdated(argThat(info ->
1652                 subtypeInstanceMatcher.matches(info)
1653                         && info.getIpv4Addresses().equals(List.of(ipV4Address))
1654                         && info.getIpv6Addresses().equals(List.of(ipV6Address))));
1655 
1656         // Goodbye: service removed callbacks are sent
1657         processResponse(createResponse(
1658                 matchingInstance, ipV6Address, 5353, SERVICE_TYPE_LABELS,
1659                 Collections.emptyMap(), 0L /* ttl */), socketKey);
1660 
1661         inOrder.verify(mockListenerOne).onServiceRemoved(matchServiceName(matchingInstance));
1662         inOrder.verify(mockListenerOne).onServiceNameRemoved(matchServiceName(matchingInstance));
1663     }
1664 
1665     @Test
testNotifySocketDestroyed()1666     public void testNotifySocketDestroyed() throws Exception {
1667         final String requestedInstance = "instance1";
1668         final String otherInstance = "instance2";
1669         final String ipV4Address = "192.0.2.0";
1670 
1671         final MdnsSearchOptions resolveOptions = MdnsSearchOptions.newBuilder()
1672                 .setNumOfQueriesBeforeBackoff(Integer.MAX_VALUE)
1673                 .setResolveInstanceName("instance1").build();
1674 
1675         startSendAndReceive(mockListenerOne, resolveOptions);
1676         // Always try to remove the task.
1677         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
1678         // Ensure the first task is executed so it schedules a future task
1679         currentThreadExecutor.getAndClearSubmittedFuture().get(
1680                 TEST_TIMEOUT_MS, TimeUnit.MILLISECONDS);
1681         startSendAndReceive(mockListenerTwo,
1682                 MdnsSearchOptions.newBuilder().setNumOfQueriesBeforeBackoff(
1683                         Integer.MAX_VALUE).build());
1684 
1685         // Filing the second request cancels the first future
1686         verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
1687 
1688         // Ensure it gets executed too
1689         currentThreadExecutor.getAndClearSubmittedFuture().get(
1690                 TEST_TIMEOUT_MS, TimeUnit.MILLISECONDS);
1691 
1692         // Complete response from instanceName
1693         processResponse(createResponse(
1694                         requestedInstance, ipV4Address, 5353, SERVICE_TYPE_LABELS,
1695                         Collections.emptyMap() /* textAttributes */, TEST_TTL),
1696                 socketKey);
1697 
1698         // Complete response from otherInstanceName
1699         processResponse(createResponse(
1700                         otherInstance, ipV4Address, 5353, SERVICE_TYPE_LABELS,
1701                         Collections.emptyMap() /* textAttributes */, TEST_TTL),
1702                 socketKey);
1703 
1704         notifySocketDestroyed();
1705         verify(mockDeps, times(3)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
1706 
1707         // mockListenerOne gets notified for the requested instance
1708         final InOrder inOrder1 = inOrder(mockListenerOne);
1709         inOrder1.verify(mockListenerOne).onServiceNameDiscovered(
1710                 matchServiceName(requestedInstance), eq(false) /* isServiceFromCache */);
1711         inOrder1.verify(mockListenerOne).onServiceFound(
1712                 matchServiceName(requestedInstance), eq(false) /* isServiceFromCache */);
1713         inOrder1.verify(mockListenerOne).onServiceRemoved(matchServiceName(requestedInstance));
1714         inOrder1.verify(mockListenerOne).onServiceNameRemoved(matchServiceName(requestedInstance));
1715         verify(mockListenerOne, never()).onServiceFound(
1716                 matchServiceName(otherInstance), anyBoolean());
1717         verify(mockListenerOne, never()).onServiceNameDiscovered(
1718                 matchServiceName(otherInstance), anyBoolean());
1719         verify(mockListenerOne, never()).onServiceRemoved(matchServiceName(otherInstance));
1720         verify(mockListenerOne, never()).onServiceNameRemoved(matchServiceName(otherInstance));
1721 
1722         // mockListenerTwo gets notified for both though
1723         final InOrder inOrder2 = inOrder(mockListenerTwo);
1724         inOrder2.verify(mockListenerTwo).onServiceNameDiscovered(
1725                 matchServiceName(requestedInstance), eq(false) /* isServiceFromCache */);
1726         inOrder2.verify(mockListenerTwo).onServiceFound(
1727                 matchServiceName(requestedInstance), eq(false) /* isServiceFromCache */);
1728         inOrder2.verify(mockListenerTwo).onServiceRemoved(matchServiceName(requestedInstance));
1729         inOrder2.verify(mockListenerTwo).onServiceNameRemoved(matchServiceName(requestedInstance));
1730         verify(mockListenerTwo).onServiceNameDiscovered(
1731                 matchServiceName(otherInstance), eq(false) /* isServiceFromCache */);
1732         verify(mockListenerTwo).onServiceFound(
1733                 matchServiceName(otherInstance), eq(false) /* isServiceFromCache */);
1734         verify(mockListenerTwo).onServiceRemoved(matchServiceName(otherInstance));
1735         verify(mockListenerTwo).onServiceNameRemoved(matchServiceName(otherInstance));
1736     }
1737 
1738     @Test
testServicesAreCached()1739     public void testServicesAreCached() throws Exception {
1740         final String serviceName = "service-instance";
1741         final String ipV4Address = "192.0.2.0";
1742         // Register a listener
1743         startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
1744         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
1745         InOrder inOrder = inOrder(mockListenerOne);
1746 
1747         // Process a response which has ip address to make response become complete.
1748 
1749         processResponse(createResponse(
1750                         serviceName, ipV4Address, 5353, SUBTYPE,
1751                         Collections.emptyMap(), TEST_TTL),
1752                 socketKey);
1753 
1754         // Verify that onServiceNameDiscovered is called.
1755         inOrder.verify(mockListenerOne).onServiceNameDiscovered(
1756                 serviceInfoCaptor.capture(), eq(false) /* isServiceFromCache */);
1757         verifyServiceInfo(serviceInfoCaptor.getAllValues().get(0),
1758                 serviceName,
1759                 SERVICE_TYPE_LABELS,
1760                 List.of(ipV4Address) /* ipv4Address */,
1761                 List.of() /* ipv6Address */,
1762                 5353 /* port */,
1763                 Collections.singletonList(SUBTYPE) /* subTypes */,
1764                 Collections.singletonMap("key", null) /* attributes */,
1765                 socketKey);
1766 
1767         // Verify that onServiceFound is called.
1768         inOrder.verify(mockListenerOne).onServiceFound(
1769                 serviceInfoCaptor.capture(), eq(false) /* isServiceFromCache */);
1770         verifyServiceInfo(serviceInfoCaptor.getAllValues().get(1),
1771                 serviceName,
1772                 SERVICE_TYPE_LABELS,
1773                 List.of(ipV4Address) /* ipv4Address */,
1774                 List.of() /* ipv6Address */,
1775                 5353 /* port */,
1776                 Collections.singletonList(SUBTYPE) /* subTypes */,
1777                 Collections.singletonMap("key", null) /* attributes */,
1778                 socketKey);
1779 
1780         // Unregister the listener
1781         stopSendAndReceive(mockListenerOne);
1782         verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
1783 
1784         // Register another listener.
1785         startSendAndReceive(mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
1786         verify(mockDeps, times(3)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
1787         InOrder inOrder2 = inOrder(mockListenerTwo);
1788 
1789         // The services are cached in MdnsServiceCache, verify that onServiceNameDiscovered is
1790         // called immediately.
1791         inOrder2.verify(mockListenerTwo).onServiceNameDiscovered(
1792                 serviceInfoCaptor.capture(), eq(true) /* isServiceFromCache */);
1793         verifyServiceInfo(serviceInfoCaptor.getAllValues().get(2),
1794                 serviceName,
1795                 SERVICE_TYPE_LABELS,
1796                 List.of(ipV4Address) /* ipv4Address */,
1797                 List.of() /* ipv6Address */,
1798                 5353 /* port */,
1799                 Collections.singletonList(SUBTYPE) /* subTypes */,
1800                 Collections.singletonMap("key", null) /* attributes */,
1801                 socketKey);
1802 
1803         // The services are cached in MdnsServiceCache, verify that onServiceFound is
1804         // called immediately.
1805         inOrder2.verify(mockListenerTwo).onServiceFound(
1806                 serviceInfoCaptor.capture(), eq(true) /* isServiceFromCache */);
1807         verifyServiceInfo(serviceInfoCaptor.getAllValues().get(3),
1808                 serviceName,
1809                 SERVICE_TYPE_LABELS,
1810                 List.of(ipV4Address) /* ipv4Address */,
1811                 List.of() /* ipv6Address */,
1812                 5353 /* port */,
1813                 Collections.singletonList(SUBTYPE) /* subTypes */,
1814                 Collections.singletonMap("key", null) /* attributes */,
1815                 socketKey);
1816 
1817         // Process a response with a different ip address, port and updated text attributes.
1818         final String ipV6Address = "2001:db8::";
1819         processResponse(createResponse(
1820                 serviceName, ipV6Address, 5354, SUBTYPE,
1821                 Collections.singletonMap("key", "value"), TEST_TTL), socketKey);
1822 
1823         // Verify the onServiceUpdated is called.
1824         inOrder2.verify(mockListenerTwo).onServiceUpdated(serviceInfoCaptor.capture());
1825         verifyServiceInfo(serviceInfoCaptor.getAllValues().get(4),
1826                 serviceName,
1827                 SERVICE_TYPE_LABELS,
1828                 List.of(ipV4Address) /* ipv4Address */,
1829                 List.of(ipV6Address) /* ipv6Address */,
1830                 5354 /* port */,
1831                 Collections.singletonList(SUBTYPE) /* subTypes */,
1832                 Collections.singletonMap("key", "value") /* attributes */,
1833                 socketKey);
1834     }
1835 
1836     @Test
sendQueries_aggressiveScanMode()1837     public void sendQueries_aggressiveScanMode() {
1838         final MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
1839                 .addSubtype(SUBTYPE).setQueryMode(AGGRESSIVE_QUERY_MODE).build();
1840         startSendAndReceive(mockListenerOne, searchOptions);
1841         // Always try to remove the task.
1842         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
1843 
1844         int burstCounter = 0;
1845         int betweenBurstTime = 0;
1846         for (int i = 0; i < expectedIPv4Packets.length; i += 3) {
1847             verifyAndSendQuery(i, betweenBurstTime, /* expectsUnicastResponse= */ true);
1848             verifyAndSendQuery(i + 1, /* timeInMs= */ 0, /* expectsUnicastResponse= */ false);
1849             verifyAndSendQuery(i + 2, TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS,
1850                     /* expectsUnicastResponse= */ false);
1851             betweenBurstTime = Math.min(
1852                     INITIAL_AGGRESSIVE_TIME_BETWEEN_BURSTS_MS * (int) Math.pow(2, burstCounter),
1853                     MAX_TIME_BETWEEN_AGGRESSIVE_BURSTS_MS);
1854             burstCounter++;
1855         }
1856         // Verify that Task is not removed before stopSendAndReceive was called.
1857         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
1858 
1859         // Stop sending packets.
1860         stopSendAndReceive(mockListenerOne);
1861         verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
1862     }
1863 
1864     @Test
sendQueries_reentry_aggressiveScanMode()1865     public void sendQueries_reentry_aggressiveScanMode() {
1866         final MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
1867                 .addSubtype(SUBTYPE).setQueryMode(AGGRESSIVE_QUERY_MODE).build();
1868         startSendAndReceive(mockListenerOne, searchOptions);
1869         // Always try to remove the task.
1870         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
1871 
1872         // First burst, first query is sent.
1873         verifyAndSendQuery(0, /* timeInMs= */ 0, /* expectsUnicastResponse= */ true);
1874 
1875         // After the first query is sent, change the subtypes, and restart.
1876         final MdnsSearchOptions searchOptions2 = MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE)
1877                 .addSubtype("_subtype2").setQueryMode(AGGRESSIVE_QUERY_MODE).build();
1878         startSendAndReceive(mockListenerOne, searchOptions2);
1879         // The previous scheduled task should be canceled.
1880         verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
1881 
1882         // Queries should continue to be sent.
1883         verifyAndSendQuery(1, /* timeInMs= */ 0, /* expectsUnicastResponse= */ true);
1884         verifyAndSendQuery(2, /* timeInMs= */ 0, /* expectsUnicastResponse= */ false);
1885         verifyAndSendQuery(3, TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS,
1886                 /* expectsUnicastResponse= */ false);
1887 
1888         // Stop sending packets.
1889         stopSendAndReceive(mockListenerOne);
1890         verify(mockDeps, times(3)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
1891     }
1892 
1893     @Test
sendQueries_blendScanWithQueryBackoff()1894     public void sendQueries_blendScanWithQueryBackoff() {
1895         final int numOfQueriesBeforeBackoff = 11;
1896         final MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
1897                 .addSubtype(SUBTYPE)
1898                 .setQueryMode(AGGRESSIVE_QUERY_MODE)
1899                 .setNumOfQueriesBeforeBackoff(numOfQueriesBeforeBackoff)
1900                 .build();
1901         startSendAndReceive(mockListenerOne, searchOptions);
1902         // Always try to remove the task.
1903         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
1904 
1905         int burstCounter = 0;
1906         int betweenBurstTime = 0;
1907         for (int i = 0; i < numOfQueriesBeforeBackoff; i += 3) {
1908             verifyAndSendQuery(i, betweenBurstTime, /* expectsUnicastResponse= */ true);
1909             verifyAndSendQuery(i + 1, /* timeInMs= */ 0, /* expectsUnicastResponse= */ false);
1910             verifyAndSendQuery(i + 2, TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS,
1911                     /* expectsUnicastResponse= */ false);
1912             betweenBurstTime = Math.min(
1913                     INITIAL_AGGRESSIVE_TIME_BETWEEN_BURSTS_MS * (int) Math.pow(2, burstCounter),
1914                     MAX_TIME_BETWEEN_AGGRESSIVE_BURSTS_MS);
1915             burstCounter++;
1916         }
1917         // In backoff mode, the current scheduled task will be canceled and reschedule if the
1918         // 0.8 * smallestRemainingTtl is larger than time to next run.
1919         long currentTime = TEST_TTL / 2 + TEST_ELAPSED_REALTIME;
1920         doReturn(currentTime).when(mockDecoderClock).elapsedRealtime();
1921         doReturn(true).when(mockDeps).hasMessages(any(), eq(EVENT_START_QUERYTASK));
1922         processResponse(createResponse(
1923                 "service-instance-1", "192.0.2.123", 5353,
1924                 SERVICE_TYPE_LABELS,
1925                 Collections.emptyMap(), TEST_TTL), socketKey);
1926         verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
1927         assertNotNull(delayMessage);
1928         verifyAndSendQuery(12 /* index */, (long) (TEST_TTL / 2 * 0.8) /* timeInMs */,
1929                 true /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
1930                 14 /* scheduledCount */);
1931         currentTime += (long) (TEST_TTL / 2 * 0.8);
1932         doReturn(currentTime).when(mockDecoderClock).elapsedRealtime();
1933         verifyAndSendQuery(13 /* index */, 0 /* timeInMs */,
1934                 false /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
1935                 15 /* scheduledCount */);
1936         verifyAndSendQuery(14 /* index */, TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS,
1937                 false /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
1938                 16 /* scheduledCount */);
1939     }
1940 
1941     @Test
testSendQueryWithKnownAnswers()1942     public void testSendQueryWithKnownAnswers() throws Exception {
1943         client = makeMdnsServiceTypeClient(
1944                 MdnsFeatureFlags.newBuilder().setIsQueryWithKnownAnswerEnabled(true).build());
1945 
1946         doCallRealMethod().when(mockDeps).getDatagramPacketsFromMdnsPacket(
1947                 any(), any(MdnsPacket.class), any(InetSocketAddress.class), anyBoolean());
1948 
1949         startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
1950         InOrder inOrder = inOrder(mockListenerOne, mockSocketClient);
1951 
1952         final ArgumentCaptor<List<DatagramPacket>> queryCaptor =
1953                 ArgumentCaptor.forClass(List.class);
1954         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
1955         // Send twice for IPv4 and IPv6
1956         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingUnicastResponse(
1957                 queryCaptor.capture(), eq(socketKey), eq(false));
1958         verify(mockDeps, times(1)).sendMessage(any(), any(Message.class));
1959         assertNotNull(delayMessage);
1960 
1961         final MdnsPacket queryPacket = MdnsPacket.parse(
1962                 new MdnsPacketReader(queryCaptor.getValue().get(0)));
1963         assertTrue(hasQuestion(queryPacket, MdnsRecord.TYPE_PTR));
1964 
1965         // Process a response
1966         final String serviceName = "service-instance";
1967         final String ipV4Address = "192.0.2.0";
1968         final String[] subtypeLabels = Stream.concat(Stream.of("_subtype", "_sub"),
1969                         Arrays.stream(SERVICE_TYPE_LABELS)).toArray(String[]::new);
1970         final MdnsPacket packetWithoutSubtype = createResponse(
1971                 serviceName, ipV4Address, 5353, SERVICE_TYPE_LABELS,
1972                 Collections.emptyMap() /* textAttributes */, TEST_TTL);
1973         final MdnsPointerRecord originalPtr = (MdnsPointerRecord) CollectionUtils.findFirst(
1974                 packetWithoutSubtype.answers, r -> r instanceof MdnsPointerRecord);
1975 
1976         // Add a subtype PTR record
1977         final ArrayList<MdnsRecord> newAnswers = new ArrayList<>(packetWithoutSubtype.answers);
1978         newAnswers.add(new MdnsPointerRecord(subtypeLabels, originalPtr.getReceiptTime(),
1979                 originalPtr.getCacheFlush(), originalPtr.getTtl(), originalPtr.getPointer()));
1980         final MdnsPacket packetWithSubtype = new MdnsPacket(
1981                 packetWithoutSubtype.flags,
1982                 packetWithoutSubtype.questions,
1983                 newAnswers,
1984                 packetWithoutSubtype.authorityRecords,
1985                 packetWithoutSubtype.additionalRecords);
1986         processResponse(packetWithSubtype, socketKey);
1987 
1988         // Expect a query with known answers
1989         dispatchMessage();
1990         final ArgumentCaptor<List<DatagramPacket>> knownAnswersQueryCaptor =
1991                 ArgumentCaptor.forClass(List.class);
1992         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
1993         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingMulticastResponse(
1994                 knownAnswersQueryCaptor.capture(), eq(socketKey), eq(false));
1995 
1996         final MdnsPacket knownAnswersQueryPacket = MdnsPacket.parse(
1997                 new MdnsPacketReader(knownAnswersQueryCaptor.getValue().get(0)));
1998         assertTrue(hasQuestion(knownAnswersQueryPacket, MdnsRecord.TYPE_PTR, SERVICE_TYPE_LABELS));
1999         assertTrue(hasAnswer(knownAnswersQueryPacket, MdnsRecord.TYPE_PTR, SERVICE_TYPE_LABELS));
2000         assertFalse(hasAnswer(knownAnswersQueryPacket, MdnsRecord.TYPE_PTR, subtypeLabels));
2001     }
2002 
2003     @Test
testSendQueryWithSubTypeWithKnownAnswers()2004     public void testSendQueryWithSubTypeWithKnownAnswers() throws Exception {
2005         client = makeMdnsServiceTypeClient(
2006                 MdnsFeatureFlags.newBuilder().setIsQueryWithKnownAnswerEnabled(true).build());
2007 
2008         doCallRealMethod().when(mockDeps).getDatagramPacketsFromMdnsPacket(
2009                 any(), any(MdnsPacket.class), any(InetSocketAddress.class), anyBoolean());
2010 
2011         final MdnsSearchOptions options = MdnsSearchOptions.newBuilder()
2012                 .addSubtype("subtype").build();
2013         startSendAndReceive(mockListenerOne, options);
2014         InOrder inOrder = inOrder(mockListenerOne, mockSocketClient);
2015 
2016         final ArgumentCaptor<List<DatagramPacket>> queryCaptor =
2017                 ArgumentCaptor.forClass(List.class);
2018         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
2019         // Send twice for IPv4 and IPv6
2020         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingUnicastResponse(
2021                 queryCaptor.capture(), eq(socketKey), eq(false));
2022         verify(mockDeps, times(1)).sendMessage(any(), any(Message.class));
2023         assertNotNull(delayMessage);
2024 
2025         final MdnsPacket queryPacket = MdnsPacket.parse(
2026                 new MdnsPacketReader(queryCaptor.getValue().get(0)));
2027         final String[] subtypeLabels = Stream.concat(Stream.of("_subtype", "_sub"),
2028                 Arrays.stream(SERVICE_TYPE_LABELS)).toArray(String[]::new);
2029         assertTrue(hasQuestion(queryPacket, MdnsRecord.TYPE_PTR, SERVICE_TYPE_LABELS));
2030         assertTrue(hasQuestion(queryPacket, MdnsRecord.TYPE_PTR, subtypeLabels));
2031 
2032         // Process a response
2033         final String serviceName = "service-instance";
2034         final String ipV4Address = "192.0.2.0";
2035         final MdnsPacket packetWithoutSubtype = createResponse(
2036                 serviceName, ipV4Address, 5353, SERVICE_TYPE_LABELS,
2037                 Collections.emptyMap() /* textAttributes */, TEST_TTL);
2038         final MdnsPointerRecord originalPtr = (MdnsPointerRecord) CollectionUtils.findFirst(
2039                 packetWithoutSubtype.answers, r -> r instanceof MdnsPointerRecord);
2040 
2041         // Add a subtype PTR record
2042         final ArrayList<MdnsRecord> newAnswers = new ArrayList<>(packetWithoutSubtype.answers);
2043         newAnswers.add(new MdnsPointerRecord(subtypeLabels, originalPtr.getReceiptTime(),
2044                 originalPtr.getCacheFlush(), originalPtr.getTtl(), originalPtr.getPointer()));
2045         final MdnsPacket packetWithSubtype = new MdnsPacket(
2046                 packetWithoutSubtype.flags,
2047                 packetWithoutSubtype.questions,
2048                 newAnswers,
2049                 packetWithoutSubtype.authorityRecords,
2050                 packetWithoutSubtype.additionalRecords);
2051         processResponse(packetWithSubtype, socketKey);
2052 
2053         // Expect a query with known answers
2054         dispatchMessage();
2055         final ArgumentCaptor<List<DatagramPacket>> knownAnswersQueryCaptor =
2056                 ArgumentCaptor.forClass(List.class);
2057         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
2058         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingMulticastResponse(
2059                 knownAnswersQueryCaptor.capture(), eq(socketKey), eq(false));
2060 
2061         final MdnsPacket knownAnswersQueryPacket = MdnsPacket.parse(
2062                 new MdnsPacketReader(knownAnswersQueryCaptor.getValue().get(0)));
2063         assertTrue(hasQuestion(knownAnswersQueryPacket, MdnsRecord.TYPE_PTR, SERVICE_TYPE_LABELS));
2064         assertTrue(hasQuestion(knownAnswersQueryPacket, MdnsRecord.TYPE_PTR, subtypeLabels));
2065         assertTrue(hasAnswer(knownAnswersQueryPacket, MdnsRecord.TYPE_PTR, SERVICE_TYPE_LABELS));
2066         assertTrue(hasAnswer(knownAnswersQueryPacket, MdnsRecord.TYPE_PTR, subtypeLabels));
2067     }
2068 
2069     @Test
sendQueries_AggressiveQueryMode_ServiceInCache()2070     public void sendQueries_AggressiveQueryMode_ServiceInCache() {
2071         final int numOfQueriesBeforeBackoff = 11;
2072         final MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
2073                 .setQueryMode(AGGRESSIVE_QUERY_MODE)
2074                 .setNumOfQueriesBeforeBackoff(numOfQueriesBeforeBackoff)
2075                 .build();
2076         startSendAndReceive(mockListenerOne, searchOptions);
2077         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
2078 
2079         int burstCounter = 0;
2080         int betweenBurstTime = 0;
2081         for (int i = 0; i < numOfQueriesBeforeBackoff; i += 3) {
2082             verifyAndSendQuery(i, betweenBurstTime, /* expectsUnicastResponse= */ true);
2083             verifyAndSendQuery(i + 1, /* timeInMs= */ 0, /* expectsUnicastResponse= */ false);
2084             verifyAndSendQuery(i + 2, TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS,
2085                     /* expectsUnicastResponse= */ false);
2086             betweenBurstTime = Math.min(
2087                     INITIAL_AGGRESSIVE_TIME_BETWEEN_BURSTS_MS * (int) Math.pow(2, burstCounter),
2088                     MAX_TIME_BETWEEN_AGGRESSIVE_BURSTS_MS);
2089             burstCounter++;
2090         }
2091         // In backoff mode, the current scheduled task will be canceled and reschedule if the
2092         // 0.8 * smallestRemainingTtl is larger than time to next run.
2093         long currentTime = TEST_TTL / 2 + TEST_ELAPSED_REALTIME;
2094         doReturn(currentTime).when(mockDecoderClock).elapsedRealtime();
2095         doReturn(true).when(mockDeps).hasMessages(any(), eq(EVENT_START_QUERYTASK));
2096         processResponse(createResponse(
2097                 "service-instance-1", "192.0.2.123", 5353,
2098                 SERVICE_TYPE_LABELS,
2099                 Collections.emptyMap(), TEST_TTL), socketKey);
2100         verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
2101         assertNotNull(delayMessage);
2102         assertEquals((long) (TEST_TTL / 2 * 0.8), latestDelayMs);
2103 
2104         // Register another listener. There is a service in cache, the query time should be
2105         // rescheduled with previous run.
2106         currentTime += (long) ((TEST_TTL / 2 * 0.8) - 500L);
2107         doReturn(currentTime).when(mockDecoderClock).elapsedRealtime();
2108         startSendAndReceive(mockListenerTwo, searchOptions);
2109         verify(mockDeps, times(3)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
2110         assertNotNull(delayMessage);
2111         assertEquals(500L, latestDelayMs);
2112 
2113         // Stop all listeners
2114         stopSendAndReceive(mockListenerOne);
2115         stopSendAndReceive(mockListenerTwo);
2116         verify(mockDeps, times(4)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
2117 
2118         // Register a new listener. There is a service in cache, the query time should be
2119         // rescheduled with remaining ttl.
2120         currentTime += 400L;
2121         doReturn(currentTime).when(mockDecoderClock).elapsedRealtime();
2122         startSendAndReceive(mockListenerOne, searchOptions);
2123         assertNotNull(delayMessage);
2124         assertEquals(9680L, latestDelayMs);
2125     }
2126 
2127     @Test
sendQueries_AccurateDelayCallback()2128     public void sendQueries_AccurateDelayCallback() {
2129         client = makeMdnsServiceTypeClient(
2130                 MdnsFeatureFlags.newBuilder().setIsAccurateDelayCallbackEnabled(true).build());
2131 
2132         final int numOfQueriesBeforeBackoff = 2;
2133         final MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
2134                 .addSubtype(SUBTYPE)
2135                 .setQueryMode(AGGRESSIVE_QUERY_MODE)
2136                 .setNumOfQueriesBeforeBackoff(numOfQueriesBeforeBackoff)
2137                 .build();
2138         startSendAndReceive(mockListenerOne, searchOptions);
2139         verify(mockScheduler, times(1)).removeDelayedMessage(EVENT_START_QUERYTASK);
2140 
2141         // Verify that the first query has been sent.
2142         verifyAndSendQuery(0 /* index */, 0 /* timeInMs */, true /* expectsUnicastResponse */,
2143                 true /* multipleSocketDiscovery */, 1 /* scheduledCount */,
2144                 1 /* sendMessageCount */, true /* useAccurateDelayCallback */);
2145 
2146         // Verify that the second query has been sent
2147         verifyAndSendQuery(1 /* index */, 0 /* timeInMs */, false /* expectsUnicastResponse */,
2148                 true /* multipleSocketDiscovery */, 2 /* scheduledCount */,
2149                 2 /* sendMessageCount */, true /* useAccurateDelayCallback */);
2150 
2151         // Verify that the third query has been sent
2152         verifyAndSendQuery(2 /* index */, TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS,
2153                 false /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
2154                 3 /* scheduledCount */, 3 /* sendMessageCount */,
2155                 true /* useAccurateDelayCallback */);
2156 
2157         // In backoff mode, the current scheduled task will be canceled and reschedule if the
2158         // 0.8 * smallestRemainingTtl is larger than time to next run.
2159         long currentTime = TEST_TTL / 2 + TEST_ELAPSED_REALTIME;
2160         doReturn(currentTime).when(mockDecoderClock).elapsedRealtime();
2161         doReturn(true).when(mockScheduler).hasDelayedMessage(EVENT_START_QUERYTASK);
2162         processResponse(createResponse(
2163                 "service-instance-1", "192.0.2.123", 5353,
2164                 SERVICE_TYPE_LABELS,
2165                 Collections.emptyMap(), TEST_TTL), socketKey);
2166         // Verify that the message removal occurred.
2167         verify(mockScheduler, times(6)).removeDelayedMessage(EVENT_START_QUERYTASK);
2168         assertNotNull(message);
2169         verifyAndSendQuery(3 /* index */, (long) (TEST_TTL / 2 * 0.8) /* timeInMs */,
2170                 true /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
2171                 5 /* scheduledCount */, 4 /* sendMessageCount */,
2172                 true /* useAccurateDelayCallback */);
2173 
2174         // Stop sending packets.
2175         stopSendAndReceive(mockListenerOne);
2176         verify(mockScheduler, times(8)).removeDelayedMessage(EVENT_START_QUERYTASK);
2177     }
2178 
2179     @Test
testTimerFdCloseProperly()2180     public void testTimerFdCloseProperly() {
2181         client = makeMdnsServiceTypeClient(
2182                 MdnsFeatureFlags.newBuilder().setIsAccurateDelayCallbackEnabled(true).build());
2183 
2184         // Start query
2185         startSendAndReceive(mockListenerOne, MdnsSearchOptions.newBuilder().build());
2186         verify(mockScheduler, times(1)).removeDelayedMessage(EVENT_START_QUERYTASK);
2187 
2188         // Stop query and verify the close() method has been called.
2189         stopSendAndReceive(mockListenerOne);
2190         verify(mockScheduler, times(2)).removeDelayedMessage(EVENT_START_QUERYTASK);
2191         verify(mockScheduler).close();
2192     }
2193 
matchServiceName(String name)2194     private static MdnsServiceInfo matchServiceName(String name) {
2195         return argThat(info -> info.getServiceInstanceName().equals(name));
2196     }
2197 
2198     // verifies that the right query was enqueued with the right delay, and send query by executing
2199     // the runnable.
verifyAndSendQuery(int index, long timeInMs, boolean expectsUnicastResponse)2200     private void verifyAndSendQuery(int index, long timeInMs, boolean expectsUnicastResponse) {
2201         verifyAndSendQuery(index, timeInMs, expectsUnicastResponse,
2202                 true /* multipleSocketDiscovery */, index + 1 /* scheduledCount */);
2203     }
2204 
verifyAndSendQuery(int index, long timeInMs, boolean expectsUnicastResponse, boolean multipleSocketDiscovery, int scheduledCount)2205     private void verifyAndSendQuery(int index, long timeInMs, boolean expectsUnicastResponse,
2206             boolean multipleSocketDiscovery, int scheduledCount) {
2207         verifyAndSendQuery(index, timeInMs, expectsUnicastResponse,
2208                 multipleSocketDiscovery, scheduledCount, index + 1 /* sendMessageCount */,
2209                 false /* useAccurateDelayCallback */);
2210     }
2211 
verifyAndSendQuery(int index, long timeInMs, boolean expectsUnicastResponse, boolean multipleSocketDiscovery, int scheduledCount, int sendMessageCount, boolean useAccurateDelayCallback)2212     private void verifyAndSendQuery(int index, long timeInMs, boolean expectsUnicastResponse,
2213             boolean multipleSocketDiscovery, int scheduledCount, int sendMessageCount,
2214             boolean useAccurateDelayCallback) {
2215         if (useAccurateDelayCallback && message != null && realHandler != null) {
2216             runOnHandler(() -> realHandler.dispatchMessage(message));
2217             message = null;
2218         } else {
2219             // Dispatch the message
2220             if (delayMessage != null && realHandler != null) {
2221                 dispatchMessage();
2222             }
2223         }
2224         assertEquals(timeInMs, latestDelayMs);
2225         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
2226         if (expectsUnicastResponse) {
2227             verify(mockSocketClient).sendPacketRequestingUnicastResponse(
2228                     argThat(pkts -> pkts.get(0).equals(expectedIPv4Packets[index])),
2229                     eq(socketKey), eq(false));
2230             if (multipleSocketDiscovery) {
2231                 verify(mockSocketClient).sendPacketRequestingUnicastResponse(
2232                         argThat(pkts -> pkts.get(0).equals(expectedIPv6Packets[index])),
2233                         eq(socketKey), eq(false));
2234             }
2235         } else {
2236             verify(mockSocketClient).sendPacketRequestingMulticastResponse(
2237                     argThat(pkts -> pkts.get(0).equals(expectedIPv4Packets[index])),
2238                     eq(socketKey), eq(false));
2239             if (multipleSocketDiscovery) {
2240                 verify(mockSocketClient).sendPacketRequestingMulticastResponse(
2241                         argThat(pkts -> pkts.get(0).equals(expectedIPv6Packets[index])),
2242                         eq(socketKey), eq(false));
2243             }
2244         }
2245         verify(mockDeps, times(sendMessageCount))
2246                 .sendMessage(any(Handler.class), any(Message.class));
2247         // Verify the task has been scheduled.
2248         if (useAccurateDelayCallback) {
2249             verify(mockScheduler, times(scheduledCount)).sendDelayedMessage(any(), anyLong());
2250         } else {
2251             verify(mockDeps, times(scheduledCount))
2252                     .sendMessageDelayed(any(Handler.class), any(Message.class), anyLong());
2253         }
2254     }
2255 
getTestServiceName(String instanceName)2256     private static String[] getTestServiceName(String instanceName) {
2257         return Stream.concat(Stream.of(instanceName),
2258                 Arrays.stream(SERVICE_TYPE_LABELS)).toArray(String[]::new);
2259     }
2260 
getServiceTypeWithSubtype(String subtype)2261     private static String[] getServiceTypeWithSubtype(String subtype) {
2262         return Stream.concat(Stream.of(subtype, "_sub"),
2263                 Arrays.stream(SERVICE_TYPE_LABELS)).toArray(String[]::new);
2264     }
2265 
hasQuestion(MdnsPacket packet, int type)2266     private static boolean hasQuestion(MdnsPacket packet, int type) {
2267         return hasQuestion(packet, type, null);
2268     }
2269 
hasQuestion(MdnsPacket packet, int type, @Nullable String[] name)2270     private static boolean hasQuestion(MdnsPacket packet, int type, @Nullable String[] name) {
2271         return packet.questions.stream().anyMatch(q -> q.getType() == type
2272                 && (name == null || Arrays.equals(q.name, name)));
2273     }
2274 
hasAnswer(MdnsPacket packet, int type, @NonNull String[] name)2275     private static boolean hasAnswer(MdnsPacket packet, int type, @NonNull String[] name) {
2276         return packet.answers.stream().anyMatch(q -> {
2277             return q.getType() == type && (Arrays.equals(q.name, name));
2278         });
2279     }
2280 
2281     // A fake ScheduledExecutorService that keeps tracking the last scheduled Runnable and its delay
2282     // time.
2283     private class FakeExecutor extends ScheduledThreadPoolExecutor {
2284         private long lastScheduledDelayInMs;
2285         private Runnable lastScheduledRunnable;
2286         private Runnable lastSubmittedRunnable;
2287         private Future<?> lastSubmittedFuture;
2288         private int futureIndex;
2289 
FakeExecutor()2290         FakeExecutor() {
2291             super(1);
2292             lastScheduledDelayInMs = -1;
2293         }
2294 
2295         @Override
submit(Runnable command)2296         public Future<?> submit(Runnable command) {
2297             Future<?> future = super.submit(command);
2298             lastSubmittedRunnable = command;
2299             lastSubmittedFuture = future;
2300             return future;
2301         }
2302 
2303         // Don't call through the real implementation, just track the scheduled Runnable, and
2304         // returns a ScheduledFuture.
2305         @Override
schedule(Runnable command, long delay, TimeUnit unit)2306         public ScheduledFuture<?> schedule(Runnable command, long delay, TimeUnit unit) {
2307             lastScheduledDelayInMs = delay;
2308             lastScheduledRunnable = command;
2309             return Mockito.mock(ScheduledFuture.class);
2310         }
2311 
2312         // Returns the delay of the last scheduled task, and clear it.
getAndClearLastScheduledDelayInMs()2313         long getAndClearLastScheduledDelayInMs() {
2314             long val = lastScheduledDelayInMs;
2315             lastScheduledDelayInMs = -1;
2316             return val;
2317         }
2318 
2319         // Returns the last scheduled task, and clear it.
getAndClearLastScheduledRunnable()2320         Runnable getAndClearLastScheduledRunnable() {
2321             Runnable val = lastScheduledRunnable;
2322             lastScheduledRunnable = null;
2323             return val;
2324         }
2325 
getAndClearSubmittedRunnable()2326         Runnable getAndClearSubmittedRunnable() {
2327             Runnable val = lastSubmittedRunnable;
2328             lastSubmittedRunnable = null;
2329             return val;
2330         }
2331 
getAndClearSubmittedFuture()2332         Future<?> getAndClearSubmittedFuture() {
2333             Future<?> val = lastSubmittedFuture;
2334             lastSubmittedFuture = null;
2335             return val;
2336         }
2337     }
2338 
createResponse( @onNull String serviceInstanceName, @Nullable String host, int port, @NonNull String subtype, @NonNull Map<String, String> textAttributes, long ptrTtlMillis)2339     private MdnsPacket createResponse(
2340             @NonNull String serviceInstanceName,
2341             @Nullable String host,
2342             int port,
2343             @NonNull String subtype,
2344             @NonNull Map<String, String> textAttributes,
2345             long ptrTtlMillis)
2346             throws Exception {
2347         final ArrayList<String> type = new ArrayList<>();
2348         type.add(subtype);
2349         type.add(MdnsConstants.SUBTYPE_LABEL);
2350         type.addAll(Arrays.asList(SERVICE_TYPE_LABELS));
2351         return createResponse(serviceInstanceName, host, port, type.toArray(new String[0]),
2352                 textAttributes, ptrTtlMillis);
2353     }
2354 
2355 
createResponse( @onNull String serviceInstanceName, @Nullable String host, int port, @NonNull String[] type, @NonNull Map<String, String> textAttributes, long ptrTtlMillis)2356     private MdnsPacket createResponse(
2357             @NonNull String serviceInstanceName,
2358             @Nullable String host,
2359             int port,
2360             @NonNull String[] type,
2361             @NonNull Map<String, String> textAttributes,
2362             long ptrTtlMillis) {
2363         return createResponse(serviceInstanceName, host, port, type, textAttributes, ptrTtlMillis,
2364                 TEST_ELAPSED_REALTIME);
2365     }
2366 
2367     // Creates a mDNS response.
createResponse( @onNull String serviceInstanceName, @Nullable String host, int port, @NonNull String[] type, @NonNull Map<String, String> textAttributes, long ptrTtlMillis, long receiptTimeMillis)2368     private MdnsPacket createResponse(
2369             @NonNull String serviceInstanceName,
2370             @Nullable String host,
2371             int port,
2372             @NonNull String[] type,
2373             @NonNull Map<String, String> textAttributes,
2374             long ptrTtlMillis,
2375             long receiptTimeMillis) {
2376 
2377         final ArrayList<MdnsRecord> answerRecords = new ArrayList<>();
2378 
2379         // Set PTR record
2380         final ArrayList<String> serviceNameList = new ArrayList<>();
2381         serviceNameList.add(serviceInstanceName);
2382         serviceNameList.addAll(Arrays.asList(type));
2383         final String[] serviceName = serviceNameList.toArray(new String[0]);
2384         final MdnsPointerRecord pointerRecord = new MdnsPointerRecord(
2385                 type,
2386                 receiptTimeMillis,
2387                 false /* cacheFlush */,
2388                 ptrTtlMillis,
2389                 serviceName);
2390         answerRecords.add(pointerRecord);
2391 
2392         // Set SRV record.
2393         final MdnsServiceRecord serviceRecord = new MdnsServiceRecord(
2394                 serviceName,
2395                 receiptTimeMillis,
2396                 false /* cacheFlush */,
2397                 TEST_TTL,
2398                 0 /* servicePriority */,
2399                 0 /* serviceWeight */,
2400                 port,
2401                 new String[]{"hostname"});
2402         answerRecords.add(serviceRecord);
2403 
2404         // Set A/AAAA record.
2405         if (host != null) {
2406             final InetAddress addr = InetAddresses.parseNumericAddress(host);
2407             final MdnsInetAddressRecord inetAddressRecord = new MdnsInetAddressRecord(
2408                     new String[] {"hostname"} /* name */,
2409                     receiptTimeMillis,
2410                     false /* cacheFlush */,
2411                     TEST_TTL,
2412                     addr);
2413             answerRecords.add(inetAddressRecord);
2414         }
2415 
2416         // Set TXT record.
2417         final List<TextEntry> textEntries = new ArrayList<>();
2418         for (Map.Entry<String, String> kv : textAttributes.entrySet()) {
2419             textEntries.add(new TextEntry(kv.getKey(), kv.getValue().getBytes(UTF_8)));
2420         }
2421         final MdnsTextRecord textRecord = new MdnsTextRecord(
2422                 serviceName,
2423                 receiptTimeMillis,
2424                 false /* cacheFlush */,
2425                 TEST_TTL,
2426                 textEntries);
2427         answerRecords.add(textRecord);
2428         return new MdnsPacket(
2429                 0 /* flags */,
2430                 Collections.emptyList() /* questions */,
2431                 answerRecords,
2432                 Collections.emptyList() /* authorityRecords */,
2433                 Collections.emptyList() /* additionalRecords */
2434         );
2435     }
2436 }