Fix unregistering listener issue

Unregistering listener on MDnsEventReporter would fail because
the bringing listener is a pointer which is not existed in the
set. So compare the binder which get from listener to find out
the real listener to be removed.

Refacotr the code to store the death recipient in the set
instead. Because we need to unlink the death recipient when
unregistering listener but we do not save the death recipient
before.

Bug: 234107254
Test: atest CtsNetTestCases
      cd system/netd & atest
Change-Id: Ic3bd793ba1d297188700b18b39f380cb7cc43f41
(cherry picked from commit e5ead45a046a920466a807630ccbc382eedc9b28)
Merged-In: Ic3bd793ba1d297188700b18b39f380cb7cc43f41

3 files changed, 44 insertions, 37 deletions

diff --git a/server/MDnsEventReporter.cpp b/server/MDnsEventReporter.cpp
index db9021ae..e94de367 100644
--- a/server/MDnsEventReporter.cpp
+++ b/server/MDnsEventReporter.cpp
@@ -18,6 +18,8 @@
 
 #include "MDnsEventReporter.h"
 
+using android::IInterface;
+using android::sp;
 using android::net::mdns::aidl::IMDnsEventListener;
 
 MDnsEventReporter& MDnsEventReporter::getInstance() {
@@ -30,11 +32,11 @@ const MDnsEventReporter::EventListenerSet& MDnsEventReporter::getEventListeners(
     return getEventListenersImpl();
 }
 
-int MDnsEventReporter::addEventListener(const android::sp<IMDnsEventListener>& listener) {
+int MDnsEventReporter::addEventListener(const sp<IMDnsEventListener>& listener) {
     return addEventListenerImpl(listener);
 }
 
-int MDnsEventReporter::removeEventListener(const android::sp<IMDnsEventListener>& listener) {
+int MDnsEventReporter::removeEventListener(const sp<IMDnsEventListener>& listener) {
     return removeEventListenerImpl(listener);
 }
 
@@ -43,7 +45,7 @@ const MDnsEventReporter::EventListenerSet& MDnsEventReporter::getEventListenersI
     return mEventListeners;
 }
 
-int MDnsEventReporter::addEventListenerImpl(const android::sp<IMDnsEventListener>& listener) {
+int MDnsEventReporter::addEventListenerImpl(const sp<IMDnsEventListener>& listener) {
     if (listener == nullptr) {
         ALOGE("The event listener should not be null");
         return -EINVAL;
@@ -52,39 +54,19 @@ int MDnsEventReporter::addEventListenerImpl(const android::sp<IMDnsEventListener
     std::lock_guard lock(mMutex);
 
     for (const auto& it : mEventListeners) {
-        if (android::IInterface::asBinder(it).get() ==
-            android::IInterface::asBinder(listener).get()) {
+        if (IInterface::asBinder(it->getListener()).get() == IInterface::asBinder(listener).get()) {
             ALOGW("The event listener was already subscribed");
             return -EEXIST;
         }
     }
 
-    // Create the death listener.
-    class DeathRecipient : public android::IBinder::DeathRecipient {
-      public:
-        DeathRecipient(MDnsEventReporter* eventReporter,
-                       const android::sp<IMDnsEventListener>& listener)
-            : mEventReporter(eventReporter), mListener(listener) {}
-        ~DeathRecipient() override = default;
-        void binderDied(const android::wp<android::IBinder>& /* who */) override {
-            mEventReporter->removeEventListenerImpl(mListener);
-        }
-
-      private:
-        MDnsEventReporter* mEventReporter;
-        android::sp<IMDnsEventListener> mListener;
-    };
-
-    android::sp<android::IBinder::DeathRecipient> deathRecipient =
-            new DeathRecipient(this, listener);
-
-    android::IInterface::asBinder(listener)->linkToDeath(deathRecipient);
-
-    mEventListeners.insert(listener);
+    auto eventListener = sp<EventListener>::make(this, listener);
+    IInterface::asBinder(listener)->linkToDeath(eventListener);
+    mEventListeners.insert(eventListener);
     return 0;
 }
 
-int MDnsEventReporter::removeEventListenerImpl(const android::sp<IMDnsEventListener>& listener) {
+int MDnsEventReporter::removeEventListenerImpl(const sp<IMDnsEventListener>& listener) {
     if (listener == nullptr) {
         ALOGE("The event listener should not be null");
         return -EINVAL;
@@ -92,6 +74,14 @@ int MDnsEventReporter::removeEventListenerImpl(const android::sp<IMDnsEventListe
 
     std::lock_guard lock(mMutex);
 
-    mEventListeners.erase(listener);
-    return 0;
-}
+    for (const auto& it : mEventListeners) {
+        const auto& binder = IInterface::asBinder(it->getListener());
+        if (binder == IInterface::asBinder(listener)) {
+            binder->unlinkToDeath(it);
+            mEventListeners.erase(it);
+            return 0;
+        }
+    }
+    ALOGE("The event listener does not exist");
+    return -ENOENT;
+}
\ No newline at end of file
diff --git a/server/MDnsEventReporter.h b/server/MDnsEventReporter.h
index e49c3e3f..cbc43ecb 100644
--- a/server/MDnsEventReporter.h
+++ b/server/MDnsEventReporter.h
@@ -23,10 +23,28 @@
 
 class MDnsEventReporter final {
   public:
+    class EventListener : public android::IBinder::DeathRecipient {
+      public:
+        EventListener(MDnsEventReporter* eventReporter,
+                      const android::sp<android::net::mdns::aidl::IMDnsEventListener>& listener)
+            : mEventReporter(eventReporter), mListener(listener) {}
+        ~EventListener() override = default;
+        void binderDied(const android::wp<android::IBinder>& /* who */) override {
+            mEventReporter->removeEventListenerImpl(mListener);
+        }
+        android::sp<android::net::mdns::aidl::IMDnsEventListener> getListener() {
+            return mListener;
+        }
+
+      private:
+        MDnsEventReporter* mEventReporter;
+        android::sp<android::net::mdns::aidl::IMDnsEventListener> mListener;
+    };
+
     MDnsEventReporter(const MDnsEventReporter&) = delete;
     MDnsEventReporter& operator=(const MDnsEventReporter&) = delete;
 
-    using EventListenerSet = std::set<android::sp<android::net::mdns::aidl::IMDnsEventListener>>;
+    using EventListenerSet = std::set<android::sp<EventListener>>;
 
     // Get the instance of the singleton MDnsEventReporter.
     static MDnsEventReporter& getInstance();
@@ -56,5 +74,4 @@ class MDnsEventReporter final {
             const android::sp<android::net::mdns::aidl::IMDnsEventListener>& listener)
             EXCLUDES(mMutex);
     const EventListenerSet& getEventListenersImpl() const EXCLUDES(mMutex);
-    void handleEventBinderDied(const void* who) EXCLUDES(mMutex);
 };
diff --git a/server/MDnsSdListener.cpp b/server/MDnsSdListener.cpp
index 16364008..1d1ea40a 100644
--- a/server/MDnsSdListener.cpp
+++ b/server/MDnsSdListener.cpp
@@ -140,7 +140,7 @@ void MDnsSdListenerDiscoverCallback(DNSServiceRef /* sdRef */, DNSServiceFlags f
     }
 
     for (const auto& it : listeners) {
-        it->onServiceDiscoveryStatus(info);
+        it->getListener()->onServiceDiscoveryStatus(info);
     }
 }
 
@@ -212,7 +212,7 @@ void MDnsSdListenerRegisterCallback(DNSServiceRef /* sdRef */, DNSServiceFlags /
     }
 
     for (const auto& it : listeners) {
-        it->onServiceRegistrationStatus(info);
+        it->getListener()->onServiceRegistrationStatus(info);
     }
 }
 
@@ -277,7 +277,7 @@ void MDnsSdListenerResolveCallback(DNSServiceRef /* sdRef */, DNSServiceFlags /*
     }
 
     for (const auto& it : listeners) {
-        it->onServiceResolutionStatus(info);
+        it->getListener()->onServiceResolutionStatus(info);
     }
 }
 
@@ -343,7 +343,7 @@ void MDnsSdListenerGetAddrInfoCallback(DNSServiceRef /* sdRef */, DNSServiceFlag
         info.result = IMDnsEventListener::SERVICE_GET_ADDR_FAILED;
     }
     for (const auto& it : listeners) {
-        it->onGettingServiceAddressStatus(info);
+        it->getListener()->onGettingServiceAddressStatus(info);
     }
 }