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