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.MdnsServiceCache.ServiceExpiredCallback;
20 import static com.android.server.connectivity.mdns.MdnsServiceCache.findMatchedResponse;
21 import static com.android.server.connectivity.mdns.util.MdnsUtils.Clock;
22 import static com.android.server.connectivity.mdns.util.MdnsUtils.ensureRunningOnHandlerThread;
23 
24 import android.annotation.NonNull;
25 import android.annotation.Nullable;
26 import android.os.Handler;
27 import android.os.Looper;
28 import android.os.Message;
29 import android.text.TextUtils;
30 import android.util.ArrayMap;
31 import android.util.Pair;
32 
33 import androidx.annotation.VisibleForTesting;
34 
35 import com.android.net.module.util.CollectionUtils;
36 import com.android.net.module.util.SharedLog;
37 import com.android.server.connectivity.mdns.util.MdnsUtils;
38 
39 import java.io.IOException;
40 import java.io.PrintWriter;
41 import java.net.DatagramPacket;
42 import java.net.Inet4Address;
43 import java.net.Inet6Address;
44 import java.net.InetSocketAddress;
45 import java.time.Instant;
46 import java.util.ArrayList;
47 import java.util.Arrays;
48 import java.util.Collection;
49 import java.util.Collections;
50 import java.util.Iterator;
51 import java.util.List;
52 import java.util.Set;
53 import java.util.concurrent.ScheduledExecutorService;
54 
55 /**
56  * Instance of this class sends and receives mDNS packets of a given service type and invoke
57  * registered {@link MdnsServiceBrowserListener} instances.
58  */
59 public class MdnsServiceTypeClient {
60 
61     private static final String TAG = MdnsServiceTypeClient.class.getSimpleName();
62     @VisibleForTesting
63     static final int EVENT_START_QUERYTASK = 1;
64     static final int EVENT_QUERY_RESULT = 2;
65     static final int INVALID_TRANSACTION_ID = -1;
66 
67     private final String serviceType;
68     private final String[] serviceTypeLabels;
69     private final MdnsSocketClientBase socketClient;
70     private final MdnsResponseDecoder responseDecoder;
71     private final ScheduledExecutorService executor;
72     @NonNull private final SocketKey socketKey;
73     @NonNull private final SharedLog sharedLog;
74     @NonNull private final Handler handler;
75     @NonNull private final MdnsQueryScheduler mdnsQueryScheduler;
76     @NonNull private final Dependencies dependencies;
77     /**
78      * The service caches for each socket. It should be accessed from looper thread only.
79      */
80     @NonNull private final MdnsServiceCache serviceCache;
81     @NonNull private final MdnsServiceCache.CacheKey cacheKey;
82     @NonNull private final ServiceExpiredCallback serviceExpiredCallback =
83             new ServiceExpiredCallback() {
84                 @Override
85                 public void onServiceRecordExpired(@NonNull MdnsResponse previousResponse,
86                         @Nullable MdnsResponse newResponse) {
87                     notifyRemovedServiceToListeners(previousResponse, "Service record expired");
88                 }
89             };
90     @NonNull private final MdnsFeatureFlags featureFlags;
91     private final ArrayMap<MdnsServiceBrowserListener, ListenerInfo> listeners =
92             new ArrayMap<>();
93     private final boolean removeServiceAfterTtlExpires =
94             MdnsConfigs.removeServiceAfterTtlExpires();
95     private final Clock clock;
96 
97     @Nullable private MdnsSearchOptions searchOptions;
98 
99     // The session ID increases when startSendAndReceive() is called where we schedule a
100     // QueryTask for
101     // new subtypes. It stays the same between packets for same subtypes.
102     private long currentSessionId = 0;
103     private long lastSentTime;
104 
105     private static class ListenerInfo {
106         @NonNull
107         final MdnsSearchOptions searchOptions;
108         final Set<String> discoveredServiceNames;
109 
ListenerInfo(@onNull MdnsSearchOptions searchOptions, @Nullable ListenerInfo previousInfo)110         ListenerInfo(@NonNull MdnsSearchOptions searchOptions,
111                 @Nullable ListenerInfo previousInfo) {
112             this.searchOptions = searchOptions;
113             this.discoveredServiceNames = previousInfo == null
114                     ? MdnsUtils.newSet() : previousInfo.discoveredServiceNames;
115         }
116 
117         /**
118          * Set the given service name as discovered.
119          *
120          * @return true if the service name was not discovered before.
121          */
setServiceDiscovered(@onNull String serviceName)122         boolean setServiceDiscovered(@NonNull String serviceName) {
123             return discoveredServiceNames.add(MdnsUtils.toDnsLowerCase(serviceName));
124         }
125 
unsetServiceDiscovered(@onNull String serviceName)126         void unsetServiceDiscovered(@NonNull String serviceName) {
127             discoveredServiceNames.remove(MdnsUtils.toDnsLowerCase(serviceName));
128         }
129     }
130 
131     private class QueryTaskHandler extends Handler {
QueryTaskHandler(Looper looper)132         QueryTaskHandler(Looper looper) {
133             super(looper);
134         }
135 
136         @Override
137         @SuppressWarnings("FutureReturnValueIgnored")
handleMessage(Message msg)138         public void handleMessage(Message msg) {
139             switch (msg.what) {
140                 case EVENT_START_QUERYTASK: {
141                     final MdnsQueryScheduler.ScheduledQueryTaskArgs taskArgs =
142                             (MdnsQueryScheduler.ScheduledQueryTaskArgs) msg.obj;
143                     // QueryTask should be run immediately after being created (not be scheduled in
144                     // advance). Because the result of "makeResponsesForResolve" depends on answers
145                     // that were received before it is called, so to take into account all answers
146                     // before sending the query, it needs to be called just before sending it.
147                     final List<MdnsResponse> servicesToResolve = makeResponsesForResolve(socketKey);
148                     final QueryTask queryTask = new QueryTask(taskArgs, servicesToResolve,
149                             getAllDiscoverySubtypes(), needSendDiscoveryQueries(listeners),
150                             getExistingServices());
151                     executor.submit(queryTask);
152                     break;
153                 }
154                 case EVENT_QUERY_RESULT: {
155                     final QuerySentArguments sentResult = (QuerySentArguments) msg.obj;
156                     // If a task is cancelled while the Executor is running it, EVENT_QUERY_RESULT
157                     // will still be sent when it ends. So use session ID to check if this task
158                     // should continue to schedule more.
159                     if (sentResult.taskArgs.sessionId != currentSessionId) {
160                         break;
161                     }
162 
163                     if ((sentResult.transactionId != INVALID_TRANSACTION_ID)) {
164                         for (int i = 0; i < listeners.size(); i++) {
165                             listeners.keyAt(i).onDiscoveryQuerySent(
166                                     sentResult.subTypes, sentResult.transactionId);
167                         }
168                     }
169 
170                     tryRemoveServiceAfterTtlExpires();
171 
172                     final long now = clock.elapsedRealtime();
173                     lastSentTime = now;
174                     final long minRemainingTtl = getMinRemainingTtl(now);
175                     MdnsQueryScheduler.ScheduledQueryTaskArgs args =
176                             mdnsQueryScheduler.scheduleNextRun(
177                                     sentResult.taskArgs.config,
178                                     minRemainingTtl,
179                                     now,
180                                     lastSentTime,
181                                     sentResult.taskArgs.sessionId
182                             );
183                     dependencies.sendMessageDelayed(
184                             handler,
185                             handler.obtainMessage(EVENT_START_QUERYTASK, args),
186                             calculateTimeToNextTask(args, now, sharedLog));
187                     break;
188                 }
189                 default:
190                     sharedLog.e("Unrecognized event " + msg.what);
191                     break;
192             }
193         }
194     }
195 
196     /**
197      * Dependencies of MdnsServiceTypeClient, for injection in tests.
198      */
199     @VisibleForTesting(otherwise = VisibleForTesting.PACKAGE_PRIVATE)
200     public static class Dependencies {
201         /**
202          * @see Handler#sendMessageDelayed(Message, long)
203          */
sendMessageDelayed(@onNull Handler handler, @NonNull Message message, long delayMillis)204         public void sendMessageDelayed(@NonNull Handler handler, @NonNull Message message,
205                 long delayMillis) {
206             handler.sendMessageDelayed(message, delayMillis);
207         }
208 
209         /**
210          * @see Handler#removeMessages(int)
211          */
removeMessages(@onNull Handler handler, int what)212         public void removeMessages(@NonNull Handler handler, int what) {
213             handler.removeMessages(what);
214         }
215 
216         /**
217          * @see Handler#hasMessages(int)
218          */
hasMessages(@onNull Handler handler, int what)219         public boolean hasMessages(@NonNull Handler handler, int what) {
220             return handler.hasMessages(what);
221         }
222 
223         /**
224          * @see Handler#post(Runnable)
225          */
sendMessage(@onNull Handler handler, @NonNull Message message)226         public void sendMessage(@NonNull Handler handler, @NonNull Message message) {
227             handler.sendMessage(message);
228         }
229 
230         /**
231          * Generate the DatagramPackets from given MdnsPacket and InetSocketAddress.
232          *
233          * <p> If the query with known answer feature is enabled and the MdnsPacket is too large for
234          *     a single DatagramPacket, it will be split into multiple DatagramPackets.
235          */
getDatagramPacketsFromMdnsPacket( @onNull byte[] packetCreationBuffer, @NonNull MdnsPacket packet, @NonNull InetSocketAddress address, boolean isQueryWithKnownAnswer)236         public List<DatagramPacket> getDatagramPacketsFromMdnsPacket(
237                 @NonNull byte[] packetCreationBuffer, @NonNull MdnsPacket packet,
238                 @NonNull InetSocketAddress address, boolean isQueryWithKnownAnswer)
239                 throws IOException {
240             if (isQueryWithKnownAnswer) {
241                 return MdnsUtils.createQueryDatagramPackets(packetCreationBuffer, packet, address);
242             } else {
243                 final byte[] queryBuffer =
244                         MdnsUtils.createRawDnsPacket(packetCreationBuffer, packet);
245                 return List.of(new DatagramPacket(queryBuffer, 0, queryBuffer.length, address));
246             }
247         }
248     }
249 
250     /**
251      * Constructor of {@link MdnsServiceTypeClient}.
252      *
253      * @param socketClient Sends and receives mDNS packet.
254      * @param executor         A {@link ScheduledExecutorService} used to schedule query tasks.
255      */
MdnsServiceTypeClient( @onNull String serviceType, @NonNull MdnsSocketClientBase socketClient, @NonNull ScheduledExecutorService executor, @NonNull SocketKey socketKey, @NonNull SharedLog sharedLog, @NonNull Looper looper, @NonNull MdnsServiceCache serviceCache, @NonNull MdnsFeatureFlags featureFlags)256     public MdnsServiceTypeClient(
257             @NonNull String serviceType,
258             @NonNull MdnsSocketClientBase socketClient,
259             @NonNull ScheduledExecutorService executor,
260             @NonNull SocketKey socketKey,
261             @NonNull SharedLog sharedLog,
262             @NonNull Looper looper,
263             @NonNull MdnsServiceCache serviceCache,
264             @NonNull MdnsFeatureFlags featureFlags) {
265         this(serviceType, socketClient, executor, new Clock(), socketKey, sharedLog, looper,
266                 new Dependencies(), serviceCache, featureFlags);
267     }
268 
269     @VisibleForTesting
MdnsServiceTypeClient( @onNull String serviceType, @NonNull MdnsSocketClientBase socketClient, @NonNull ScheduledExecutorService executor, @NonNull Clock clock, @NonNull SocketKey socketKey, @NonNull SharedLog sharedLog, @NonNull Looper looper, @NonNull Dependencies dependencies, @NonNull MdnsServiceCache serviceCache, @NonNull MdnsFeatureFlags featureFlags)270     public MdnsServiceTypeClient(
271             @NonNull String serviceType,
272             @NonNull MdnsSocketClientBase socketClient,
273             @NonNull ScheduledExecutorService executor,
274             @NonNull Clock clock,
275             @NonNull SocketKey socketKey,
276             @NonNull SharedLog sharedLog,
277             @NonNull Looper looper,
278             @NonNull Dependencies dependencies,
279             @NonNull MdnsServiceCache serviceCache,
280             @NonNull MdnsFeatureFlags featureFlags) {
281         this.serviceType = serviceType;
282         this.socketClient = socketClient;
283         this.executor = executor;
284         this.serviceTypeLabels = TextUtils.split(serviceType, "\\.");
285         this.responseDecoder = new MdnsResponseDecoder(clock, serviceTypeLabels);
286         this.clock = clock;
287         this.socketKey = socketKey;
288         this.sharedLog = sharedLog;
289         this.handler = new QueryTaskHandler(looper);
290         this.dependencies = dependencies;
291         this.serviceCache = serviceCache;
292         this.mdnsQueryScheduler = new MdnsQueryScheduler();
293         this.cacheKey = new MdnsServiceCache.CacheKey(serviceType, socketKey);
294         this.featureFlags = featureFlags;
295     }
296 
297     /**
298      * Do the cleanup of the MdnsServiceTypeClient
299      */
shutDown()300     private void shutDown() {
301         removeScheduledTask();
302         mdnsQueryScheduler.cancelScheduledRun();
303         serviceCache.unregisterServiceExpiredCallback(cacheKey);
304     }
305 
buildMdnsServiceInfoFromResponse( @onNull MdnsResponse response, @NonNull String[] serviceTypeLabels)306     private static MdnsServiceInfo buildMdnsServiceInfoFromResponse(
307             @NonNull MdnsResponse response, @NonNull String[] serviceTypeLabels) {
308         String[] hostName = null;
309         int port = 0;
310         if (response.hasServiceRecord()) {
311             hostName = response.getServiceRecord().getServiceHost();
312             port = response.getServiceRecord().getServicePort();
313         }
314 
315         final List<String> ipv4Addresses = new ArrayList<>();
316         final List<String> ipv6Addresses = new ArrayList<>();
317         if (response.hasInet4AddressRecord()) {
318             for (MdnsInetAddressRecord inetAddressRecord : response.getInet4AddressRecords()) {
319                 final Inet4Address inet4Address = inetAddressRecord.getInet4Address();
320                 ipv4Addresses.add((inet4Address == null) ? null : inet4Address.getHostAddress());
321             }
322         }
323         if (response.hasInet6AddressRecord()) {
324             for (MdnsInetAddressRecord inetAddressRecord : response.getInet6AddressRecords()) {
325                 final Inet6Address inet6Address = inetAddressRecord.getInet6Address();
326                 ipv6Addresses.add((inet6Address == null) ? null : inet6Address.getHostAddress());
327             }
328         }
329         String serviceInstanceName = response.getServiceInstanceName();
330         if (serviceInstanceName == null) {
331             throw new IllegalStateException(
332                     "mDNS response must have non-null service instance name");
333         }
334         List<String> textStrings = null;
335         List<MdnsServiceInfo.TextEntry> textEntries = null;
336         if (response.hasTextRecord()) {
337             textStrings = response.getTextRecord().getStrings();
338             textEntries = response.getTextRecord().getEntries();
339         }
340         Instant now = Instant.now();
341         // TODO: Throw an error message if response doesn't have Inet6 or Inet4 address.
342         return new MdnsServiceInfo(
343                 serviceInstanceName,
344                 serviceTypeLabels,
345                 response.getSubtypes(),
346                 hostName,
347                 port,
348                 ipv4Addresses,
349                 ipv6Addresses,
350                 textStrings,
351                 textEntries,
352                 response.getInterfaceIndex(),
353                 response.getNetwork(),
354                 now.plusMillis(response.getMinRemainingTtl(now.toEpochMilli())));
355     }
356 
getExistingServices()357     private List<MdnsResponse> getExistingServices() {
358         return featureFlags.isQueryWithKnownAnswerEnabled()
359                 ? serviceCache.getCachedServices(cacheKey) : Collections.emptyList();
360     }
361 
362     /**
363      * Registers {@code listener} for receiving discovery event of mDNS service instances, and
364      * starts
365      * (or continue) to send mDNS queries periodically.
366      *
367      * @param listener      The {@link MdnsServiceBrowserListener} to register.
368      * @param searchOptions {@link MdnsSearchOptions} contains the list of subtypes to discover.
369      */
370     @SuppressWarnings("FutureReturnValueIgnored")
startSendAndReceive( @onNull MdnsServiceBrowserListener listener, @NonNull MdnsSearchOptions searchOptions)371     public void startSendAndReceive(
372             @NonNull MdnsServiceBrowserListener listener,
373             @NonNull MdnsSearchOptions searchOptions) {
374         ensureRunningOnHandlerThread(handler);
375         this.searchOptions = searchOptions;
376         boolean hadReply = false;
377         final ListenerInfo existingInfo = listeners.get(listener);
378         final ListenerInfo listenerInfo = new ListenerInfo(searchOptions, existingInfo);
379         listeners.put(listener, listenerInfo);
380         if (existingInfo == null) {
381             for (MdnsResponse existingResponse : serviceCache.getCachedServices(cacheKey)) {
382                 if (!responseMatchesOptions(existingResponse, searchOptions)) continue;
383                 final MdnsServiceInfo info =
384                         buildMdnsServiceInfoFromResponse(existingResponse, serviceTypeLabels);
385                 listener.onServiceNameDiscovered(info, true /* isServiceFromCache */);
386                 listenerInfo.setServiceDiscovered(info.getServiceInstanceName());
387                 if (existingResponse.isComplete()) {
388                     listener.onServiceFound(info, true /* isServiceFromCache */);
389                     hadReply = true;
390                 }
391             }
392         }
393         // Remove the next scheduled periodical task.
394         removeScheduledTask();
395         mdnsQueryScheduler.cancelScheduledRun();
396         // Keep tracking the ScheduledFuture for the task so we can cancel it if caller is not
397         // interested anymore.
398         final QueryTaskConfig taskConfig = new QueryTaskConfig(
399                 searchOptions.getQueryMode(),
400                 searchOptions.onlyUseIpv6OnIpv6OnlyNetworks(),
401                 searchOptions.numOfQueriesBeforeBackoff(),
402                 socketKey);
403         final long now = clock.elapsedRealtime();
404         if (lastSentTime == 0) {
405             lastSentTime = now;
406         }
407         final long minRemainingTtl = getMinRemainingTtl(now);
408         if (hadReply) {
409             MdnsQueryScheduler.ScheduledQueryTaskArgs args =
410                     mdnsQueryScheduler.scheduleNextRun(
411                             taskConfig,
412                             minRemainingTtl,
413                             now,
414                             lastSentTime,
415                             currentSessionId
416                     );
417             dependencies.sendMessageDelayed(
418                     handler,
419                     handler.obtainMessage(EVENT_START_QUERYTASK, args),
420                     calculateTimeToNextTask(args, now, sharedLog));
421         } else {
422             final List<MdnsResponse> servicesToResolve = makeResponsesForResolve(socketKey);
423             final QueryTask queryTask = new QueryTask(
424                     mdnsQueryScheduler.scheduleFirstRun(taskConfig, now,
425                             minRemainingTtl, currentSessionId), servicesToResolve,
426                     getAllDiscoverySubtypes(), needSendDiscoveryQueries(listeners),
427                     getExistingServices());
428             executor.submit(queryTask);
429         }
430 
431         serviceCache.registerServiceExpiredCallback(cacheKey, serviceExpiredCallback);
432     }
433 
getAllDiscoverySubtypes()434     private Set<String> getAllDiscoverySubtypes() {
435         final Set<String> subtypes = MdnsUtils.newSet();
436         for (int i = 0; i < listeners.size(); i++) {
437             final MdnsSearchOptions listenerOptions = listeners.valueAt(i).searchOptions;
438             subtypes.addAll(listenerOptions.getSubtypes());
439         }
440         return subtypes;
441     }
442 
443     /**
444      * Get the executor service.
445      */
getExecutor()446     public ScheduledExecutorService getExecutor() {
447         return executor;
448     }
449 
removeScheduledTask()450     private void removeScheduledTask() {
451         dependencies.removeMessages(handler, EVENT_START_QUERYTASK);
452         sharedLog.log("Remove EVENT_START_QUERYTASK"
453                 + ", current session: " + currentSessionId);
454         ++currentSessionId;
455     }
456 
responseMatchesOptions(@onNull MdnsResponse response, @NonNull MdnsSearchOptions options)457     private boolean responseMatchesOptions(@NonNull MdnsResponse response,
458             @NonNull MdnsSearchOptions options) {
459         final boolean matchesInstanceName = options.getResolveInstanceName() == null
460                 // DNS is case-insensitive, so ignore case in the comparison
461                 || MdnsUtils.equalsIgnoreDnsCase(options.getResolveInstanceName(),
462                 response.getServiceInstanceName());
463 
464         // If discovery is requiring some subtypes, the response must have one that matches a
465         // requested one.
466         final List<String> responseSubtypes = response.getSubtypes() == null
467                 ? Collections.emptyList() : response.getSubtypes();
468         final boolean matchesSubtype = options.getSubtypes().size() == 0
469                 || CollectionUtils.any(options.getSubtypes(), requiredSub ->
470                 CollectionUtils.any(responseSubtypes, actualSub ->
471                         MdnsUtils.equalsIgnoreDnsCase(
472                                 MdnsConstants.SUBTYPE_PREFIX + requiredSub, actualSub)));
473 
474         return matchesInstanceName && matchesSubtype;
475     }
476 
477     /**
478      * Unregisters {@code listener} from receiving discovery event of mDNS service instances.
479      *
480      * @param listener The {@link MdnsServiceBrowserListener} to unregister.
481      * @return {@code true} if no listener is registered with this client after unregistering {@code
482      * listener}. Otherwise returns {@code false}.
483      */
stopSendAndReceive(@onNull MdnsServiceBrowserListener listener)484     public boolean stopSendAndReceive(@NonNull MdnsServiceBrowserListener listener) {
485         ensureRunningOnHandlerThread(handler);
486         if (listeners.remove(listener) == null) {
487             return listeners.isEmpty();
488         }
489         if (listeners.isEmpty()) {
490             shutDown();
491         }
492         return listeners.isEmpty();
493     }
494 
495     /**
496      * Process an incoming response packet.
497      */
processResponse(@onNull MdnsPacket packet, @NonNull SocketKey socketKey)498     public synchronized void processResponse(@NonNull MdnsPacket packet,
499             @NonNull SocketKey socketKey) {
500         ensureRunningOnHandlerThread(handler);
501         // Augment the list of current known responses, and generated responses for resolve
502         // requests if there is no known response
503         final List<MdnsResponse> cachedList = serviceCache.getCachedServices(cacheKey);
504         final List<MdnsResponse> currentList = new ArrayList<>(cachedList);
505         List<MdnsResponse> additionalResponses = makeResponsesForResolve(socketKey);
506         for (MdnsResponse additionalResponse : additionalResponses) {
507             if (findMatchedResponse(
508                     cachedList, additionalResponse.getServiceInstanceName()) == null) {
509                 currentList.add(additionalResponse);
510             }
511         }
512         final Pair<Set<MdnsResponse>, ArrayList<MdnsResponse>> augmentedResult =
513                 responseDecoder.augmentResponses(packet, currentList,
514                         socketKey.getInterfaceIndex(), socketKey.getNetwork());
515 
516         final Set<MdnsResponse> modifiedResponse = augmentedResult.first;
517         final ArrayList<MdnsResponse> allResponses = augmentedResult.second;
518 
519         for (MdnsResponse response : allResponses) {
520             final String serviceInstanceName = response.getServiceInstanceName();
521             if (modifiedResponse.contains(response)) {
522                 if (response.isGoodbye()) {
523                     onGoodbyeReceived(serviceInstanceName);
524                 } else {
525                     onResponseModified(response);
526                 }
527             } else if (findMatchedResponse(cachedList, serviceInstanceName) != null) {
528                 // If the response is not modified and already in the cache. The cache will
529                 // need to be updated to refresh the last receipt time.
530                 serviceCache.addOrUpdateService(cacheKey, response);
531             }
532         }
533         if (dependencies.hasMessages(handler, EVENT_START_QUERYTASK)) {
534             final long now = clock.elapsedRealtime();
535             final long minRemainingTtl = getMinRemainingTtl(now);
536             MdnsQueryScheduler.ScheduledQueryTaskArgs args =
537                     mdnsQueryScheduler.maybeRescheduleCurrentRun(now, minRemainingTtl,
538                             lastSentTime, currentSessionId + 1);
539             if (args != null) {
540                 removeScheduledTask();
541                 dependencies.sendMessageDelayed(
542                         handler,
543                         handler.obtainMessage(EVENT_START_QUERYTASK, args),
544                         calculateTimeToNextTask(args, now, sharedLog));
545             }
546         }
547     }
548 
onFailedToParseMdnsResponse(int receivedPacketNumber, int errorCode)549     public synchronized void onFailedToParseMdnsResponse(int receivedPacketNumber, int errorCode) {
550         ensureRunningOnHandlerThread(handler);
551         for (int i = 0; i < listeners.size(); i++) {
552             listeners.keyAt(i).onFailedToParseMdnsResponse(receivedPacketNumber, errorCode);
553         }
554     }
555 
notifyRemovedServiceToListeners(@onNull MdnsResponse response, @NonNull String message)556     private void notifyRemovedServiceToListeners(@NonNull MdnsResponse response,
557             @NonNull String message) {
558         for (int i = 0; i < listeners.size(); i++) {
559             if (!responseMatchesOptions(response, listeners.valueAt(i).searchOptions)) continue;
560             final MdnsServiceBrowserListener listener = listeners.keyAt(i);
561             if (response.getServiceInstanceName() != null) {
562                 listeners.valueAt(i).unsetServiceDiscovered(response.getServiceInstanceName());
563                 final MdnsServiceInfo serviceInfo = buildMdnsServiceInfoFromResponse(
564                         response, serviceTypeLabels);
565                 if (response.isComplete()) {
566                     sharedLog.log(message + ". onServiceRemoved: " + serviceInfo);
567                     listener.onServiceRemoved(serviceInfo);
568                 }
569                 sharedLog.log(message + ". onServiceNameRemoved: " + serviceInfo);
570                 listener.onServiceNameRemoved(serviceInfo);
571             }
572         }
573     }
574 
575     /** Notify all services are removed because the socket is destroyed. */
notifySocketDestroyed()576     public void notifySocketDestroyed() {
577         ensureRunningOnHandlerThread(handler);
578         for (MdnsResponse response : serviceCache.getCachedServices(cacheKey)) {
579             final String name = response.getServiceInstanceName();
580             if (name == null) continue;
581             notifyRemovedServiceToListeners(response, "Socket destroyed");
582         }
583         shutDown();
584     }
585 
onResponseModified(@onNull MdnsResponse response)586     private void onResponseModified(@NonNull MdnsResponse response) {
587         final String serviceInstanceName = response.getServiceInstanceName();
588         final MdnsResponse currentResponse =
589                 serviceCache.getCachedService(serviceInstanceName, cacheKey);
590 
591         final boolean newInCache = currentResponse == null;
592         boolean serviceBecomesComplete = false;
593         if (newInCache) {
594             if (serviceInstanceName != null) {
595                 serviceCache.addOrUpdateService(cacheKey, response);
596             }
597         } else {
598             boolean before = currentResponse.isComplete();
599             serviceCache.addOrUpdateService(cacheKey, response);
600             boolean after = response.isComplete();
601             serviceBecomesComplete = !before && after;
602         }
603         sharedLog.i(String.format(
604                 "Handling response from service: %s, newInCache: %b, serviceBecomesComplete:"
605                         + " %b, responseIsComplete: %b",
606                 serviceInstanceName, newInCache, serviceBecomesComplete,
607                 response.isComplete()));
608         MdnsServiceInfo serviceInfo =
609                 buildMdnsServiceInfoFromResponse(response, serviceTypeLabels);
610 
611         for (int i = 0; i < listeners.size(); i++) {
612             // If a service stops matching the options (currently can only happen if it loses a
613             // subtype), service lost callbacks should also be sent; this is not done today as
614             // only expiration of SRV records is used, not PTR records used for subtypes, so
615             // services never lose PTR record subtypes.
616             if (!responseMatchesOptions(response, listeners.valueAt(i).searchOptions)) continue;
617             final MdnsServiceBrowserListener listener = listeners.keyAt(i);
618             final ListenerInfo listenerInfo = listeners.valueAt(i);
619             final boolean newServiceFound = listenerInfo.setServiceDiscovered(serviceInstanceName);
620             if (newServiceFound) {
621                 sharedLog.log("onServiceNameDiscovered: " + serviceInfo);
622                 listener.onServiceNameDiscovered(serviceInfo, false /* isServiceFromCache */);
623             }
624 
625             if (response.isComplete()) {
626                 if (newServiceFound || serviceBecomesComplete) {
627                     sharedLog.log("onServiceFound: " + serviceInfo);
628                     listener.onServiceFound(serviceInfo, false /* isServiceFromCache */);
629                 } else {
630                     sharedLog.log("onServiceUpdated: " + serviceInfo);
631                     listener.onServiceUpdated(serviceInfo);
632                 }
633             }
634         }
635     }
636 
onGoodbyeReceived(@ullable String serviceInstanceName)637     private void onGoodbyeReceived(@Nullable String serviceInstanceName) {
638         final MdnsResponse response =
639                 serviceCache.removeService(serviceInstanceName, cacheKey);
640         if (response == null) {
641             return;
642         }
643         notifyRemovedServiceToListeners(response, "Goodbye received");
644     }
645 
shouldRemoveServiceAfterTtlExpires()646     private boolean shouldRemoveServiceAfterTtlExpires() {
647         if (removeServiceAfterTtlExpires) {
648             return true;
649         }
650         return searchOptions != null && searchOptions.removeExpiredService();
651     }
652 
makeResponsesForResolve(@onNull SocketKey socketKey)653     private List<MdnsResponse> makeResponsesForResolve(@NonNull SocketKey socketKey) {
654         final List<MdnsResponse> resolveResponses = new ArrayList<>();
655         for (int i = 0; i < listeners.size(); i++) {
656             final String resolveName = listeners.valueAt(i).searchOptions.getResolveInstanceName();
657             if (resolveName == null) {
658                 continue;
659             }
660             if (CollectionUtils.any(resolveResponses,
661                     r -> MdnsUtils.equalsIgnoreDnsCase(resolveName, r.getServiceInstanceName()))) {
662                 continue;
663             }
664             MdnsResponse knownResponse =
665                     serviceCache.getCachedService(resolveName, cacheKey);
666             if (knownResponse == null) {
667                 final ArrayList<String> instanceFullName = new ArrayList<>(
668                         serviceTypeLabels.length + 1);
669                 instanceFullName.add(resolveName);
670                 instanceFullName.addAll(Arrays.asList(serviceTypeLabels));
671                 knownResponse = new MdnsResponse(
672                         0L /* lastUpdateTime */, instanceFullName.toArray(new String[0]),
673                         socketKey.getInterfaceIndex(), socketKey.getNetwork());
674             }
675             resolveResponses.add(knownResponse);
676         }
677         return resolveResponses;
678     }
679 
needSendDiscoveryQueries( @onNull ArrayMap<MdnsServiceBrowserListener, ListenerInfo> listeners)680     private static boolean needSendDiscoveryQueries(
681             @NonNull ArrayMap<MdnsServiceBrowserListener, ListenerInfo> listeners) {
682         // Note iterators are discouraged on ArrayMap as per its documentation
683         for (int i = 0; i < listeners.size(); i++) {
684             if (listeners.valueAt(i).searchOptions.getResolveInstanceName() == null) {
685                 return true;
686             }
687         }
688         return false;
689     }
690 
tryRemoveServiceAfterTtlExpires()691     private void tryRemoveServiceAfterTtlExpires() {
692         if (!shouldRemoveServiceAfterTtlExpires()) return;
693 
694         final Iterator<MdnsResponse> iter = serviceCache.getCachedServices(cacheKey).iterator();
695         while (iter.hasNext()) {
696             MdnsResponse existingResponse = iter.next();
697             if (existingResponse.hasServiceRecord()
698                     && existingResponse.getServiceRecord()
699                     .getRemainingTTL(clock.elapsedRealtime()) == 0) {
700                 serviceCache.removeService(existingResponse.getServiceInstanceName(), cacheKey);
701                 notifyRemovedServiceToListeners(existingResponse, "TTL expired");
702             }
703         }
704     }
705 
706     private static class QuerySentArguments {
707         private final int transactionId;
708         private final List<String> subTypes = new ArrayList<>();
709         private final MdnsQueryScheduler.ScheduledQueryTaskArgs taskArgs;
710 
QuerySentArguments(int transactionId, @NonNull List<String> subTypes, @NonNull MdnsQueryScheduler.ScheduledQueryTaskArgs taskArgs)711         QuerySentArguments(int transactionId, @NonNull List<String> subTypes,
712                 @NonNull MdnsQueryScheduler.ScheduledQueryTaskArgs taskArgs) {
713             this.transactionId = transactionId;
714             this.subTypes.addAll(subTypes);
715             this.taskArgs = taskArgs;
716         }
717     }
718 
719     // A FutureTask that enqueues a single query, and schedule a new FutureTask for the next task.
720     private class QueryTask implements Runnable {
721         private final MdnsQueryScheduler.ScheduledQueryTaskArgs taskArgs;
722         private final List<MdnsResponse> servicesToResolve = new ArrayList<>();
723         private final List<String> subtypes = new ArrayList<>();
724         private final boolean sendDiscoveryQueries;
725         private final List<MdnsResponse> existingServices = new ArrayList<>();
QueryTask(@onNull MdnsQueryScheduler.ScheduledQueryTaskArgs taskArgs, @NonNull Collection<MdnsResponse> servicesToResolve, @NonNull Collection<String> subtypes, boolean sendDiscoveryQueries, @NonNull Collection<MdnsResponse> existingServices)726         QueryTask(@NonNull MdnsQueryScheduler.ScheduledQueryTaskArgs taskArgs,
727                 @NonNull Collection<MdnsResponse> servicesToResolve,
728                 @NonNull Collection<String> subtypes, boolean sendDiscoveryQueries,
729                 @NonNull Collection<MdnsResponse> existingServices) {
730             this.taskArgs = taskArgs;
731             this.servicesToResolve.addAll(servicesToResolve);
732             this.subtypes.addAll(subtypes);
733             this.sendDiscoveryQueries = sendDiscoveryQueries;
734             this.existingServices.addAll(existingServices);
735         }
736 
737         @Override
run()738         public void run() {
739             Pair<Integer, List<String>> result;
740             try {
741                 result =
742                         new EnqueueMdnsQueryCallable(
743                                 socketClient,
744                                 serviceType,
745                                 subtypes,
746                                 taskArgs.config.expectUnicastResponse,
747                                 taskArgs.config.transactionId,
748                                 taskArgs.config.socketKey,
749                                 taskArgs.config.onlyUseIpv6OnIpv6OnlyNetworks,
750                                 sendDiscoveryQueries,
751                                 servicesToResolve,
752                                 clock,
753                                 sharedLog,
754                                 dependencies,
755                                 existingServices,
756                                 featureFlags.isQueryWithKnownAnswerEnabled())
757                                 .call();
758             } catch (RuntimeException e) {
759                 sharedLog.e(String.format("Failed to run EnqueueMdnsQueryCallable for subtype: %s",
760                         TextUtils.join(",", subtypes)), e);
761                 result = Pair.create(INVALID_TRANSACTION_ID, new ArrayList<>());
762             }
763             dependencies.sendMessage(
764                     handler, handler.obtainMessage(EVENT_QUERY_RESULT,
765                             new QuerySentArguments(result.first, result.second, taskArgs)));
766         }
767     }
768 
getMinRemainingTtl(long now)769     private long getMinRemainingTtl(long now) {
770         long minRemainingTtl = Long.MAX_VALUE;
771         for (MdnsResponse response : serviceCache.getCachedServices(cacheKey)) {
772             if (!response.isComplete()) {
773                 continue;
774             }
775             long remainingTtl =
776                     response.getServiceRecord().getRemainingTTL(now);
777             // remainingTtl is <= 0 means the service expired.
778             if (remainingTtl <= 0) {
779                 return 0;
780             }
781             if (remainingTtl < minRemainingTtl) {
782                 minRemainingTtl = remainingTtl;
783             }
784         }
785         return minRemainingTtl == Long.MAX_VALUE ? 0 : minRemainingTtl;
786     }
787 
calculateTimeToNextTask(MdnsQueryScheduler.ScheduledQueryTaskArgs args, long now, SharedLog sharedLog)788     private static long calculateTimeToNextTask(MdnsQueryScheduler.ScheduledQueryTaskArgs args,
789             long now, SharedLog sharedLog) {
790         long timeToNextTasksWithBackoffInMs = Math.max(args.timeToRun - now, 0);
791         sharedLog.log(String.format("Next run: sessionId: %d, in %d ms",
792                 args.sessionId, timeToNextTasksWithBackoffInMs));
793         return timeToNextTasksWithBackoffInMs;
794     }
795 
796     /**
797      * Dump ServiceTypeClient state.
798      */
dump(PrintWriter pw)799     public void dump(PrintWriter pw) {
800         ensureRunningOnHandlerThread(handler);
801         pw.println("ServiceTypeClient: Type{" + serviceType + "} " + socketKey + " with "
802                 + listeners.size() + " listeners.");
803     }
804 }