[CDM] retroactively add restored associations upon app installation

Bug: 314992577
Test: atest CtsCompanionDeviceManagerCoreTestCases:BackupAndRestoreTest
Change-Id: Idad3a345746c73e38061d6702585debd69c4973b

6 files changed, 112 insertions, 16 deletions

diff --git a/core/java/android/companion/AssociationInfo.java b/core/java/android/companion/AssociationInfo.java
index cdb92acc5256..843158c0e9fb 100644
--- a/core/java/android/companion/AssociationInfo.java
+++ b/core/java/android/companion/AssociationInfo.java
@@ -71,6 +71,12 @@ public final class AssociationInfo implements Parcelable {
      * @see CompanionDeviceManager#disassociate(int)
      */
     private final boolean mRevoked;
+    /**
+     * Indicates that the association is waiting for its corresponding companion app to be installed
+     * before it can be added to CDM. This is likely because it was restored onto the device from a
+     * backup.
+     */
+    private final boolean mPending;
     private final long mTimeApprovedMs;
     /**
      * A long value indicates the last time connected reported by selfManaged devices
@@ -88,7 +94,7 @@ public final class AssociationInfo implements Parcelable {
             @Nullable String tag, @Nullable MacAddress macAddress,
             @Nullable CharSequence displayName, @Nullable String deviceProfile,
             @Nullable AssociatedDevice associatedDevice, boolean selfManaged,
-            boolean notifyOnDeviceNearby, boolean revoked, long timeApprovedMs,
+            boolean notifyOnDeviceNearby, boolean revoked, boolean pending, long timeApprovedMs,
             long lastTimeConnectedMs, int systemDataSyncFlags) {
         if (id <= 0) {
             throw new IllegalArgumentException("Association ID should be greater than 0");
@@ -109,6 +115,7 @@ public final class AssociationInfo implements Parcelable {
         mSelfManaged = selfManaged;
         mNotifyOnDeviceNearby = notifyOnDeviceNearby;
         mRevoked = revoked;
+        mPending = pending;
         mTimeApprovedMs = timeApprovedMs;
         mLastTimeConnectedMs = lastTimeConnectedMs;
         mSystemDataSyncFlags = systemDataSyncFlags;
@@ -236,6 +243,15 @@ public final class AssociationInfo implements Parcelable {
     }
 
     /**
+     * @return true if the association is waiting for its corresponding app to be installed
+     * before it can be added to CDM.
+     * @hide
+     */
+    public boolean isPending() {
+        return mPending;
+    }
+
+    /**
      * @return the last time self reported disconnected for selfManaged only.
      * @hide
      */
@@ -318,6 +334,7 @@ public final class AssociationInfo implements Parcelable {
                 + ", mAssociatedDevice=" + mAssociatedDevice
                 + ", mNotifyOnDeviceNearby=" + mNotifyOnDeviceNearby
                 + ", mRevoked=" + mRevoked
+                + ", mPending=" + mPending
                 + ", mTimeApprovedMs=" + new Date(mTimeApprovedMs)
                 + ", mLastTimeConnectedMs=" + (
                     mLastTimeConnectedMs == Long.MAX_VALUE
@@ -336,6 +353,7 @@ public final class AssociationInfo implements Parcelable {
                 && mSelfManaged == that.mSelfManaged
                 && mNotifyOnDeviceNearby == that.mNotifyOnDeviceNearby
                 && mRevoked == that.mRevoked
+                && mPending == that.mPending
                 && mTimeApprovedMs == that.mTimeApprovedMs
                 && mLastTimeConnectedMs == that.mLastTimeConnectedMs
                 && Objects.equals(mPackageName, that.mPackageName)
@@ -351,7 +369,7 @@ public final class AssociationInfo implements Parcelable {
     public int hashCode() {
         return Objects.hash(mId, mUserId, mPackageName, mTag, mDeviceMacAddress, mDisplayName,
                 mDeviceProfile, mAssociatedDevice, mSelfManaged, mNotifyOnDeviceNearby, mRevoked,
-                mTimeApprovedMs, mLastTimeConnectedMs, mSystemDataSyncFlags);
+                mPending, mTimeApprovedMs, mLastTimeConnectedMs, mSystemDataSyncFlags);
     }
 
     @Override
@@ -372,6 +390,7 @@ public final class AssociationInfo implements Parcelable {
         dest.writeBoolean(mSelfManaged);
         dest.writeBoolean(mNotifyOnDeviceNearby);
         dest.writeBoolean(mRevoked);
+        dest.writeBoolean(mPending);
         dest.writeLong(mTimeApprovedMs);
         dest.writeLong(mLastTimeConnectedMs);
         dest.writeInt(mSystemDataSyncFlags);
@@ -389,6 +408,7 @@ public final class AssociationInfo implements Parcelable {
         mSelfManaged = in.readBoolean();
         mNotifyOnDeviceNearby = in.readBoolean();
         mRevoked = in.readBoolean();
+        mPending = in.readBoolean();
         mTimeApprovedMs = in.readLong();
         mLastTimeConnectedMs = in.readLong();
         mSystemDataSyncFlags = in.readInt();
@@ -427,6 +447,7 @@ public final class AssociationInfo implements Parcelable {
         private boolean mSelfManaged;
         private boolean mNotifyOnDeviceNearby;
         private boolean mRevoked;
+        private boolean mPending;
         private long mTimeApprovedMs;
         private long mLastTimeConnectedMs;
         private int mSystemDataSyncFlags;
@@ -453,6 +474,7 @@ public final class AssociationInfo implements Parcelable {
             mSelfManaged = info.mSelfManaged;
             mNotifyOnDeviceNearby = info.mNotifyOnDeviceNearby;
             mRevoked = info.mRevoked;
+            mPending = info.mPending;
             mTimeApprovedMs = info.mTimeApprovedMs;
             mLastTimeConnectedMs = info.mLastTimeConnectedMs;
             mSystemDataSyncFlags = info.mSystemDataSyncFlags;
@@ -476,6 +498,7 @@ public final class AssociationInfo implements Parcelable {
             mSelfManaged = info.mSelfManaged;
             mNotifyOnDeviceNearby = info.mNotifyOnDeviceNearby;
             mRevoked = info.mRevoked;
+            mPending = info.mPending;
             mTimeApprovedMs = info.mTimeApprovedMs;
             mLastTimeConnectedMs = info.mLastTimeConnectedMs;
             mSystemDataSyncFlags = info.mSystemDataSyncFlags;
@@ -549,6 +572,14 @@ public final class AssociationInfo implements Parcelable {
         }
 
         /** @hide */
+        @NonNull
+        @SuppressLint("MissingGetterMatchingBuilder")
+        public Builder setPending(boolean pending) {
+            mPending = pending;
+            return this;
+        }
+
+        /** @hide */
         @TestApi
         @NonNull
         @SuppressLint("MissingGetterMatchingBuilder")
@@ -606,6 +637,7 @@ public final class AssociationInfo implements Parcelable {
                     mSelfManaged,
                     mNotifyOnDeviceNearby,
                     mRevoked,
+                    mPending,
                     mTimeApprovedMs,
                     mLastTimeConnectedMs,
                     mSystemDataSyncFlags
diff --git a/services/companion/java/com/android/server/companion/AssociationRequestsProcessor.java b/services/companion/java/com/android/server/companion/AssociationRequestsProcessor.java
index a6ed8464128a..4b3772a7a54d 100644
--- a/services/companion/java/com/android/server/companion/AssociationRequestsProcessor.java
+++ b/services/companion/java/com/android/server/companion/AssociationRequestsProcessor.java
@@ -284,7 +284,7 @@ class AssociationRequestsProcessor {
         final AssociationInfo association = new AssociationInfo(id, userId, packageName,
                 /* tag */ null, macAddress, displayName, deviceProfile, associatedDevice,
                 selfManaged, /* notifyOnDeviceNearby */ false, /* revoked */ false,
-                timestamp, Long.MAX_VALUE, /* systemDataSyncFlags */ 0);
+                /* pending */ false, timestamp, Long.MAX_VALUE, /* systemDataSyncFlags */ 0);
 
         // Add role holder for association (if specified) and add new association to store.
         maybeGrantRoleAndStoreAssociation(association, callback, resultReceiver);
diff --git a/services/companion/java/com/android/server/companion/BackupRestoreProcessor.java b/services/companion/java/com/android/server/companion/BackupRestoreProcessor.java
index a7dbd1c15aec..18cf46f8af24 100644
--- a/services/companion/java/com/android/server/companion/BackupRestoreProcessor.java
+++ b/services/companion/java/com/android/server/companion/BackupRestoreProcessor.java
@@ -18,6 +18,8 @@ package com.android.server.companion;
 
 import static android.os.UserHandle.getCallingUserId;
 
+import static com.android.server.companion.CompanionDeviceManagerService.PerUserAssociationSet;
+
 import android.annotation.NonNull;
 import android.annotation.SuppressLint;
 import android.annotation.UserIdInt;
@@ -26,9 +28,11 @@ import android.companion.Flags;
 import android.companion.datatransfer.SystemDataTransferRequest;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManagerInternal;
+import android.util.ArraySet;
 import android.util.Log;
 import android.util.Slog;
 
+import com.android.internal.annotations.GuardedBy;
 import com.android.internal.util.CollectionUtils;
 import com.android.server.companion.datatransfer.SystemDataTransferRequestStore;
 
@@ -58,6 +62,14 @@ class BackupRestoreProcessor {
     @NonNull
     private final AssociationRequestsProcessor mAssociationRequestsProcessor;
 
+    /**
+     * A structure that consists of a set of restored associations that are pending corresponding
+     * companion app to be installed.
+     */
+    @GuardedBy("mAssociationsPendingAppInstall")
+    private final PerUserAssociationSet mAssociationsPendingAppInstall =
+            new PerUserAssociationSet();
+
     BackupRestoreProcessor(@NonNull CompanionDeviceManagerService service,
                            @NonNull AssociationStoreImpl associationStore,
                            @NonNull PersistentDataStore persistentStore,
@@ -170,7 +182,7 @@ class BackupRestoreProcessor {
                 mAssociationRequestsProcessor.maybeGrantRoleAndStoreAssociation(newAssociation,
                         null, null);
             } else {
-                // TODO(b/314992577): Check if package is installed before granting
+                addToPendingAppInstall(newAssociation);
             }
 
             // Re-map restored system data transfer requests to newly created associations
@@ -185,6 +197,30 @@ class BackupRestoreProcessor {
         mService.persistStateForUser(userId);
     }
 
+    void addToPendingAppInstall(@NonNull AssociationInfo association) {
+        association = (new AssociationInfo.Builder(association))
+                .setPending(true)
+                .build();
+
+        synchronized (mAssociationsPendingAppInstall) {
+            mAssociationsPendingAppInstall.forUser(association.getUserId()).add(association);
+        }
+    }
+
+    void removeFromPendingAppInstall(@NonNull AssociationInfo association) {
+        synchronized (mAssociationsPendingAppInstall) {
+            mAssociationsPendingAppInstall.forUser(association.getUserId()).remove(association);
+        }
+    }
+
+    @NonNull
+    Set<AssociationInfo> getAssociationsPendingAppInstallForUser(@UserIdInt int userId) {
+        synchronized (mAssociationsPendingAppInstall) {
+            // Return a copy.
+            return new ArraySet<>(mAssociationsPendingAppInstall.forUser(userId));
+        }
+    }
+
     /**
      * Detects and handles collision between restored association and local association. Returns
      * true if there has been a collision and false otherwise.
diff --git a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
index 858887ae20c6..056ec895821d 100644
--- a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
+++ b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
@@ -287,7 +287,9 @@ public class CompanionDeviceManagerService extends SystemService {
         final Set<Integer> usersToPersistStateFor = new ArraySet<>();
 
         for (AssociationInfo association : allAssociations) {
-            if (!association.isRevoked()) {
+            if (association.isPending()) {
+                mBackupRestoreProcessor.addToPendingAppInstall(association);
+            } else if (!association.isRevoked()) {
                 activeAssociations.add(association);
             } else if (maybeRemoveRoleHolderForAssociation(association)) {
                 // Nothing more to do here, but we'll need to persist all the associations to the
@@ -514,6 +516,9 @@ public class CompanionDeviceManagerService extends SystemService {
                 mAssociationStore.getAssociationsForUser(userId));
         // ... and add the revoked (removed) association, that are yet to be permanently removed.
         allAssociations.addAll(getPendingRoleHolderRemovalAssociationsForUser(userId));
+        // ... and add the restored associations that are pending missing package installation.
+        allAssociations.addAll(mBackupRestoreProcessor
+                .getAssociationsPendingAppInstallForUser(userId));
 
         final Map<String, Set<Integer>> usedIdsForUser = getPreviouslyUsedIdsForUser(userId);
 
@@ -583,7 +588,19 @@ public class CompanionDeviceManagerService extends SystemService {
 
     private void onPackageAddedInternal(@UserIdInt int userId, @NonNull String packageName) {
         if (DEBUG) Log.i(TAG, "onPackageAddedInternal() u" + userId + "/" + packageName);
-        // TODO(b/314992577): Retroactively grant roles for restored associations
+
+        Set<AssociationInfo> associationsPendingAppInstall = mBackupRestoreProcessor
+                .getAssociationsPendingAppInstallForUser(userId);
+        for (AssociationInfo association : associationsPendingAppInstall) {
+            if (!packageName.equals(association.getPackageName())) continue;
+
+            AssociationInfo newAssociation = new AssociationInfo.Builder(association)
+                    .setPending(false)
+                    .build();
+            mAssociationRequestsProcessor.maybeGrantRoleAndStoreAssociation(newAssociation,
+                    null, null);
+            mBackupRestoreProcessor.removeFromPendingAppInstall(association);
+        }
     }
 
     // Revoke associations if the selfManaged companion device does not connect for 3 months.
@@ -1152,6 +1169,15 @@ public class CompanionDeviceManagerService extends SystemService {
                 usedIds.put(it.getId(), true);
             }
 
+            // Some IDs may be reserved by associations that aren't stored yet due to missing
+            // package after a backup restoration. We don't want the ID to have been taken by
+            // another association by the time when it is activated from the package installation.
+            final Set<AssociationInfo> pendingAssociations = mBackupRestoreProcessor
+                    .getAssociationsPendingAppInstallForUser(userId);
+            for (AssociationInfo it: pendingAssociations) {
+                usedIds.put(it.getId(), true);
+            }
+
             // Second: collect all IDs that have been previously used for this package (and user).
             final Set<Integer> previouslyUsedIds =
                     getPreviouslyUsedIdsForPackageLocked(userId, packageName);
@@ -1718,7 +1744,7 @@ public class CompanionDeviceManagerService extends SystemService {
         }
     }
 
-    private static class PerUserAssociationSet extends PerUser<Set<AssociationInfo>> {
+    static class PerUserAssociationSet extends PerUser<Set<AssociationInfo>> {
         @Override
         protected @NonNull Set<AssociationInfo> create(int userId) {
             return new ArraySet<>();
diff --git a/services/companion/java/com/android/server/companion/PersistentDataStore.java b/services/companion/java/com/android/server/companion/PersistentDataStore.java
index dbaf7e85b7fa..1ebe65c6aa5f 100644
--- a/services/companion/java/com/android/server/companion/PersistentDataStore.java
+++ b/services/companion/java/com/android/server/companion/PersistentDataStore.java
@@ -189,6 +189,7 @@ final class PersistentDataStore {
     private static final String XML_ATTR_SELF_MANAGED = "self_managed";
     private static final String XML_ATTR_NOTIFY_DEVICE_NEARBY = "notify_device_nearby";
     private static final String XML_ATTR_REVOKED = "revoked";
+    private static final String XML_ATTR_PENDING = "pending";
     private static final String XML_ATTR_TIME_APPROVED = "time_approved";
     private static final String XML_ATTR_LAST_TIME_CONNECTED = "last_time_connected";
     private static final String XML_ATTR_SYSTEM_DATA_SYNC_FLAGS = "system_data_sync_flags";
@@ -464,8 +465,8 @@ final class PersistentDataStore {
 
         out.add(new AssociationInfo(associationId, userId, appPackage, tag,
                 MacAddress.fromString(deviceAddress), null, profile, null,
-                /* managedByCompanionApp */ false, notify, /* revoked */ false, timeApproved,
-                Long.MAX_VALUE, /* systemDataSyncFlags */ 0));
+                /* managedByCompanionApp */ false, notify, /* revoked */ false, /* pending */ false,
+                timeApproved, Long.MAX_VALUE, /* systemDataSyncFlags */ 0));
     }
 
     private static void readAssociationsV1(@NonNull TypedXmlPullParser parser,
@@ -496,6 +497,7 @@ final class PersistentDataStore {
         final boolean selfManaged = readBooleanAttribute(parser, XML_ATTR_SELF_MANAGED);
         final boolean notify = readBooleanAttribute(parser, XML_ATTR_NOTIFY_DEVICE_NEARBY);
         final boolean revoked = readBooleanAttribute(parser, XML_ATTR_REVOKED, false);
+        final boolean pending = readBooleanAttribute(parser, XML_ATTR_PENDING, false);
         final long timeApproved = readLongAttribute(parser, XML_ATTR_TIME_APPROVED, 0L);
         final long lastTimeConnected = readLongAttribute(
                 parser, XML_ATTR_LAST_TIME_CONNECTED, Long.MAX_VALUE);
@@ -504,7 +506,7 @@ final class PersistentDataStore {
 
         final AssociationInfo associationInfo = createAssociationInfoNoThrow(associationId, userId,
                 appPackage, tag, macAddress, displayName, profile, selfManaged, notify, revoked,
-                timeApproved, lastTimeConnected, systemDataSyncFlags);
+                pending, timeApproved, lastTimeConnected, systemDataSyncFlags);
         if (associationInfo != null) {
             out.add(associationInfo);
         }
@@ -558,8 +560,8 @@ final class PersistentDataStore {
         writeBooleanAttribute(serializer, XML_ATTR_SELF_MANAGED, a.isSelfManaged());
         writeBooleanAttribute(
                 serializer, XML_ATTR_NOTIFY_DEVICE_NEARBY, a.isNotifyOnDeviceNearby());
-        writeBooleanAttribute(
-                serializer, XML_ATTR_REVOKED, a.isRevoked());
+        writeBooleanAttribute(serializer, XML_ATTR_REVOKED, a.isRevoked());
+        writeBooleanAttribute(serializer, XML_ATTR_PENDING, a.isPending());
         writeLongAttribute(serializer, XML_ATTR_TIME_APPROVED, a.getTimeApprovedMs());
         writeLongAttribute(
                 serializer, XML_ATTR_LAST_TIME_CONNECTED, a.getLastTimeConnectedMs());
@@ -603,14 +605,14 @@ final class PersistentDataStore {
             @UserIdInt int userId, @NonNull String appPackage, @Nullable String tag,
             @Nullable MacAddress macAddress, @Nullable CharSequence displayName,
             @Nullable String profile, boolean selfManaged, boolean notify, boolean revoked,
-            long timeApproved, long lastTimeConnected, int systemDataSyncFlags) {
+            boolean pending, long timeApproved, long lastTimeConnected, int systemDataSyncFlags) {
         AssociationInfo associationInfo = null;
         try {
             // We do not persist AssociatedDevice, which means that AssociationInfo retrieved from
             // datastore is not guaranteed to be identical to the one from initial association.
             associationInfo = new AssociationInfo(associationId, userId, appPackage, tag,
                     macAddress, displayName, profile, null, selfManaged, notify,
-                    revoked, timeApproved, lastTimeConnected, systemDataSyncFlags);
+                    revoked, pending, timeApproved, lastTimeConnected, systemDataSyncFlags);
         } catch (Exception e) {
             if (DEBUG) Log.w(TAG, "Could not create AssociationInfo", e);
         }
diff --git a/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java b/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java
index 995d1f4d5520..276c8321fb65 100644
--- a/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java
@@ -1982,8 +1982,8 @@ public class VirtualDeviceManagerServiceTest {
         return new AssociationInfo(associationId, /* userId= */ 0, /* packageName=*/ null,
                 /* tag= */ null, MacAddress.BROADCAST_ADDRESS, /* displayName= */ "", deviceProfile,
                 /* associatedDevice= */ null, /* selfManaged= */ true,
-                /* notifyOnDeviceNearby= */ false, /* revoked= */false,  /* timeApprovedMs= */0,
-                /* lastTimeConnectedMs= */0, /* systemDataSyncFlags= */ -1);
+                /* notifyOnDeviceNearby= */ false, /* revoked= */ false, /* pending= */ false,
+                /* timeApprovedMs= */0, /* lastTimeConnectedMs= */0, /* systemDataSyncFlags= */ -1);
     }
 
     /** Helper class to drop permissions temporarily and restore them at the end of a test. */