Use userId and package installer name when committing rollbacks

Currently when committing a rollback, the system user is used
for the package context for the downgrade session, even when
the package was not installed as this user. This leads to a
package being installed to the system user in the cases where
it is only installed to a secondary user.

The installer package
name is also currently retrieved from package manager (and does
not work for apex), so storing this information when enabling
the rollback will be useful.

This change should be compatible with previous rollbacks stored
on the device.

Bug: 133210624
Test: atest RollbackTest
Test: atest SecondaryUserRollbackTest
Test: Install test APK v1 and v2 with rollback enabled for a secondary
user, rollback and verify the package is not installed for
the system user

Change-Id: I16346120cf9ba7b04254ab5e8fa5090be7d562ac

7 files changed, 124 insertions, 46 deletions

diff --git a/services/core/java/com/android/server/rollback/Rollback.java b/services/core/java/com/android/server/rollback/Rollback.java
index 1cf07cb47f2f..8b79c3ff66c6 100644
--- a/services/core/java/com/android/server/rollback/Rollback.java
+++ b/services/core/java/com/android/server/rollback/Rollback.java
@@ -21,6 +21,7 @@ import static com.android.server.rollback.RollbackManagerServiceImpl.sendFailure
 import android.Manifest;
 import android.annotation.IntDef;
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentSender;
@@ -33,7 +34,9 @@ import android.content.rollback.RollbackInfo;
 import android.content.rollback.RollbackManager;
 import android.os.Binder;
 import android.os.ParcelFileDescriptor;
+import android.os.UserHandle;
 import android.os.UserManager;
+import android.text.TextUtils;
 import android.util.IntArray;
 import android.util.Slog;
 import android.util.SparseLongArray;
@@ -142,18 +145,37 @@ class Rollback {
     private final Object mLock = new Object();
 
     /**
+     * The user that performed the install with rollback enabled.
+     */
+    public final int mUserId;
+
+    /**
+     * The installer package name from the install session that enabled the rollback. May be null if
+     * that session did not set this value.
+     *
+     * If this is an empty string then the installer package name will be resolved by
+     * PackageManager.
+     */
+    @Nullable public final String mInstallerPackageName;
+
+    /**
      * Constructs a new, empty Rollback instance.
      *
      * @param rollbackId the id of the rollback.
      * @param backupDir the directory where the rollback data is stored.
      * @param stagedSessionId the session id if this is a staged rollback, -1 otherwise.
+     * @param userId the user that performed the install with rollback enabled.
+     * @param installerPackageName the installer package name from the original install session.
      */
-    Rollback(int rollbackId, File backupDir, int stagedSessionId) {
+    Rollback(int rollbackId, File backupDir, int stagedSessionId, int userId,
+            String installerPackageName) {
         this.info = new RollbackInfo(rollbackId,
                 /* packages */ new ArrayList<>(),
                 /* isStaged */ stagedSessionId != -1,
                 /* causePackages */ new ArrayList<>(),
                 /* committedSessionId */ -1);
+        mUserId = userId;
+        mInstallerPackageName = installerPackageName;
         mBackupDir = backupDir;
         mStagedSessionId = stagedSessionId;
         mState = ROLLBACK_STATE_ENABLING;
@@ -164,8 +186,11 @@ class Rollback {
      * Constructs a pre-populated Rollback instance.
      */
     Rollback(RollbackInfo info, File backupDir, Instant timestamp, int stagedSessionId,
-            @RollbackState int state, int apkSessionId, boolean restoreUserDataInProgress) {
+            @RollbackState int state, int apkSessionId, boolean restoreUserDataInProgress,
+            int userId, String installerPackageName) {
         this.info = info;
+        mUserId = userId;
+        mInstallerPackageName = installerPackageName;
         mBackupDir = backupDir;
         mTimestamp = timestamp;
         mStagedSessionId = stagedSessionId;
@@ -216,6 +241,21 @@ class Rollback {
     }
 
     /**
+     * Returns the ID of the user that performed the install with rollback enabled.
+     */
+    int getUserId() {
+        return mUserId;
+    }
+
+    /**
+     * Returns the installer package name from the install session that enabled the rollback. In the
+     * case that this is called on a rollback from an older version, returns the empty string.
+     */
+    @Nullable String getInstallerPackageName() {
+        return mInstallerPackageName;
+    }
+
+    /**
      * Returns true if the rollback is in the ENABLING state.
      */
     boolean isEnabling() {
@@ -360,7 +400,8 @@ class Rollback {
             // Get a context to use to install the downgraded version of the package.
             Context pkgContext;
             try {
-                pkgContext = context.createPackageContext(callerPackageName, 0);
+                pkgContext = context.createPackageContextAsUser(callerPackageName, 0,
+                        UserHandle.of(mUserId));
             } catch (PackageManager.NameNotFoundException e) {
                 sendFailure(context, statusReceiver, RollbackManager.STATUS_FAILURE,
                         "Invalid callerPackageName");
@@ -385,15 +426,13 @@ class Rollback {
                 for (PackageRollbackInfo pkgRollbackInfo : info.getPackages()) {
                     PackageInstaller.SessionParams params = new PackageInstaller.SessionParams(
                             PackageInstaller.SessionParams.MODE_FULL_INSTALL);
-                    // TODO: We can't get the installerPackageName for apex
-                    // (b/123920130). Is it okay to ignore the installer package
-                    // for apex?
-                    if (!pkgRollbackInfo.isApex()) {
-                        String installerPackageName =
-                                pm.getInstallerPackageName(pkgRollbackInfo.getPackageName());
-                        if (installerPackageName != null) {
-                            params.setInstallerPackageName(installerPackageName);
-                        }
+                    String installerPackageName = mInstallerPackageName;
+                    if (TextUtils.isEmpty(mInstallerPackageName)) {
+                        installerPackageName = pm.getInstallerPackageName(
+                                pkgRollbackInfo.getPackageName());
+                    }
+                    if (installerPackageName != null) {
+                        params.setInstallerPackageName(installerPackageName);
                     }
                     params.setRequestDowngrade(true);
                     params.setRequiredInstalledVersionCode(
diff --git a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
index cd44f64ada4f..ef4c12e8145a 100644
--- a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
+++ b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
@@ -1248,13 +1248,22 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
     @GuardedBy("mLock")
     private NewRollback createNewRollbackLocked(PackageInstaller.SessionInfo parentSession) {
         int rollbackId = allocateRollbackIdLocked();
+        final int userId;
+        if (parentSession.getUser() == UserHandle.ALL) {
+            userId = UserHandle.USER_SYSTEM;
+        } else {
+            userId = parentSession.getUser().getIdentifier();
+        }
+        String installerPackageName = parentSession.getInstallerPackageName();
         final Rollback rollback;
         int parentSessionId = parentSession.getSessionId();
 
         if (parentSession.isStaged()) {
-            rollback = mRollbackStore.createStagedRollback(rollbackId, parentSessionId);
+            rollback = mRollbackStore.createStagedRollback(rollbackId, parentSessionId, userId,
+                    installerPackageName);
         } else {
-            rollback = mRollbackStore.createNonStagedRollback(rollbackId);
+            rollback = mRollbackStore.createNonStagedRollback(rollbackId, userId,
+                    installerPackageName);
         }
 
         int[] packageSessionIds;
diff --git a/services/core/java/com/android/server/rollback/RollbackStore.java b/services/core/java/com/android/server/rollback/RollbackStore.java
index 425da374a5a8..550c75442892 100644
--- a/services/core/java/com/android/server/rollback/RollbackStore.java
+++ b/services/core/java/com/android/server/rollback/RollbackStore.java
@@ -16,6 +16,8 @@
 
 package com.android.server.rollback;
 
+import static android.os.UserHandle.USER_SYSTEM;
+
 import static com.android.server.rollback.Rollback.rollbackStateFromString;
 
 import android.annotation.NonNull;
@@ -196,18 +198,19 @@ class RollbackStore {
      * Creates a new Rollback instance for a non-staged rollback with
      * backupDir assigned.
      */
-    Rollback createNonStagedRollback(int rollbackId) {
+    Rollback createNonStagedRollback(int rollbackId, int userId, String installerPackageName) {
         File backupDir = new File(mRollbackDataDir, Integer.toString(rollbackId));
-        return new Rollback(rollbackId, backupDir, -1);
+        return new Rollback(rollbackId, backupDir, -1, userId, installerPackageName);
     }
 
     /**
      * Creates a new Rollback instance for a staged rollback with
      * backupDir assigned.
      */
-    Rollback createStagedRollback(int rollbackId, int stagedSessionId) {
+    Rollback createStagedRollback(int rollbackId, int stagedSessionId, int userId,
+            String installerPackageName) {
         File backupDir = new File(mRollbackDataDir, Integer.toString(rollbackId));
-        return new Rollback(rollbackId, backupDir, stagedSessionId);
+        return new Rollback(rollbackId, backupDir, stagedSessionId, userId, installerPackageName);
     }
 
     /**
@@ -263,6 +266,8 @@ class RollbackStore {
             dataJson.put("state", rollback.getStateAsString());
             dataJson.put("apkSessionId", rollback.getApkSessionId());
             dataJson.put("restoreUserDataInProgress", rollback.isRestoreUserDataInProgress());
+            dataJson.put("userId", rollback.getUserId());
+            dataJson.putOpt("installerPackageName", rollback.getInstallerPackageName());
 
             PrintWriter pw = new PrintWriter(new File(rollback.getBackupDir(), "rollback.json"));
             pw.println(dataJson.toString());
@@ -305,7 +310,9 @@ class RollbackStore {
                 dataJson.getInt("stagedSessionId"),
                 rollbackStateFromString(dataJson.getString("state")),
                 dataJson.getInt("apkSessionId"),
-                dataJson.getBoolean("restoreUserDataInProgress"));
+                dataJson.getBoolean("restoreUserDataInProgress"),
+                dataJson.optInt("userId", USER_SYSTEM),
+                dataJson.optString("installerPackageName", ""));
     }
 
     private static JSONObject toJson(VersionedPackage pkg) throws JSONException {
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 0c5451f4ec31..a83d94001cf8 100644
--- a/services/tests/servicestests/src/com/android/server/rollback/AppDataRollbackHelperTest.java
+++ b/services/tests/servicestests/src/com/android/server/rollback/AppDataRollbackHelperTest.java
@@ -106,6 +106,11 @@ public class AppDataRollbackHelperTest {
         return createPackageRollbackInfo(packageName, new int[] {});
     }
 
+    private static Rollback createRollbackForId(int rollbackId) {
+        return new Rollback(rollbackId, new File("/does/not/exist"), -1,
+                0, "com.xyz");
+    }
+
     @Test
     public void testRestoreAppDataSnapshot_pendingBackupForUser() throws Exception {
         Installer installer = mock(Installer.class);
@@ -232,18 +237,16 @@ public class AppDataRollbackHelperTest {
         wasRecentlyRestored.getPendingRestores().add(
                 new RestoreInfo(73 /* userId */, 239 /* appId*/, "seInfo"));
 
-        Rollback dataWithPendingBackup = new Rollback(101, new File("/does/not/exist"), -1);
+        Rollback dataWithPendingBackup = createRollbackForId(101);
         dataWithPendingBackup.info.getPackages().add(pendingBackup);
 
-        Rollback dataWithRecentRestore = new Rollback(17239, new File("/does/not/exist"),
-                -1);
+        Rollback dataWithRecentRestore = createRollbackForId(17239);
         dataWithRecentRestore.info.getPackages().add(wasRecentlyRestored);
 
-        Rollback dataForDifferentUser = new Rollback(17239, new File("/does/not/exist"),
-                -1);
+        Rollback dataForDifferentUser = createRollbackForId(17239);
         dataForDifferentUser.info.getPackages().add(ignoredInfo);
 
-        Rollback dataForRestore = new Rollback(17239, new File("/does/not/exist"), -1);
+        Rollback dataForRestore = createRollbackForId(17239);
         dataForRestore.info.getPackages().add(pendingRestore);
         dataForRestore.info.getPackages().add(wasRecentlyRestored);
 
diff --git a/services/tests/servicestests/src/com/android/server/rollback/RollbackStoreTest.java b/services/tests/servicestests/src/com/android/server/rollback/RollbackStoreTest.java
index ee3b15a97f26..aec489cd1ab4 100644
--- a/services/tests/servicestests/src/com/android/server/rollback/RollbackStoreTest.java
+++ b/services/tests/servicestests/src/com/android/server/rollback/RollbackStoreTest.java
@@ -43,6 +43,8 @@ import java.util.Objects;
 public class RollbackStoreTest {
 
     private static final int ID = 123;
+    private static final int USER = 0;
+    private static final String INSTALLER = "some.installer";
 
     private static final Correspondence<VersionedPackage, VersionedPackage> VER_PKG_CORR =
             new Correspondence<VersionedPackage, VersionedPackage>() {
@@ -97,7 +99,8 @@ public class RollbackStoreTest {
             + "'longVersionCode':23},{'packageName':'something','longVersionCode':999}],"
             + "'committedSessionId':45654465},'timestamp':'2019-10-01T12:29:08.855Z',"
             + "'stagedSessionId':-1,'state':'enabling','apkSessionId':-1,"
-            + "'restoreUserDataInProgress':true}";
+            + "'restoreUserDataInProgress':true, 'userId':0,"
+            + "'installerPackageName':'some.installer'}";
 
     @Rule
     public TemporaryFolder mFolder = new TemporaryFolder();
@@ -115,7 +118,7 @@ public class RollbackStoreTest {
 
     @Test
     public void createNonStaged() {
-        Rollback rollback = mRollbackStore.createNonStagedRollback(ID);
+        Rollback rollback = mRollbackStore.createNonStagedRollback(ID, USER, INSTALLER);
 
         assertThat(rollback.getBackupDir().getAbsolutePath())
                 .isEqualTo(mFolder.getRoot().getAbsolutePath() + "/" + ID);
@@ -128,7 +131,7 @@ public class RollbackStoreTest {
 
     @Test
     public void createStaged() {
-        Rollback rollback = mRollbackStore.createStagedRollback(ID, 897);
+        Rollback rollback = mRollbackStore.createStagedRollback(ID, 897, USER, INSTALLER);
 
         assertThat(rollback.getBackupDir().getAbsolutePath())
                 .isEqualTo(mFolder.getRoot().getAbsolutePath() + "/" + ID);
@@ -143,7 +146,7 @@ public class RollbackStoreTest {
 
     @Test
     public void saveAndLoadRollback() {
-        Rollback origRb = mRollbackStore.createNonStagedRollback(ID);
+        Rollback origRb = mRollbackStore.createNonStagedRollback(ID, USER, INSTALLER);
 
         origRb.setRestoreUserDataInProgress(true);
         origRb.info.getCausePackages().add(new VersionedPackage("com.made.up", 2));
@@ -193,7 +196,7 @@ public class RollbackStoreTest {
 
     @Test
     public void loadFromJson() throws Exception {
-        Rollback expectedRb = mRollbackStore.createNonStagedRollback(ID);
+        Rollback expectedRb = mRollbackStore.createNonStagedRollback(ID, USER, INSTALLER);
 
         expectedRb.setTimestamp(Instant.parse("2019-10-01T12:29:08.855Z"));
         expectedRb.setRestoreUserDataInProgress(true);
@@ -242,7 +245,7 @@ public class RollbackStoreTest {
 
     @Test
     public void saveAndDelete() {
-        Rollback rollback = mRollbackStore.createNonStagedRollback(ID);
+        Rollback rollback = mRollbackStore.createNonStagedRollback(ID, USER, INSTALLER);
 
         RollbackStore.saveRollback(rollback);
 
@@ -287,6 +290,9 @@ public class RollbackStoreTest {
             assertPackageRollbacksAreEquivalent(
                     b.info.getPackages().get(i), a.info.getPackages().get(i));
         }
+
+        assertThat(a.getUserId()).isEqualTo(b.getUserId());
+        assertThat(a.getInstallerPackageName()).isEqualTo(b.getInstallerPackageName());
     }
 
     private void assertPackageRollbacksAreEquivalent(PackageRollbackInfo b, PackageRollbackInfo a) {
diff --git a/services/tests/servicestests/src/com/android/server/rollback/RollbackUnitTest.java b/services/tests/servicestests/src/com/android/server/rollback/RollbackUnitTest.java
index 151b6e2890bb..e368d634b968 100644
--- a/services/tests/servicestests/src/com/android/server/rollback/RollbackUnitTest.java
+++ b/services/tests/servicestests/src/com/android/server/rollback/RollbackUnitTest.java
@@ -52,6 +52,8 @@ public class RollbackUnitTest {
     private static final String PKG_2 = "test.testpackage.pkg2";
     private static final String PKG_3 = "com.blah.hello.three";
     private static final String PKG_4 = "com.something.4pack";
+    private static final int USER = 0;
+    private static final String INSTALLER = "some.installer";
 
     @Mock private AppDataRollbackHelper mMockDataHelper;
 
@@ -66,7 +68,7 @@ public class RollbackUnitTest {
         int sessionId = 567;
         File file = new File("/test/testing");
 
-        Rollback rollback = new Rollback(rollbackId, file, sessionId);
+        Rollback rollback = new Rollback(rollbackId, file, sessionId, USER, INSTALLER);
 
         assertThat(rollback.isEnabling()).isTrue();
         assertThat(rollback.getBackupDir().getAbsolutePath()).isEqualTo("/test/testing");
@@ -79,7 +81,7 @@ public class RollbackUnitTest {
         int rollbackId = 123;
         File file = new File("/test/testing");
 
-        Rollback rollback = new Rollback(rollbackId, file, -1);
+        Rollback rollback = new Rollback(rollbackId, file, -1, USER, INSTALLER);
 
         assertThat(rollback.isEnabling()).isTrue();
         assertThat(rollback.getBackupDir().getAbsolutePath()).isEqualTo("/test/testing");
@@ -88,7 +90,8 @@ public class RollbackUnitTest {
 
     @Test
     public void rollbackMadeAvailable() {
-        Rollback rollback = new Rollback(123, new File("/test/testing"), -1);
+        Rollback rollback = new Rollback(123, new File("/test/testing"), -1, USER,
+                INSTALLER);
 
         assertThat(rollback.isEnabling()).isTrue();
         assertThat(rollback.isAvailable()).isFalse();
@@ -106,7 +109,7 @@ public class RollbackUnitTest {
 
     @Test
     public void deletedRollbackCannotBeMadeAvailable() {
-        Rollback rollback = new Rollback(123, new File("/test/testing"), -1);
+        Rollback rollback = new Rollback(123, new File("/test/testing"), -1, USER, INSTALLER);
 
         rollback.delete(mMockDataHelper);
 
@@ -120,7 +123,7 @@ public class RollbackUnitTest {
 
     @Test
     public void getPackageNamesAllAndJustApex() {
-        Rollback rollback = new Rollback(123, new File("/test/testing"), -1);
+        Rollback rollback = new Rollback(123, new File("/test/testing"), -1, USER, INSTALLER);
         PackageRollbackInfo pkgInfo1 = newPkgInfoFor(PKG_1, 12, 10, false);
         PackageRollbackInfo pkgInfo2 = newPkgInfoFor(PKG_2, 18, 11, true);
         PackageRollbackInfo pkgInfo3 = newPkgInfoFor(PKG_3, 19, 1, false);
@@ -134,7 +137,7 @@ public class RollbackUnitTest {
 
     @Test
     public void includesPackagesAfterEnable() {
-        Rollback rollback = new Rollback(123, new File("/test/testing"), -1);
+        Rollback rollback = new Rollback(123, new File("/test/testing"), -1, USER, INSTALLER);
         PackageRollbackInfo pkgInfo1 = newPkgInfoFor(PKG_1, 12, 10, false);
         PackageRollbackInfo pkgInfo2 = newPkgInfoFor(PKG_2, 18, 12, true);
         PackageRollbackInfo pkgInfo3 = newPkgInfoFor(PKG_3, 157, 156, false);
@@ -162,7 +165,7 @@ public class RollbackUnitTest {
 
     @Test
     public void snapshotWhenEnabling() {
-        Rollback rollback = new Rollback(123, new File("/test/testing"), -1);
+        Rollback rollback = new Rollback(123, new File("/test/testing"), -1, USER, INSTALLER);
         PackageRollbackInfo pkgInfo1 = newPkgInfoFor(PKG_1, 12, 10, false);
         PackageRollbackInfo pkgInfo2 = newPkgInfoFor(PKG_2, 18, 12, true);
         rollback.info.getPackages().addAll(Arrays.asList(pkgInfo1, pkgInfo2));
@@ -180,7 +183,7 @@ public class RollbackUnitTest {
 
     @Test
     public void snapshotWhenAvailable() {
-        Rollback rollback = new Rollback(123, new File("/test/testing"), -1);
+        Rollback rollback = new Rollback(123, new File("/test/testing"), -1, USER, INSTALLER);
         PackageRollbackInfo pkgInfo1 = newPkgInfoFor(PKG_1, 12, 10, false);
         PackageRollbackInfo pkgInfo2 = newPkgInfoFor(PKG_2, 18, 12, true);
         rollback.info.getPackages().addAll(Arrays.asList(pkgInfo1, pkgInfo2));
@@ -201,7 +204,7 @@ public class RollbackUnitTest {
 
     @Test
     public void snapshotWhenDeleted() {
-        Rollback rollback = new Rollback(123, new File("/test/testing"), -1);
+        Rollback rollback = new Rollback(123, new File("/test/testing"), -1, USER, INSTALLER);
         PackageRollbackInfo pkgInfo1 = newPkgInfoFor(PKG_1, 12, 10, false);
         PackageRollbackInfo pkgInfo2 = newPkgInfoFor(PKG_2, 18, 12, true);
         rollback.info.getPackages().addAll(Arrays.asList(pkgInfo1, pkgInfo2));
@@ -222,7 +225,7 @@ public class RollbackUnitTest {
 
     @Test
     public void snapshotThenDelete() {
-        Rollback rollback = new Rollback(123, new File("/test/testing"), -1);
+        Rollback rollback = new Rollback(123, new File("/test/testing"), -1, USER, INSTALLER);
         PackageRollbackInfo pkgInfo1 = newPkgInfoFor(PKG_1, 12, 10, false);
         PackageRollbackInfo pkgInfo2 = newPkgInfoFor(PKG_2, 18, 12, true);
         rollback.info.getPackages().addAll(Arrays.asList(pkgInfo1, pkgInfo2));
@@ -242,7 +245,7 @@ public class RollbackUnitTest {
 
     @Test
     public void restoreUserDataDoesNothingIfNotInProgress() {
-        Rollback rollback = new Rollback(123, new File("/test/testing"), -1);
+        Rollback rollback = new Rollback(123, new File("/test/testing"), -1, USER, INSTALLER);
         PackageRollbackInfo pkgInfo1 = newPkgInfoFor(PKG_1, 12, 10, false);
         PackageRollbackInfo pkgInfo2 = newPkgInfoFor(PKG_2, 18, 12, true);
         rollback.info.getPackages().addAll(Arrays.asList(pkgInfo1, pkgInfo2));
@@ -257,7 +260,7 @@ public class RollbackUnitTest {
 
     @Test
     public void restoreUserDataDoesNothingIfPackageNotFound() {
-        Rollback rollback = new Rollback(123, new File("/test/testing"), -1);
+        Rollback rollback = new Rollback(123, new File("/test/testing"), -1, USER, INSTALLER);
         PackageRollbackInfo pkgInfo1 = newPkgInfoFor(PKG_1, 12, 10, false);
         PackageRollbackInfo pkgInfo2 = newPkgInfoFor(PKG_2, 18, 12, true);
         rollback.info.getPackages().addAll(Arrays.asList(pkgInfo1, pkgInfo2));
@@ -273,7 +276,7 @@ public class RollbackUnitTest {
 
     @Test
     public void restoreUserDataRestoresIfInProgressAndPackageFound() {
-        Rollback rollback = new Rollback(123, new File("/test/testing"), -1);
+        Rollback rollback = new Rollback(123, new File("/test/testing"), -1, USER, INSTALLER);
         PackageRollbackInfo pkgInfo1 = newPkgInfoFor(PKG_1, 12, 10, false);
         PackageRollbackInfo pkgInfo2 = newPkgInfoFor(PKG_2, 18, 12, true);
         rollback.info.getPackages().addAll(Arrays.asList(pkgInfo1, pkgInfo2));
diff --git a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java
index ed8f272874e1..599ee56afdb3 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java
@@ -33,6 +33,7 @@ import android.content.Intent;
 import android.content.IntentFilter;
 import android.content.rollback.RollbackInfo;
 import android.content.rollback.RollbackManager;
+import android.os.UserManager;
 import android.provider.DeviceConfig;
 import android.util.Log;
 
@@ -52,7 +53,9 @@ import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
 import java.util.Collections;
+import java.util.List;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
 
 /**
  * Test system Rollback APIs.
@@ -96,7 +99,10 @@ public class RollbackTest {
                     Manifest.permission.INSTALL_PACKAGES,
                     Manifest.permission.DELETE_PACKAGES,
                     Manifest.permission.TEST_MANAGE_ROLLBACKS,
-                    Manifest.permission.MANAGE_ROLLBACKS);
+                    Manifest.permission.MANAGE_ROLLBACKS,
+                    Manifest.permission.MANAGE_USERS,
+                    Manifest.permission.CREATE_USERS,
+                    Manifest.permission.INTERACT_ACROSS_USERS);
 
             // Register a broadcast receiver for notification when the
             // rollback has been committed.
@@ -106,7 +112,6 @@ public class RollbackTest {
             // Uninstall TestApp.A
             Uninstall.packages(TestApp.A);
             assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(-1);
-
             // TODO: There is currently a race condition between when the app is
             // uninstalled and when rollback manager deletes the rollback. Fix it
             // so that's not the case!
@@ -149,6 +154,12 @@ public class RollbackTest {
             RollbackUtils.rollback(available.getRollbackId());
             assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(1);
 
+            UserManager um = (UserManager) context.getSystemService(context.USER_SERVICE);
+            List<Integer> userIds = um.getUsers(true)
+                    .stream().map(user -> user.id).collect(Collectors.toList());
+            assertThat(InstallUtils.isOnlyInstalledForUser(TestApp.A,
+                    context.getUserId(), userIds)).isTrue();
+
             // Verify we received a broadcast for the rollback.
             // TODO: Race condition between the timeout and when the broadcast is
             // received could lead to test flakiness.