Do not check active socket from other callbacks

Not all SocketTypeClients are cleaned up when their corresponding
sockets are closed. This occurs because the active socket check
for each callback examines the socket from all callbacks, even
those scheduled for removal in later notifications. To fix this,
check the active socket from the current callback only.

Fix: 317461509
Test: atest FrameworksNetTests NsdManagerTest
      Manually test with multiple discovery requests and
      Wifi on/off.
Change-Id: I2f81bb8e75f2c038d5e9dc025b208dcc6b441112
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java b/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java
index e7b0eaa..869ac9b 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java
@@ -96,7 +96,8 @@
         @Override
         public void onInterfaceDestroyed(@NonNull SocketKey socketKey,
                 @NonNull MdnsInterfaceSocket socket) {
-            notifySocketDestroyed(socketKey);
+            mActiveSockets.remove(socketKey);
+            mSocketCreationCallback.onSocketDestroyed(socketKey);
             maybeCleanupPacketHandler(socketKey);
         }
 
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java
index ad30ce0..9474464 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java
@@ -101,11 +101,17 @@
 
     private SocketCallback expectSocketCallback(MdnsServiceBrowserListener listener,
             Network requestedNetwork) {
+        return expectSocketCallback(listener, requestedNetwork, mSocketCreationCallback,
+                1 /* requestSocketCount */);
+    }
+
+    private SocketCallback expectSocketCallback(MdnsServiceBrowserListener listener,
+                Network requestedNetwork, SocketCreationCallback callback, int requestSocketCount) {
         final ArgumentCaptor<SocketCallback> callbackCaptor =
                 ArgumentCaptor.forClass(SocketCallback.class);
         mHandler.post(() -> mSocketClient.notifyNetworkRequested(
-                listener, requestedNetwork, mSocketCreationCallback));
-        verify(mProvider, timeout(DEFAULT_TIMEOUT))
+                listener, requestedNetwork, callback));
+        verify(mProvider, timeout(DEFAULT_TIMEOUT).times(requestSocketCount))
                 .requestSocket(eq(requestedNetwork), callbackCaptor.capture());
         return callbackCaptor.getValue();
     }
@@ -365,4 +371,40 @@
         callback.onInterfaceDestroyed(otherSocketKey, otherSocket);
         verify(mSocketCreationCallback).onSocketDestroyed(otherSocketKey);
     }
+
+    @Test
+    public void testSocketDestroyed_MultipleCallbacks() {
+        final MdnsInterfaceSocket socket2 = mock(MdnsInterfaceSocket.class);
+        final SocketKey socketKey2 = new SocketKey(1001 /* interfaceIndex */);
+        final SocketCreationCallback creationCallback1 = mock(SocketCreationCallback.class);
+        final SocketCreationCallback creationCallback2 = mock(SocketCreationCallback.class);
+        final SocketCreationCallback creationCallback3 = mock(SocketCreationCallback.class);
+        final SocketCallback callback1 = expectSocketCallback(
+                mock(MdnsServiceBrowserListener.class), mNetwork, creationCallback1,
+                1 /* requestSocketCount */);
+        final SocketCallback callback2 = expectSocketCallback(
+                mock(MdnsServiceBrowserListener.class), mNetwork, creationCallback2,
+                2 /* requestSocketCount */);
+        final SocketCallback callback3 = expectSocketCallback(
+                mock(MdnsServiceBrowserListener.class), null /* requestedNetwork */,
+                creationCallback3, 1 /* requestSocketCount */);
+
+        doReturn(createEmptyNetworkInterface()).when(mSocket).getInterface();
+        callback1.onSocketCreated(mSocketKey, mSocket, List.of());
+        callback2.onSocketCreated(mSocketKey, mSocket, List.of());
+        callback3.onSocketCreated(mSocketKey, mSocket, List.of());
+        callback3.onSocketCreated(socketKey2, socket2, List.of());
+        verify(creationCallback1).onSocketCreated(mSocketKey);
+        verify(creationCallback2).onSocketCreated(mSocketKey);
+        verify(creationCallback3).onSocketCreated(mSocketKey);
+        verify(creationCallback3).onSocketCreated(socketKey2);
+
+        callback1.onInterfaceDestroyed(mSocketKey, mSocket);
+        callback2.onInterfaceDestroyed(mSocketKey, mSocket);
+        callback3.onInterfaceDestroyed(mSocketKey, mSocket);
+        verify(creationCallback1).onSocketDestroyed(mSocketKey);
+        verify(creationCallback2).onSocketDestroyed(mSocketKey);
+        verify(creationCallback3).onSocketDestroyed(mSocketKey);
+        verify(creationCallback3, never()).onSocketDestroyed(socketKey2);
+    }
 }