Ensure that binders are unlinked from death recipients

- Add a callback to assist data requester to notify owners that the request
  is completed so they can clean up accordingly

Bug: 67864419
Test: com.android.server.am.AssistDataRequesterTest
Change-Id: I775aaf78ac7d9632cd9ca0a4cf05f2a87211d5fe
diff --git a/services/core/java/com/android/server/am/AssistDataReceiverProxy.java b/services/core/java/com/android/server/am/AssistDataReceiverProxy.java
index 8306731..f22fe37 100644
--- a/services/core/java/com/android/server/am/AssistDataReceiverProxy.java
+++ b/services/core/java/com/android/server/am/AssistDataReceiverProxy.java
@@ -37,17 +37,12 @@
     private static final String TAG = TAG_WITH_CLASS_NAME ? "AssistDataReceiverProxy" : TAG_AM;
 
     private String mCallerPackage;
-    private boolean mBinderDied;
     private IAssistDataReceiver mReceiver;
 
     public AssistDataReceiverProxy(IAssistDataReceiver receiver, String callerPackage) {
-        try {
-            receiver.asBinder().linkToDeath(this, 0);
-        } catch (RemoteException e) {
-            Log.w(TAG, "Could not link to client death", e);
-        }
         mReceiver = receiver;
         mCallerPackage = callerPackage;
+        linkToDeath();
     }
 
     @Override
@@ -58,7 +53,7 @@
 
     @Override
     public void onAssistDataReceivedLocked(Bundle data, int activityIndex, int activityCount) {
-        if (!mBinderDied) {
+        if (mReceiver != null) {
             try {
                 mReceiver.onHandleAssistData(data);
             } catch (RemoteException e) {
@@ -70,7 +65,7 @@
 
     @Override
     public void onAssistScreenshotReceivedLocked(Bitmap screenshot) {
-        if (!mBinderDied) {
+        if (mReceiver != null) {
             try {
                 mReceiver.onHandleAssistScreenshot(screenshot);
             } catch (RemoteException e) {
@@ -81,7 +76,27 @@
     }
 
     @Override
+    public void onAssistRequestCompleted() {
+        unlinkToDeath();
+    }
+
+    @Override
     public void binderDied() {
-        mBinderDied = true;
+        unlinkToDeath();
+    }
+
+    private void linkToDeath() {
+        try {
+            mReceiver.asBinder().linkToDeath(this, 0);
+        } catch (RemoteException e) {
+            Log.w(TAG, "Could not link to client death", e);
+        }
+    }
+
+    private void unlinkToDeath() {
+        if (mReceiver != null) {
+            mReceiver.asBinder().unlinkToDeath(this, 0);
+        }
+        mReceiver = null;
     }
 }
\ No newline at end of file
diff --git a/services/core/java/com/android/server/am/AssistDataRequester.java b/services/core/java/com/android/server/am/AssistDataRequester.java
index 0377cc3..a8f829f 100644
--- a/services/core/java/com/android/server/am/AssistDataRequester.java
+++ b/services/core/java/com/android/server/am/AssistDataRequester.java
@@ -94,6 +94,17 @@
          */
         @GuardedBy("mCallbacksLock")
         void onAssistScreenshotReceivedLocked(Bitmap screenshot);
+
+        /**
+         * Called when there is no more pending assist data or screenshots for the last request.
+         * If the request was canceled, then this callback will not be made. In addition, the
+         * callback will be made with the {@param mCallbacksLock} held, and only if
+         * {@link #canHandleReceivedAssistDataLocked()} is true.
+         */
+        @GuardedBy("mCallbacksLock")
+        default void onAssistRequestCompleted() {
+            // Do nothing
+        }
     }
 
     /**
@@ -139,12 +150,13 @@
     public void requestAssistData(List<IBinder> activityTokens, final boolean fetchData,
             final boolean fetchScreenshot, boolean allowFetchData, boolean allowFetchScreenshot,
             int callingUid, String callingPackage) {
-        // TODO: Better handle the cancel case if a request can be reused
-        // TODO: Known issue, if the assist data is not allowed on the current activity, then no
-        //       assist data is requested for any of the other activities
+        // TODO(b/34090158): Known issue, if the assist data is not allowed on the current activity,
+        //                   then no assist data is requested for any of the other activities
 
         // Early exit if there are no activity to fetch for
         if (activityTokens.isEmpty()) {
+            // No activities, just dispatch request-complete
+            tryDispatchRequestComplete();
             return;
         }
 
@@ -223,6 +235,9 @@
                 }
             }
         }
+        // For the cases where we dispatch null data/screenshot due to permissions, just dispatch
+        // request-complete after those are made
+        tryDispatchRequestComplete();
     }
 
     /**
@@ -230,6 +245,11 @@
      * data. The owner is also responsible for locking before calling this method.
      */
     public void processPendingAssistData() {
+        flushPendingAssistData();
+        tryDispatchRequestComplete();
+    }
+
+    private void flushPendingAssistData() {
         final int dataCount = mAssistData.size();
         for (int i = 0; i < dataCount; i++) {
             dispatchAssistDataReceived(mAssistData.get(i));
@@ -273,8 +293,9 @@
 
             if (mCallbacks.canHandleReceivedAssistDataLocked()) {
                 // Process any pending data and dispatch the new data as well
-                processPendingAssistData();
+                flushPendingAssistData();
                 dispatchAssistDataReceived(data);
+                tryDispatchRequestComplete();
             } else {
                 // Queue up the data for processing later
                 mAssistData.add(data);
@@ -292,8 +313,9 @@
 
             if (mCallbacks.canHandleReceivedAssistDataLocked()) {
                 // Process any pending data and dispatch the new data as well
-                processPendingAssistData();
+                flushPendingAssistData();
                 dispatchAssistScreenshotReceived(screenshot);
+                tryDispatchRequestComplete();
             } else {
                 // Queue up the data for processing later
                 mAssistScreenshot.add(screenshot);
@@ -317,6 +339,13 @@
         mCallbacks.onAssistScreenshotReceivedLocked(screenshot);
     }
 
+    private void tryDispatchRequestComplete() {
+        if (mPendingDataCount == 0 && mPendingScreenshotCount == 0 &&
+                mAssistData.isEmpty() && mAssistScreenshot.isEmpty()) {
+            mCallbacks.onAssistRequestCompleted();
+        }
+    }
+
     public void dump(String prefix, PrintWriter pw) {
         pw.print(prefix); pw.print("mPendingDataCount="); pw.println(mPendingDataCount);
         pw.print(prefix); pw.print("mAssistData="); pw.println(mAssistData);
diff --git a/services/core/java/com/android/server/wm/PinnedStackController.java b/services/core/java/com/android/server/wm/PinnedStackController.java
index 365366a..d8726bf 100644
--- a/services/core/java/com/android/server/wm/PinnedStackController.java
+++ b/services/core/java/com/android/server/wm/PinnedStackController.java
@@ -149,6 +149,9 @@
         @Override
         public void binderDied() {
             // Clean up the state if the listener dies
+            if (mPinnedStackListener != null) {
+                mPinnedStackListener.asBinder().unlinkToDeath(mPinnedStackListenerDeathHandler, 0);
+            }
             mPinnedStackListener = null;
         }
     }
diff --git a/services/tests/servicestests/src/com/android/server/am/AssistDataRequesterTest.java b/services/tests/servicestests/src/com/android/server/am/AssistDataRequesterTest.java
index 0d16ef3..ce88d84 100644
--- a/services/tests/servicestests/src/com/android/server/am/AssistDataRequesterTest.java
+++ b/services/tests/servicestests/src/com/android/server/am/AssistDataRequesterTest.java
@@ -22,6 +22,7 @@
 import static android.app.AppOpsManager.OP_ASSIST_STRUCTURE;
 import static android.graphics.Bitmap.Config.ARGB_8888;
 
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
@@ -130,8 +131,7 @@
             mHandler.post(() -> {
                 try {
                     mGate.await(10, TimeUnit.SECONDS);
-                    mDataRequester.onHandleAssistScreenshot(Bitmap.createBitmap(1, 1,
-                            ARGB_8888));
+                    mDataRequester.onHandleAssistScreenshot(Bitmap.createBitmap(1, 1, ARGB_8888));
                 } catch (InterruptedException e) {
                     Log.e(TAG, "Failed to wait", e);
                 }
@@ -197,11 +197,16 @@
         assertTrue(mDataRequester.getPendingScreenshotCount() == 0);
         assertTrue(mCallbacks.receivedData.isEmpty());
         assertTrue(mCallbacks.receivedScreenshots.isEmpty());
+        assertFalse(mCallbacks.requestCompleted);
 
         mCallbacks.canHandleReceivedData = true;
         mDataRequester.processPendingAssistData();
+        // Since we are posting the callback for the request-complete, flush the handler as well
+        mGate.countDown();
+        waitForIdle(mHandler);
         assertTrue(mCallbacks.receivedData.size() == 5);
         assertTrue(mCallbacks.receivedScreenshots.size() == 1);
+        assertTrue(mCallbacks.requestCompleted);
 
         // Clear the state and ensure that we only process pending data once
         mCallbacks.reset();
@@ -297,6 +302,7 @@
         assertTrue("Expected " + numPendingScreenshots + " pending screenshots, got "
                         + mDataRequester.getPendingScreenshotCount(),
                 mDataRequester.getPendingScreenshotCount() == numPendingScreenshots);
+        assertFalse("Expected request NOT completed", mCallbacks.requestCompleted);
         mGate.countDown();
         waitForIdle(mHandler);
         assertTrue("Expected " + numReceivedData + " data, received "
@@ -305,6 +311,7 @@
         assertTrue("Expected " + numReceivedScreenshots + " screenshots, received "
                         + mCallbacks.receivedScreenshots.size(),
                 mCallbacks.receivedScreenshots.size() == numReceivedScreenshots);
+        assertTrue("Expected request completed", mCallbacks.requestCompleted);
     }
 
     private List<IBinder> createActivityList(int size) {
@@ -324,9 +331,10 @@
         latch.await(2, TimeUnit.SECONDS);
     }
 
-    private static class Callbacks implements AssistDataRequesterCallbacks {
+    private class Callbacks implements AssistDataRequesterCallbacks {
 
         boolean canHandleReceivedData = true;
+        boolean requestCompleted = false;
         ArrayList<Bundle> receivedData = new ArrayList<>();
         ArrayList<Bitmap> receivedScreenshots = new ArrayList<>();
 
@@ -350,5 +358,17 @@
         public void onAssistScreenshotReceivedLocked(Bitmap screenshot) {
             receivedScreenshots.add(screenshot);
         }
+
+        @Override
+        public void onAssistRequestCompleted() {
+            mHandler.post(() -> {
+                try {
+                    mGate.await(10, TimeUnit.SECONDS);
+                    requestCompleted = true;
+                } catch (InterruptedException e) {
+                    Log.e(TAG, "Failed to wait", e);
+                }
+            });
+        }
     }
 }
\ No newline at end of file