From b3aa6562e5680f40842658abd2a40030b35381e9 Mon Sep 17 00:00:00 2001
From: George Hodulik <georgehodulik@google.com>
Date: Fri, 3 May 2019 13:41:39 -0700
Subject: [PATCH] Add null check in AppPredictionService callback wrapper.

Bug: 131917369
Test: Killed launcher and confirmed the binder died callback was called and callback was removed.
Test: compiles
Change-Id: I40612aeca2de2f197e17379355e6d9e46193ae5e
---
 .../appprediction/AppPredictionService.java   | 35 +++++++++++++++----
 1 file changed, 28 insertions(+), 7 deletions(-)

diff --git a/core/java/android/service/appprediction/AppPredictionService.java b/core/java/android/service/appprediction/AppPredictionService.java
index c1323bc8fd4c4..1391d43b00ca0 100644
--- a/core/java/android/service/appprediction/AppPredictionService.java
+++ b/core/java/android/service/appprediction/AppPredictionService.java
@@ -20,6 +20,7 @@ import static com.android.internal.util.function.pooled.PooledLambda.obtainMessa
 import android.annotation.CallSuper;
 import android.annotation.MainThread;
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.annotation.SystemApi;
 import android.annotation.TestApi;
 import android.app.Service;
@@ -100,7 +101,7 @@ public abstract class AppPredictionService extends Service {
             mHandler.sendMessage(
                     obtainMessage(AppPredictionService::onSortAppTargets,
                             AppPredictionService.this, sessionId, targets.getList(), null,
-                            new CallbackWrapper(callback)));
+                            new CallbackWrapper(callback, null)));
         }
 
         @Override
@@ -196,7 +197,9 @@ public abstract class AppPredictionService extends Service {
 
         final CallbackWrapper wrapper = findCallbackWrapper(callbacks, callback);
         if (wrapper == null) {
-            callbacks.add(new CallbackWrapper(callback));
+            callbacks.add(new CallbackWrapper(callback,
+                    callbackWrapper ->
+                        mHandler.post(() -> removeCallbackWrapper(callbacks, callbackWrapper))));
             if (callbacks.size() == 1) {
                 onStartPredictionUpdates();
             }
@@ -219,10 +222,18 @@ public abstract class AppPredictionService extends Service {
 
         final CallbackWrapper wrapper = findCallbackWrapper(callbacks, callback);
         if (wrapper != null) {
-            callbacks.remove(wrapper);
-            if (callbacks.isEmpty()) {
-                onStopPredictionUpdates();
-            }
+            removeCallbackWrapper(callbacks, wrapper);
+        }
+    }
+
+    private void removeCallbackWrapper(
+                ArrayList<CallbackWrapper> callbacks, CallbackWrapper wrapper) {
+        if (callbacks == null) {
+            return;
+        }
+        callbacks.remove(wrapper);
+        if (callbacks.isEmpty()) {
+            onStopPredictionUpdates();
         }
     }
 
@@ -295,9 +306,12 @@ public abstract class AppPredictionService extends Service {
             IBinder.DeathRecipient {
 
         private IPredictionCallback mCallback;
+        private final Consumer<CallbackWrapper> mOnBinderDied;
 
-        CallbackWrapper(IPredictionCallback callback) {
+        CallbackWrapper(IPredictionCallback callback,
+                @Nullable Consumer<CallbackWrapper> onBinderDied) {
             mCallback = callback;
+            mOnBinderDied = onBinderDied;
             try {
                 mCallback.asBinder().linkToDeath(this, 0);
             } catch (RemoteException e) {
@@ -306,6 +320,10 @@ public abstract class AppPredictionService extends Service {
         }
 
         public boolean isCallback(@NonNull IPredictionCallback callback) {
+            if (mCallback == null) {
+                Slog.e(TAG, "Callback is null, likely the binder has died.");
+                return false;
+            }
             return mCallback.equals(callback);
         }
 
@@ -323,6 +341,9 @@ public abstract class AppPredictionService extends Service {
         @Override
         public void binderDied() {
             mCallback = null;
+            if (mOnBinderDied != null) {
+                mOnBinderDied.accept(this);
+            }
         }
     }
 }