From df82beec723b409b9c269d8ca6656a6faec14cab Mon Sep 17 00:00:00 2001
From: Remi NGUYEN VAN <reminv@google.com>
Date: Fri, 4 Apr 2025 13:11:08 +0900
Subject: [PATCH] Ensure PacketHandler is removed from sockets

When MdnsMultinetworkSocketClient stops using a socket (when
isSocketActive starts returning false), the packet handler should be
removed from that socket. The code only removed it from the
mSocketPacketHandlers map without clearing it from the socket.

This caused packets to be handled potentially many times, affecting
performance.

Bug: 406623719
Bug: 408324621
Test: atest
Flag: EXEMPT bugfix
(cherry picked from https://googleplex-android-review.googlesource.com/q/commit:b5fbc0143af3409c5c8c660c3ecb34d4e964efdb)
Merged-In: I3d1117a4b847f3b95a01fc8a5997a53655ef608a
Change-Id: I3d1117a4b847f3b95a01fc8a5997a53655ef608a
---
 .../mdns/MdnsMultinetworkSocketClient.java          | 13 +++++++++----
 .../mdns/MdnsMultinetworkSocketClientTest.java      |  5 +++++
 2 files changed, 14 insertions(+), 4 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 36fad31746..4c760d6d18 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java
@@ -101,7 +101,7 @@ public class MdnsMultinetworkSocketClient implements MdnsSocketClientBase {
                 @NonNull MdnsInterfaceSocket socket) {
             mActiveSockets.remove(socketKey);
             mSocketCreationCallback.onSocketDestroyed(socketKey);
-            maybeCleanupPacketHandler(socketKey);
+            maybeCleanupPacketHandler(socketKey, socket);
         }
 
         private void notifySocketDestroyed(@NonNull SocketKey socketKey) {
@@ -115,8 +115,9 @@ public class MdnsMultinetworkSocketClient implements MdnsSocketClientBase {
             for (int i = mActiveSockets.size() - 1; i >= 0; i--) {
                 // Iterate from the end so the socket can be removed
                 final SocketKey socketKey = mActiveSockets.keyAt(i);
+                final MdnsInterfaceSocket socket = mActiveSockets.valueAt(i);
                 notifySocketDestroyed(socketKey);
-                maybeCleanupPacketHandler(socketKey);
+                maybeCleanupPacketHandler(socketKey, socket);
             }
         }
     }
@@ -143,9 +144,13 @@ public class MdnsMultinetworkSocketClient implements MdnsSocketClientBase {
         return null;
     }
 
-    private void maybeCleanupPacketHandler(@NonNull SocketKey socketKey) {
+    private void maybeCleanupPacketHandler(@NonNull SocketKey socketKey,
+            @NonNull MdnsInterfaceSocket socket) {
         if (isSocketActive(socketKey)) return;
-        mSocketPacketHandlers.remove(socketKey);
+        final ReadPacketHandler handler = mSocketPacketHandlers.remove(socketKey);
+        if (handler != null) {
+            socket.removePacketHandler(handler);
+        }
     }
 
     private class ReadPacketHandler implements MulticastPacketReader.PacketHandler {
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 4c71991474..e115d5a882 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java
@@ -357,6 +357,8 @@ public class MdnsMultinetworkSocketClientTest {
         verify(mProvider).unrequestSocket(callback);
         verify(mSocketCreationCallback).onSocketDestroyed(mSocketKey);
         verify(mSocketCreationCallback).onSocketDestroyed(otherSocketKey);
+        verify(mSocket).removePacketHandler(any());
+        verify(otherSocket).removePacketHandler(any());
     }
 
     @Test
@@ -375,8 +377,10 @@ public class MdnsMultinetworkSocketClientTest {
         // Notify socket destroyed
         callback.onInterfaceDestroyed(mSocketKey, mSocket);
         verify(mSocketCreationCallback).onSocketDestroyed(mSocketKey);
+        verify(mSocket).removePacketHandler(any());
         callback.onInterfaceDestroyed(otherSocketKey, otherSocket);
         verify(mSocketCreationCallback).onSocketDestroyed(otherSocketKey);
+        verify(otherSocket).removePacketHandler(any());
     }
 
     @Test
@@ -413,6 +417,7 @@ public class MdnsMultinetworkSocketClientTest {
         verify(creationCallback2).onSocketDestroyed(mSocketKey);
         verify(creationCallback3).onSocketDestroyed(mSocketKey);
         verify(creationCallback3, never()).onSocketDestroyed(socketKey2);
+        verify(mSocket).removePacketHandler(any());
     }
 
     @Test