Avoid activity leak via Companion callback

Test: invoke associate() API and ensure it still works
Change-Id: I9aedb043b4b1f0d77f076d8753cd60ff7c98a7d6

1 files changed, 28 insertions, 3 deletions

diff --git a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
index 6093241a6681..73f1705e35df 100644
--- a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
+++ b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
@@ -110,10 +110,15 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
     private final CompanionDeviceManagerImpl mImpl;
     private final ConcurrentMap<Integer, AtomicFile> mUidToStorage = new ConcurrentHashMap<>();
     private IDeviceIdleController mIdleController;
-    private IFindDeviceCallback mFindDeviceCallback;
     private ServiceConnection mServiceConnection;
     private IAppOpsService mAppOpsManager;
 
+    private IFindDeviceCallback mFindDeviceCallback;
+    private AssociationRequest mRequest;
+    private String mCallingPackage;
+
+    private final Object mLock = new Object();
+
     public CompanionDeviceManagerService(Context context) {
         super(context);
         mImpl = new CompanionDeviceManagerImpl();
@@ -156,8 +161,12 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
     }
 
     private void cleanup() {
-        mServiceConnection = unbind(mServiceConnection);
-        mFindDeviceCallback = unlinkToDeath(mFindDeviceCallback, this, 0);
+        synchronized (mLock) {
+            mServiceConnection = unbind(mServiceConnection);
+            mFindDeviceCallback = unlinkToDeath(mFindDeviceCallback, this, 0);
+            mRequest = null;
+            mCallingPackage = null;
+        }
     }
 
     /**
@@ -222,6 +231,17 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
         }
 
         @Override
+        public void stopScan(AssociationRequest request,
+                IFindDeviceCallback callback,
+                String callingPackage) {
+            if (Objects.equals(request, mRequest)
+                    && Objects.equals(callback, mFindDeviceCallback)
+                    && Objects.equals(callingPackage, mCallingPackage)) {
+                cleanup();
+            }
+        }
+
+        @Override
         public List<String> getAssociations(String callingPackage, int userId)
                 throws RemoteException {
             checkCallerIsSystemOr(callingPackage, userId);
@@ -340,7 +360,11 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
                             "onServiceConnected(name = " + name + ", service = "
                                     + service + ")");
                 }
+
                 mFindDeviceCallback = findDeviceCallback;
+                mRequest = request;
+                mCallingPackage = callingPackage;
+
                 try {
                     mFindDeviceCallback.asBinder().linkToDeath(
                             CompanionDeviceManagerService.this, 0);
@@ -348,6 +372,7 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
                     cleanup();
                     return;
                 }
+
                 try {
                     ICompanionDeviceDiscoveryService.Stub
                             .asInterface(service)