From 01d2f643be11c95618b72f31d89ca59564515567 Mon Sep 17 00:00:00 2001
From: Paul Hu <paulhu@google.com>
Date: Fri, 22 Dec 2023 07:24:24 +0000
Subject: [PATCH] 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
---
 .../mdns/MdnsMultinetworkSocketClient.java    |  3 +-
 .../MdnsMultinetworkSocketClientTest.java     | 46 ++++++++++++++++++-
 2 files changed, 46 insertions(+), 3 deletions(-)

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 e7b0eaa8a5..869ac9b11e 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 @@ public class MdnsMultinetworkSocketClient implements MdnsSocketClientBase {
         @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 ad30ce09bf..9474464c8b 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 @@ public class MdnsMultinetworkSocketClientTest {
 
     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 @@ public class MdnsMultinetworkSocketClientTest {
         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);
+    }
 }