[RTT] Add cleanup on app death

Add a binder death listener to register for app death and execute
cleanup code.

This is a replacement for the death listener added to the AsyncChannel
which caused a memory leak in apps.

Bug: 37443149
Test: verified using test app ResponderToolActivity
Change-Id: I90b28fcb2d7d61cf40807960c3bf56a262b595b9

2 files changed, 62 insertions, 2 deletions

diff --git a/service/java/com/android/server/wifi/RttService.java b/service/java/com/android/server/wifi/RttService.java
index 7e4648b10..89de870bb 100644
--- a/service/java/com/android/server/wifi/RttService.java
+++ b/service/java/com/android/server/wifi/RttService.java
@@ -26,6 +26,7 @@ import android.os.RemoteException;
 import android.util.ArrayMap;
 import android.util.Log;
 import android.util.Slog;
+import android.util.SparseArray;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.util.AsyncChannel;
@@ -50,12 +51,56 @@ public final class RttService extends SystemService {
     private static final String WIFICOND_SERVICE_NAME = "wificond";
 
     static class RttServiceImpl extends IRttManager.Stub {
+        private int mCurrentKey = 100; // increment on each usage
+        private final SparseArray<IBinder> mBinderByKey = new SparseArray<>();
 
         @Override
-        public Messenger getMessenger() {
+        public Messenger getMessenger(IBinder binder, int[] key) {
+            if (key != null && key.length != 0) {
+                final int keyToUse = mCurrentKey++;
+                if (binder != null) {
+                    try {
+                        binder.linkToDeath(() -> {
+                            // clean-up here if didn't get final registration
+                            Slog.d(TAG, "Binder death on key=" + keyToUse);
+                            mBinderByKey.delete(keyToUse);
+                        }, 0);
+                        mBinderByKey.put(keyToUse, binder);
+                    } catch (RemoteException e) {
+                        Slog.e(TAG, "getMessenger: can't link to death on binder: " + e);
+                        return null;
+                    }
+                }
+
+                key[0] = keyToUse;
+            }
             return new Messenger(mClientHandler);
         }
 
+        private class RttDeathListener implements IBinder.DeathRecipient {
+            private final IBinder mBinder;
+            private final Messenger mReplyTo;
+
+            RttDeathListener(IBinder binder, Messenger replyTo) {
+                mBinder = binder;
+                mReplyTo = replyTo;
+            }
+
+            @Override
+            public void binderDied() {
+                if (DBG) Slog.d(TAG, "binder death for client mReplyTo=" + mReplyTo);
+                synchronized (mLock) {
+                    ClientInfo ci = mClients.remove(mReplyTo);
+                    if (ci != null) {
+                        ci.cleanup();
+                    } else {
+                        Slog.w(TAG,
+                                "ClientInfo not found for terminated app -- mReplyTo=" + mReplyTo);
+                    }
+                }
+            }
+        }
+
         private class ClientHandler extends Handler {
 
             ClientHandler(android.os.Looper looper) {
@@ -93,6 +138,21 @@ public final class RttService extends SystemService {
                         ac.replyToMessage(msg, AsyncChannel.CMD_CHANNEL_FULLY_CONNECTED,
                                 AsyncChannel.STATUS_SUCCESSFUL);
                         return;
+                    case RttManager.CMD_OP_REG_BINDER: {
+                        int key = msg.arg1;
+                        IBinder binder = mBinderByKey.get(key);
+                        if (binder == null) {
+                            Slog.e(TAG, "Can't find binder registered with key=" + key + " - no "
+                                    + "death listener!");
+                            return;
+                        }
+                        try {
+                            binder.linkToDeath(new RttDeathListener(binder, msg.replyTo), 0);
+                        } catch (RemoteException e) {
+                            Slog.e(TAG, "Can't link to death for binder on key=" + key);
+                        }
+                        return;
+                    }
                 }
 
                 ClientInfo ci;
diff --git a/service/java/com/android/server/wifi/aware/WifiAwareRttStateManager.java b/service/java/com/android/server/wifi/aware/WifiAwareRttStateManager.java
index 74f34b362..afc044c1b 100644
--- a/service/java/com/android/server/wifi/aware/WifiAwareRttStateManager.java
+++ b/service/java/com/android/server/wifi/aware/WifiAwareRttStateManager.java
@@ -74,7 +74,7 @@ public class WifiAwareRttStateManager {
     public void startWithRttService(Context context, Looper looper, IRttManager service) {
         Messenger messenger;
         try {
-            messenger = service.getMessenger();
+            messenger = service.getMessenger(null, new int[1]);
         } catch (RemoteException e) {
             Log.e(TAG, "start(): not able to getMessenger() of WIFI_RTT_SERVICE");
             return;