From 3730a1bcf27cf4020305a0e31fc1eeb509d250ac Mon Sep 17 00:00:00 2001
From: Rikka <rikkanyaaa@gmail.com>
Date: Sun, 21 Aug 2022 17:04:44 +0000
Subject: [PATCH] Fix
 com.android.server.wm.TaskFpsCallbackController#unregisterListener method
 NEVER works

In com.android.server.wm.TaskFpsCallbackController class, ITaskFpsCallback (which extends IInterface) is used as the key for the HashMap. As ITaskFpsCallback instance is created every time, the unregisterListener method will NEVER work.

The correct usage is to use ITaskFpsCallback#asBinder as the key, just like what RemoteCallbackList do.

Change-Id: I16321d6784402105aea90a816f9ab23bd72e4635
---
 .../server/wm/TaskFpsCallbackController.java  | 32 ++++++++++++-------
 1 file changed, 21 insertions(+), 11 deletions(-)

diff --git a/services/core/java/com/android/server/wm/TaskFpsCallbackController.java b/services/core/java/com/android/server/wm/TaskFpsCallbackController.java
index c099628438900..8c798759c890e 100644
--- a/services/core/java/com/android/server/wm/TaskFpsCallbackController.java
+++ b/services/core/java/com/android/server/wm/TaskFpsCallbackController.java
@@ -26,8 +26,8 @@ import java.util.HashMap;
 final class TaskFpsCallbackController {
 
     private final Context mContext;
-    private final HashMap<ITaskFpsCallback, Long> mTaskFpsCallbacks;
-    private final HashMap<ITaskFpsCallback, IBinder.DeathRecipient> mDeathRecipients;
+    private final HashMap<IBinder, Long> mTaskFpsCallbacks;
+    private final HashMap<IBinder, IBinder.DeathRecipient> mDeathRecipients;
 
     TaskFpsCallbackController(Context context) {
         mContext = context;
@@ -36,32 +36,42 @@ final class TaskFpsCallbackController {
     }
 
     void registerListener(int taskId, ITaskFpsCallback callback) {
-        if (mTaskFpsCallbacks.containsKey(callback)) {
+        if (callback == null) {
+            return;
+        }
+
+        IBinder binder = callback.asBinder();
+        if (mTaskFpsCallbacks.containsKey(binder)) {
             return;
         }
 
         final long nativeListener = nativeRegister(callback, taskId);
-        mTaskFpsCallbacks.put(callback, nativeListener);
+        mTaskFpsCallbacks.put(binder, nativeListener);
 
         final IBinder.DeathRecipient deathRecipient = () -> unregisterListener(callback);
         try {
-            callback.asBinder().linkToDeath(deathRecipient, 0);
-            mDeathRecipients.put(callback, deathRecipient);
+            binder.linkToDeath(deathRecipient, 0);
+            mDeathRecipients.put(binder, deathRecipient);
         } catch (RemoteException e) {
             // ignore
         }
     }
 
     void unregisterListener(ITaskFpsCallback callback) {
-        if (!mTaskFpsCallbacks.containsKey(callback)) {
+        if (callback == null) {
             return;
         }
 
-        callback.asBinder().unlinkToDeath(mDeathRecipients.get(callback), 0);
-        mDeathRecipients.remove(callback);
+        IBinder binder = callback.asBinder();
+        if (!mTaskFpsCallbacks.containsKey(binder)) {
+            return;
+        }
 
-        nativeUnregister(mTaskFpsCallbacks.get(callback));
-        mTaskFpsCallbacks.remove(callback);
+        binder.unlinkToDeath(mDeathRecipients.get(binder), 0);
+        mDeathRecipients.remove(binder);
+
+        nativeUnregister(mTaskFpsCallbacks.get(binder));
+        mTaskFpsCallbacks.remove(binder);
     }
 
     private static native long nativeRegister(ITaskFpsCallback callback, int taskId);