Merge "Implement the stop resolution with MdnsDiscoveryManager"
diff --git a/service-t/src/com/android/server/NsdService.java b/service-t/src/com/android/server/NsdService.java
index 8e06fde..54ec093 100644
--- a/service-t/src/com/android/server/NsdService.java
+++ b/service-t/src/com/android/server/NsdService.java
@@ -545,6 +545,12 @@
                 return new String(out.array(), 0, out.position(), utf8);
             }
 
+            private void stopDiscoveryManagerRequest(ClientRequest request, int clientId, int id,
+                    ClientInfo clientInfo) {
+                clientInfo.unregisterMdnsListenerFromRequest(request);
+                removeRequestMap(clientId, id, clientInfo);
+            }
+
             @Override
             public boolean processMessage(Message msg) {
                 final ClientInfo clientInfo;
@@ -631,11 +637,7 @@
                         // point, so this needs to check the type of the original request to
                         // unregister instead of looking at the flag value.
                         if (request instanceof DiscoveryManagerRequest) {
-                            final MdnsListener listener =
-                                    ((DiscoveryManagerRequest) request).mListener;
-                            mMdnsDiscoveryManager.unregisterListener(
-                                    listener.getListenedServiceType(), listener);
-                            removeRequestMap(clientId, id, clientInfo);
+                            stopDiscoveryManagerRequest(request, clientId, id, clientInfo);
                             clientInfo.onStopDiscoverySucceeded(clientId);
                         } else {
                             removeRequestMap(clientId, id, clientInfo);
@@ -803,15 +805,22 @@
                             break;
                         }
                         id = request.mGlobalId;
-                        removeRequestMap(clientId, id, clientInfo);
-                        if (stopResolveService(id)) {
+                        // Note isMdnsDiscoveryManagerEnabled may have changed to false at this
+                        // point, so this needs to check the type of the original request to
+                        // unregister instead of looking at the flag value.
+                        if (request instanceof DiscoveryManagerRequest) {
+                            stopDiscoveryManagerRequest(request, clientId, id, clientInfo);
                             clientInfo.onStopResolutionSucceeded(clientId);
                         } else {
-                            clientInfo.onStopResolutionFailed(
-                                    clientId, NsdManager.FAILURE_OPERATION_NOT_RUNNING);
+                            removeRequestMap(clientId, id, clientInfo);
+                            if (stopResolveService(id)) {
+                                clientInfo.onStopResolutionSucceeded(clientId);
+                            } else {
+                                clientInfo.onStopResolutionFailed(
+                                        clientId, NsdManager.FAILURE_OPERATION_NOT_RUNNING);
+                            }
+                            clientInfo.mResolvedService = null;
                         }
-                        clientInfo.mResolvedService = null;
-                        // TODO: Implement the stop resolution with MdnsDiscoveryManager.
                         break;
                     }
                     case NsdManager.REGISTER_SERVICE_CALLBACK:
@@ -1164,10 +1173,7 @@
                             Log.wtf(TAG, "non-DiscoveryManager request in DiscoveryManager event");
                             break;
                         }
-                        final MdnsListener listener = ((DiscoveryManagerRequest) request).mListener;
-                        mMdnsDiscoveryManager.unregisterListener(
-                                listener.getListenedServiceType(), listener);
-                        removeRequestMap(clientId, transactionId, clientInfo);
+                        stopDiscoveryManagerRequest(request, clientId, transactionId, clientInfo);
                         break;
                     }
                     default:
@@ -1694,6 +1700,13 @@
             mIsPreSClient = true;
         }
 
+        private void unregisterMdnsListenerFromRequest(ClientRequest request) {
+            final MdnsListener listener =
+                    ((DiscoveryManagerRequest) request).mListener;
+            mMdnsDiscoveryManager.unregisterListener(
+                    listener.getListenedServiceType(), listener);
+        }
+
         // Remove any pending requests from the global map when we get rid of a client,
         // and send cancellations to the daemon.
         private void expungeAllRequests() {
@@ -1709,10 +1722,7 @@
                 }
 
                 if (request instanceof DiscoveryManagerRequest) {
-                    final MdnsListener listener =
-                            ((DiscoveryManagerRequest) request).mListener;
-                    mMdnsDiscoveryManager.unregisterListener(
-                            listener.getListenedServiceType(), listener);
+                    unregisterMdnsListenerFromRequest(request);
                     continue;
                 }
 
diff --git a/tests/unit/java/com/android/server/NsdServiceTest.java b/tests/unit/java/com/android/server/NsdServiceTest.java
index 8fc9252..08221ec 100644
--- a/tests/unit/java/com/android/server/NsdServiceTest.java
+++ b/tests/unit/java/com/android/server/NsdServiceTest.java
@@ -1150,6 +1150,37 @@
                 argThat(info -> matches(info, new NsdServiceInfo(regInfo.getServiceName(), null))));
     }
 
+    @Test
+    public void testStopServiceResolutionWithMdnsDiscoveryManager() {
+        setMdnsDiscoveryManagerEnabled();
+
+        final NsdManager client = connectClient(mService);
+        final ResolveListener resolveListener = mock(ResolveListener.class);
+        final Network network = new Network(999);
+        final String serviceType = "_nsd._service._tcp";
+        final String constructedServiceType = "_nsd._sub._service._tcp.local";
+        final ArgumentCaptor<MdnsServiceBrowserListener> listenerCaptor =
+                ArgumentCaptor.forClass(MdnsServiceBrowserListener.class);
+        final NsdServiceInfo request = new NsdServiceInfo(SERVICE_NAME, serviceType);
+        request.setNetwork(network);
+        client.resolveService(request, resolveListener);
+        waitForIdle();
+        verify(mSocketProvider).startMonitoringSockets();
+        verify(mDiscoveryManager).registerListener(eq(constructedServiceType),
+                listenerCaptor.capture(), argThat(options -> network.equals(options.getNetwork())));
+
+        client.stopServiceResolution(resolveListener);
+        waitForIdle();
+
+        // Verify the listener has been unregistered.
+        verify(mDiscoveryManager, timeout(TIMEOUT_MS))
+                .unregisterListener(eq(constructedServiceType), eq(listenerCaptor.getValue()));
+        verify(resolveListener, timeout(TIMEOUT_MS)).onResolutionStopped(argThat(ns ->
+                request.getServiceName().equals(ns.getServiceName())
+                        && request.getServiceType().equals(ns.getServiceType())));
+        verify(mSocketProvider, timeout(CLEANUP_DELAY_MS + TIMEOUT_MS)).stopMonitoringSockets();
+    }
+
     private void waitForIdle() {
         HandlerUtils.waitForIdle(mHandler, TIMEOUT_MS);
     }