Remove INSTALLED_USERS extra from ACTION_PACKAGE_ENABLE_ROLLBACK.

This changes RollbackManagerServiceImpl to use the list of installed users
that is passed in to snapshotAndRestoreUserData for the purpose of snapshotting.

With this change, it is no longer necessary for the list of installed users to be
added as an extra to the ENABLE_ROLLBACK intent, so that extra
(EXTRA_ENABLE_ROLLBACK_INSTALLED_USERS) is removed.

Bug: 139181444
Test: atest RollbackTest
Test: atest AppDataRollbackHelperTest
Change-Id: Iffd63977a7498e3cba20d2039af623c921d38f57
diff --git a/core/java/android/content/pm/PackageManagerInternal.java b/core/java/android/content/pm/PackageManagerInternal.java
index 2f198ac..f50502e 100644
--- a/core/java/android/content/pm/PackageManagerInternal.java
+++ b/core/java/android/content/pm/PackageManagerInternal.java
@@ -675,12 +675,6 @@
             "android.content.pm.extra.ENABLE_ROLLBACK_INSTALL_FLAGS";
 
     /**
-     * Extra field name for the set of installed users for a given rollback package.
-     */
-    public static final String EXTRA_ENABLE_ROLLBACK_INSTALLED_USERS =
-            "android.content.pm.extra.ENABLE_ROLLBACK_INSTALLED_USERS";
-
-    /**
      * Extra field name for the user id an install is associated with when
      * enabling rollback.
      */
diff --git a/core/java/android/content/rollback/PackageRollbackInfo.java b/core/java/android/content/rollback/PackageRollbackInfo.java
index 2014751..c89796d 100644
--- a/core/java/android/content/rollback/PackageRollbackInfo.java
+++ b/core/java/android/content/rollback/PackageRollbackInfo.java
@@ -76,10 +76,10 @@
     private final boolean mIsApex;
 
     /*
-     * The list of users the package is installed for.
+     * The list of users for which snapshots have been saved.
      */
     // NOTE: Not a part of the Parcelable representation of this object.
-    private final IntArray mInstalledUsers;
+    private final IntArray mSnapshottedUsers;
 
     /**
      * A mapping between user and an inode of theirs CE data snapshot.
@@ -148,8 +148,8 @@
     }
 
     /** @hide */
-    public IntArray getInstalledUsers() {
-        return mInstalledUsers;
+    public IntArray getSnapshottedUsers() {
+        return mSnapshottedUsers;
     }
 
     /** @hide */
@@ -179,14 +179,14 @@
     public PackageRollbackInfo(VersionedPackage packageRolledBackFrom,
             VersionedPackage packageRolledBackTo,
             @NonNull IntArray pendingBackups, @NonNull ArrayList<RestoreInfo> pendingRestores,
-            boolean isApex, @NonNull IntArray installedUsers,
+            boolean isApex, @NonNull IntArray snapshottedUsers,
             @NonNull SparseLongArray ceSnapshotInodes) {
         this.mVersionRolledBackFrom = packageRolledBackFrom;
         this.mVersionRolledBackTo = packageRolledBackTo;
         this.mPendingBackups = pendingBackups;
         this.mPendingRestores = pendingRestores;
         this.mIsApex = isApex;
-        this.mInstalledUsers = installedUsers;
+        this.mSnapshottedUsers = snapshottedUsers;
         this.mCeSnapshotInodes = ceSnapshotInodes;
     }
 
@@ -196,7 +196,7 @@
         this.mIsApex = in.readBoolean();
         this.mPendingRestores = null;
         this.mPendingBackups = null;
-        this.mInstalledUsers = null;
+        this.mSnapshottedUsers = null;
         this.mCeSnapshotInodes = null;
     }
 
diff --git a/services/core/java/com/android/server/pm/PackageManagerService.java b/services/core/java/com/android/server/pm/PackageManagerService.java
index 3f6d4df..f51e0cc 100644
--- a/services/core/java/com/android/server/pm/PackageManagerService.java
+++ b/services/core/java/com/android/server/pm/PackageManagerService.java
@@ -14580,17 +14580,6 @@
                             TRACE_TAG_PACKAGE_MANAGER, "enable_rollback", enableRollbackToken);
                     mPendingEnableRollback.append(enableRollbackToken, this);
 
-                    final int[] installedUsers;
-                    synchronized (mLock) {
-                        PackageSetting ps = mSettings.getPackageLPr(pkgLite.packageName);
-                        if (ps != null) {
-                            installedUsers = ps.queryInstalledUsers(mUserManager.getUserIds(),
-                                    true);
-                        } else {
-                            installedUsers = new int[0];
-                        }
-                    }
-
                     Intent enableRollbackIntent = new Intent(Intent.ACTION_PACKAGE_ENABLE_ROLLBACK);
                     enableRollbackIntent.putExtra(
                             PackageManagerInternal.EXTRA_ENABLE_ROLLBACK_TOKEN,
@@ -14599,9 +14588,6 @@
                             PackageManagerInternal.EXTRA_ENABLE_ROLLBACK_INSTALL_FLAGS,
                             installFlags);
                     enableRollbackIntent.putExtra(
-                            PackageManagerInternal.EXTRA_ENABLE_ROLLBACK_INSTALLED_USERS,
-                            installedUsers);
-                    enableRollbackIntent.putExtra(
                             PackageManagerInternal.EXTRA_ENABLE_ROLLBACK_USER,
                             getRollbackUser().getIdentifier());
                     enableRollbackIntent.setDataAndType(Uri.fromFile(new File(origin.resolvedPath)),
diff --git a/services/core/java/com/android/server/rollback/AppDataRollbackHelper.java b/services/core/java/com/android/server/rollback/AppDataRollbackHelper.java
index f5f7d67..563ccbe 100644
--- a/services/core/java/com/android/server/rollback/AppDataRollbackHelper.java
+++ b/services/core/java/com/android/server/rollback/AppDataRollbackHelper.java
@@ -52,12 +52,13 @@
     }
 
     /**
-     * Creates an app data snapshot for a specified {@code packageRollbackInfo}. Updates said {@code
-     * packageRollbackInfo} with the inodes of the CE user data snapshot folders.
+     * Creates an app data snapshot for a specified {@code packageRollbackInfo} and the specified
+     * {@code userIds}. Updates said {@code packageRollbackInfo} with the inodes of the CE user data
+     * snapshot folders.
      */
-    public void snapshotAppData(int snapshotId, PackageRollbackInfo packageRollbackInfo) {
-        final int[] installedUsers = packageRollbackInfo.getInstalledUsers().toArray();
-        for (int user : installedUsers) {
+    public void snapshotAppData(
+            int snapshotId, PackageRollbackInfo packageRollbackInfo, int[] userIds) {
+        for (int user : userIds) {
             final int storageFlags;
             if (isUserCredentialLocked(user)) {
                 // We've encountered a user that hasn't unlocked on a FBE device, so we can't copy
@@ -80,6 +81,7 @@
                         + packageRollbackInfo.getPackageName() + ", userId: " + user, ie);
             }
         }
+        packageRollbackInfo.getSnapshottedUsers().addAll(IntArray.wrap(userIds));
     }
 
     /**
diff --git a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
index f0cf9cf..1b6af40 100644
--- a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
+++ b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
@@ -177,16 +177,14 @@
                             PackageManagerInternal.EXTRA_ENABLE_ROLLBACK_TOKEN, -1);
                     int installFlags = intent.getIntExtra(
                             PackageManagerInternal.EXTRA_ENABLE_ROLLBACK_INSTALL_FLAGS, 0);
-                    int[] installedUsers = intent.getIntArrayExtra(
-                            PackageManagerInternal.EXTRA_ENABLE_ROLLBACK_INSTALLED_USERS);
                     int user = intent.getIntExtra(
                             PackageManagerInternal.EXTRA_ENABLE_ROLLBACK_USER, 0);
 
                     File newPackageCodePath = new File(intent.getData().getPath());
 
                     getHandler().post(() -> {
-                        boolean success = enableRollback(installFlags, newPackageCodePath,
-                                installedUsers, user, token);
+                        boolean success =
+                                enableRollback(installFlags, newPackageCodePath, user, token);
                         int ret = PackageManagerInternal.ENABLE_ROLLBACK_SUCCEEDED;
                         if (!success) {
                             ret = PackageManagerInternal.ENABLE_ROLLBACK_FAILED;
@@ -821,13 +819,12 @@
      *
      * @param installFlags information about what is being installed.
      * @param newPackageCodePath path to the package about to be installed.
-     * @param installedUsers the set of users for which a given package is installed.
      * @param user the user that owns the install session to enable rollback on.
      * @param token the distinct rollback token sent by package manager.
      * @return true if enabling the rollback succeeds, false otherwise.
      */
-    private boolean enableRollback(int installFlags, File newPackageCodePath,
-            int[] installedUsers, @UserIdInt int user, int token) {
+    private boolean enableRollback(
+            int installFlags, File newPackageCodePath, @UserIdInt int user, int token) {
 
         // Find the session id associated with this install.
         // TODO: It would be nice if package manager or package installer told
@@ -872,38 +869,15 @@
 
         // Check to see if this is the apk session for a staged session with
         // rollback enabled.
-        // TODO: This check could be made more efficient.
-        RollbackData rd = null;
         synchronized (mLock) {
             for (int i = 0; i < mRollbacks.size(); ++i) {
                 RollbackData data = mRollbacks.get(i);
                 if (data.apkSessionId == parentSession.getSessionId()) {
-                    rd = data;
-                    break;
-                }
-            }
-        }
-
-        if (rd != null) {
-            // This is the apk session for a staged session. We do not need to create a new rollback
-            // for this session.
-            PackageParser.PackageLite newPackage = null;
-            try {
-                newPackage = PackageParser.parsePackageLite(
-                        new File(packageSession.resolvedBaseCodePath), 0);
-            } catch (PackageParser.PackageParserException e) {
-                Slog.e(TAG, "Unable to parse new package", e);
-                return false;
-            }
-            String packageName = newPackage.packageName;
-            for (PackageRollbackInfo info : rd.info.getPackages()) {
-                if (info.getPackageName().equals(packageName)) {
-                    info.getInstalledUsers().addAll(IntArray.wrap(installedUsers));
+                    // This is the apk session for a staged session with rollback enabled. We do not
+                    // need to create a new rollback for this session.
                     return true;
                 }
             }
-            Slog.e(TAG, "Unable to find package in apk session");
-            return false;
         }
 
         NewRollback newRollback;
@@ -919,7 +893,7 @@
         }
         newRollback.addToken(token);
 
-        return enableRollbackForPackageSession(newRollback.data, packageSession, installedUsers);
+        return enableRollbackForPackageSession(newRollback.data, packageSession);
     }
 
     /**
@@ -930,7 +904,7 @@
      * @return true on success, false on failure.
      */
     private boolean enableRollbackForPackageSession(RollbackData data,
-            PackageInstaller.SessionInfo session, @NonNull int[] installedUsers) {
+            PackageInstaller.SessionInfo session) {
         // TODO: Don't attempt to enable rollback for split installs.
         final int installFlags = session.installFlags;
         if ((installFlags & PackageManager.INSTALL_ENABLE_ROLLBACK) == 0) {
@@ -988,8 +962,7 @@
         PackageRollbackInfo packageRollbackInfo = new PackageRollbackInfo(
                 newVersion, installedVersion,
                 new IntArray() /* pendingBackups */, new ArrayList<>() /* pendingRestores */,
-                isApex, IntArray.wrap(installedUsers),
-                new SparseLongArray() /* ceSnapshotInodes */);
+                isApex, new IntArray(), new SparseLongArray() /* ceSnapshotInodes */);
 
         try {
             ApplicationInfo appInfo = pkgInfo.applicationInfo;
@@ -1019,7 +992,7 @@
         }
 
         getHandler().post(() -> {
-            snapshotUserDataInternal(packageName);
+            snapshotUserDataInternal(packageName, userIds);
             restoreUserDataInternal(packageName, userIds, appId, ceDataInode, seInfo, token);
             final PackageManagerInternal pmi = LocalServices.getService(
                     PackageManagerInternal.class);
@@ -1027,7 +1000,7 @@
         });
     }
 
-    private void snapshotUserDataInternal(String packageName) {
+    private void snapshotUserDataInternal(String packageName, int[] userIds) {
         synchronized (mLock) {
             // staged installs
             for (int i = 0; i < mRollbacks.size(); i++) {
@@ -1038,7 +1011,8 @@
 
                 for (PackageRollbackInfo info : data.info.getPackages()) {
                     if (info.getPackageName().equals(packageName)) {
-                        mAppDataRollbackHelper.snapshotAppData(data.info.getRollbackId(), info);
+                        mAppDataRollbackHelper.snapshotAppData(
+                                data.info.getRollbackId(), info, userIds);
                         saveRollbackData(data);
                         break;
                     }
@@ -1049,8 +1023,8 @@
             for (NewRollback rollback : mNewRollbacks) {
                 info = getPackageRollbackInfo(rollback.data, packageName);
                 if (info != null) {
-                    mAppDataRollbackHelper.snapshotAppData(rollback.data.info.getRollbackId(),
-                            info);
+                    mAppDataRollbackHelper.snapshotAppData(
+                            rollback.data.info.getRollbackId(), info, userIds);
                     saveRollbackData(rollback.data);
                 }
             }
@@ -1114,8 +1088,7 @@
             }
 
             if (!session.isMultiPackage()) {
-                if (!enableRollbackForPackageSession(newRollback.data, session,
-                            new int[0])) {
+                if (!enableRollbackForPackageSession(newRollback.data, session)) {
                     Slog.e(TAG, "Unable to enable rollback for session: " + sessionId);
                     result.offer(false);
                     return;
@@ -1129,8 +1102,7 @@
                         result.offer(false);
                         return;
                     }
-                    if (!enableRollbackForPackageSession(newRollback.data, childSession,
-                                new int[0])) {
+                    if (!enableRollbackForPackageSession(newRollback.data, childSession)) {
                         Slog.e(TAG, "Unable to enable rollback for session: " + sessionId);
                         result.offer(false);
                         return;
@@ -1407,9 +1379,9 @@
 
     private void deleteRollback(RollbackData rollbackData) {
         for (PackageRollbackInfo info : rollbackData.info.getPackages()) {
-            IntArray installedUsers = info.getInstalledUsers();
-            for (int i = 0; i < installedUsers.size(); i++) {
-                int userId = installedUsers.get(i);
+            IntArray snapshottedUsers = info.getSnapshottedUsers();
+            for (int i = 0; i < snapshottedUsers.size(); i++) {
+                int userId = snapshottedUsers.get(i);
                 mAppDataRollbackHelper.destroyAppDataSnapshot(rollbackData.info.getRollbackId(),
                         info, userId);
             }
diff --git a/services/core/java/com/android/server/rollback/RollbackStore.java b/services/core/java/com/android/server/rollback/RollbackStore.java
index 1c36dc7..d106825 100644
--- a/services/core/java/com/android/server/rollback/RollbackStore.java
+++ b/services/core/java/com/android/server/rollback/RollbackStore.java
@@ -319,13 +319,14 @@
 
         IntArray pendingBackups = info.getPendingBackups();
         List<RestoreInfo> pendingRestores = info.getPendingRestores();
-        IntArray installedUsers = info.getInstalledUsers();
+        IntArray snapshottedUsers = info.getSnapshottedUsers();
         json.put("pendingBackups", convertToJsonArray(pendingBackups));
         json.put("pendingRestores", convertToJsonArray(pendingRestores));
 
         json.put("isApex", info.isApex());
 
-        json.put("installedUsers", convertToJsonArray(installedUsers));
+        // Field is named 'installedUsers' for legacy reasons.
+        json.put("installedUsers", convertToJsonArray(snapshottedUsers));
         json.put("ceSnapshotInodes", ceSnapshotInodesToJson(info.getCeSnapshotInodes()));
 
         return json;
@@ -345,12 +346,13 @@
 
         final boolean isApex = json.getBoolean("isApex");
 
-        final IntArray installedUsers = convertToIntArray(json.getJSONArray("installedUsers"));
+        // Field is named 'installedUsers' for legacy reasons.
+        final IntArray snapshottedUsers = convertToIntArray(json.getJSONArray("installedUsers"));
         final SparseLongArray ceSnapshotInodes = ceSnapshotInodesFromJson(
                 json.getJSONArray("ceSnapshotInodes"));
 
         return new PackageRollbackInfo(versionRolledBackFrom, versionRolledBackTo,
-                pendingBackups, pendingRestores, isApex, installedUsers, ceSnapshotInodes);
+                pendingBackups, pendingRestores, isApex, snapshottedUsers, ceSnapshotInodes);
     }
 
     private static JSONArray versionedPackagesToJson(List<VersionedPackage> packages)
diff --git a/services/tests/servicestests/src/com/android/server/rollback/AppDataRollbackHelperTest.java b/services/tests/servicestests/src/com/android/server/rollback/AppDataRollbackHelperTest.java
index 7986055..1f0ef6e 100644
--- a/services/tests/servicestests/src/com/android/server/rollback/AppDataRollbackHelperTest.java
+++ b/services/tests/servicestests/src/com/android/server/rollback/AppDataRollbackHelperTest.java
@@ -16,6 +16,7 @@
 
 package com.android.server.rollback;
 
+import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
@@ -58,8 +59,8 @@
         // All users are unlocked so we should snapshot data for them.
         doReturn(true).when(helper).isUserCredentialLocked(eq(10));
         doReturn(true).when(helper).isUserCredentialLocked(eq(11));
-        PackageRollbackInfo info = createPackageRollbackInfo("com.foo.bar", new int[]{10, 11});
-        helper.snapshotAppData(5, info);
+        PackageRollbackInfo info = createPackageRollbackInfo("com.foo.bar");
+        helper.snapshotAppData(5, info, new int[]{10, 11});
 
         assertEquals(2, info.getPendingBackups().size());
         assertEquals(10, info.getPendingBackups().get(0));
@@ -79,8 +80,8 @@
         doReturn(true).when(helper).isUserCredentialLocked(eq(11));
         when(installer.snapshotAppData(anyString(), anyInt(), anyInt(), anyInt())).thenReturn(239L);
 
-        PackageRollbackInfo info2 = createPackageRollbackInfo("com.foo.bar", new int[]{10, 11});
-        helper.snapshotAppData(7, info2);
+        PackageRollbackInfo info2 = createPackageRollbackInfo("com.foo.bar");
+        helper.snapshotAppData(7, info2, new int[]{10, 11});
         assertEquals(1, info2.getPendingBackups().size());
         assertEquals(11, info2.getPendingBackups().get(0));
 
@@ -278,4 +279,15 @@
 
         inOrder.verifyNoMoreInteractions();
     }
+
+    @Test
+    public void snapshotAddDataSavesSnapshottedUsersToInfo() {
+        Installer installer = mock(Installer.class);
+        AppDataRollbackHelper helper = new AppDataRollbackHelper(installer);
+
+        PackageRollbackInfo info = createPackageRollbackInfo("com.foo.bar");
+        helper.snapshotAppData(5, info, new int[]{10, 11});
+
+        assertArrayEquals(info.getSnapshottedUsers().toArray(), new int[]{10, 11});
+    }
 }