Only allow rollback to be enabled on modules.

Only allow rollback to be enabled on the modules included in a mainline
update. We don't want to support rollbacks for all apks in general yet.

Enforce that only installers granted the MANAGE_ROLLBACKS permission can
enable rollback for a package.

Introduce a new TEST_MANAGE_ROLLBACKS permission that can be used to
enable rollback on packages that are not modules. This allows us to
continue testing rollbacks, given we can't do a mainline update as part
of the rollback tests.

Test: atest RollbackTest, with new tests for permissions added.
Bug: 128277794
Change-Id: I29ab9a750a1283592b8a855322ece516e42260ca
diff --git a/api/system-current.txt b/api/system-current.txt
index 61a7f58..b9301ba 100644
--- a/api/system-current.txt
+++ b/api/system-current.txt
@@ -186,6 +186,7 @@
     field public static final String SUBSTITUTE_NOTIFICATION_APP_NAME = "android.permission.SUBSTITUTE_NOTIFICATION_APP_NAME";
     field public static final String SUBSTITUTE_SHARE_TARGET_APP_NAME_AND_ICON = "android.permission.SUBSTITUTE_SHARE_TARGET_APP_NAME_AND_ICON";
     field public static final String SUSPEND_APPS = "android.permission.SUSPEND_APPS";
+    field public static final String TEST_MANAGE_ROLLBACKS = "android.permission.TEST_MANAGE_ROLLBACKS";
     field public static final String TETHER_PRIVILEGED = "android.permission.TETHER_PRIVILEGED";
     field public static final String TV_INPUT_HARDWARE = "android.permission.TV_INPUT_HARDWARE";
     field public static final String TV_VIRTUAL_REMOTE_CONTROLLER = "android.permission.TV_VIRTUAL_REMOTE_CONTROLLER";
@@ -1835,9 +1836,9 @@
   }
 
   public final class RollbackManager {
-    method @RequiresPermission(android.Manifest.permission.MANAGE_ROLLBACKS) public void commitRollback(int, @NonNull java.util.List<android.content.pm.VersionedPackage>, @NonNull android.content.IntentSender);
-    method @RequiresPermission(android.Manifest.permission.MANAGE_ROLLBACKS) @NonNull public java.util.List<android.content.rollback.RollbackInfo> getAvailableRollbacks();
-    method @RequiresPermission(android.Manifest.permission.MANAGE_ROLLBACKS) @NonNull public java.util.List<android.content.rollback.RollbackInfo> getRecentlyCommittedRollbacks();
+    method @RequiresPermission(anyOf={android.Manifest.permission.MANAGE_ROLLBACKS, android.Manifest.permission.TEST_MANAGE_ROLLBACKS}) public void commitRollback(int, @NonNull java.util.List<android.content.pm.VersionedPackage>, @NonNull android.content.IntentSender);
+    method @RequiresPermission(anyOf={android.Manifest.permission.MANAGE_ROLLBACKS, android.Manifest.permission.TEST_MANAGE_ROLLBACKS}) @NonNull public java.util.List<android.content.rollback.RollbackInfo> getAvailableRollbacks();
+    method @RequiresPermission(anyOf={android.Manifest.permission.MANAGE_ROLLBACKS, android.Manifest.permission.TEST_MANAGE_ROLLBACKS}) @NonNull public java.util.List<android.content.rollback.RollbackInfo> getRecentlyCommittedRollbacks();
     field public static final String EXTRA_STATUS = "android.content.rollback.extra.STATUS";
     field public static final String EXTRA_STATUS_MESSAGE = "android.content.rollback.extra.STATUS_MESSAGE";
     field public static final int STATUS_FAILURE = 1; // 0x1
diff --git a/api/test-current.txt b/api/test-current.txt
index 4ccfa1c..7ff1c13 100644
--- a/api/test-current.txt
+++ b/api/test-current.txt
@@ -14,6 +14,7 @@
     field public static final String MANAGE_ROLLBACKS = "android.permission.MANAGE_ROLLBACKS";
     field public static final String READ_CELL_BROADCASTS = "android.permission.READ_CELL_BROADCASTS";
     field public static final String REMOVE_TASKS = "android.permission.REMOVE_TASKS";
+    field public static final String TEST_MANAGE_ROLLBACKS = "android.permission.TEST_MANAGE_ROLLBACKS";
     field public static final String WRITE_DEVICE_CONFIG = "android.permission.WRITE_DEVICE_CONFIG";
     field public static final String WRITE_MEDIA_STORAGE = "android.permission.WRITE_MEDIA_STORAGE";
     field public static final String WRITE_OBB = "android.permission.WRITE_OBB";
@@ -720,11 +721,11 @@
   }
 
   public final class RollbackManager {
-    method @RequiresPermission(android.Manifest.permission.MANAGE_ROLLBACKS) public void commitRollback(int, @NonNull java.util.List<android.content.pm.VersionedPackage>, @NonNull android.content.IntentSender);
-    method @RequiresPermission(android.Manifest.permission.MANAGE_ROLLBACKS) public void expireRollbackForPackage(@NonNull String);
-    method @RequiresPermission(android.Manifest.permission.MANAGE_ROLLBACKS) @NonNull public java.util.List<android.content.rollback.RollbackInfo> getAvailableRollbacks();
-    method @RequiresPermission(android.Manifest.permission.MANAGE_ROLLBACKS) @NonNull public java.util.List<android.content.rollback.RollbackInfo> getRecentlyCommittedRollbacks();
-    method @RequiresPermission(android.Manifest.permission.MANAGE_ROLLBACKS) public void reloadPersistedData();
+    method @RequiresPermission(anyOf={android.Manifest.permission.MANAGE_ROLLBACKS, android.Manifest.permission.TEST_MANAGE_ROLLBACKS}) public void commitRollback(int, @NonNull java.util.List<android.content.pm.VersionedPackage>, @NonNull android.content.IntentSender);
+    method @RequiresPermission(android.Manifest.permission.TEST_MANAGE_ROLLBACKS) public void expireRollbackForPackage(@NonNull String);
+    method @RequiresPermission(anyOf={android.Manifest.permission.MANAGE_ROLLBACKS, android.Manifest.permission.TEST_MANAGE_ROLLBACKS}) @NonNull public java.util.List<android.content.rollback.RollbackInfo> getAvailableRollbacks();
+    method @RequiresPermission(anyOf={android.Manifest.permission.MANAGE_ROLLBACKS, android.Manifest.permission.TEST_MANAGE_ROLLBACKS}) @NonNull public java.util.List<android.content.rollback.RollbackInfo> getRecentlyCommittedRollbacks();
+    method @RequiresPermission(android.Manifest.permission.TEST_MANAGE_ROLLBACKS) public void reloadPersistedData();
     field public static final String EXTRA_STATUS = "android.content.rollback.extra.STATUS";
     field public static final String EXTRA_STATUS_MESSAGE = "android.content.rollback.extra.STATUS_MESSAGE";
     field public static final int STATUS_FAILURE = 1; // 0x1
diff --git a/core/java/android/content/rollback/RollbackManager.java b/core/java/android/content/rollback/RollbackManager.java
index 9038b03..799fb50 100644
--- a/core/java/android/content/rollback/RollbackManager.java
+++ b/core/java/android/content/rollback/RollbackManager.java
@@ -57,10 +57,12 @@
     /**
      * Returns a list of all currently available rollbacks.
      *
-     * @throws SecurityException if the caller does not have the
-     *            MANAGE_ROLLBACKS permission.
+     * @throws SecurityException if the caller does not have appropriate permissions.
      */
-    @RequiresPermission(android.Manifest.permission.MANAGE_ROLLBACKS)
+    @RequiresPermission(anyOf = {
+            android.Manifest.permission.MANAGE_ROLLBACKS,
+            android.Manifest.permission.TEST_MANAGE_ROLLBACKS
+    })
     @NonNull
     public List<RollbackInfo> getAvailableRollbacks() {
         try {
@@ -85,10 +87,12 @@
      * rolled back from.
      *
      * @return the recently committed rollbacks
-     * @throws SecurityException if the caller does not have the
-     *            MANAGE_ROLLBACKS permission.
+     * @throws SecurityException if the caller does not have appropriate permissions.
      */
-    @RequiresPermission(android.Manifest.permission.MANAGE_ROLLBACKS)
+    @RequiresPermission(anyOf = {
+            android.Manifest.permission.MANAGE_ROLLBACKS,
+            android.Manifest.permission.TEST_MANAGE_ROLLBACKS
+    })
     public @NonNull List<RollbackInfo> getRecentlyCommittedRollbacks() {
         try {
             return mBinder.getRecentlyExecutedRollbacks().getList();
@@ -171,10 +175,12 @@
      * @param statusReceiver where to deliver the results. Intents sent to
      *                       this receiver contain {@link #EXTRA_STATUS}
      *                       and {@link #EXTRA_STATUS_MESSAGE}.
-     * @throws SecurityException if the caller does not have the
-     *            MANAGE_ROLLBACKS permission.
+     * @throws SecurityException if the caller does not have appropriate permissions.
      */
-    @RequiresPermission(android.Manifest.permission.MANAGE_ROLLBACKS)
+    @RequiresPermission(anyOf = {
+            android.Manifest.permission.MANAGE_ROLLBACKS,
+            android.Manifest.permission.TEST_MANAGE_ROLLBACKS
+    })
     public void commitRollback(int rollbackId, @NonNull List<VersionedPackage> causePackages,
             @NonNull IntentSender statusReceiver) {
         try {
@@ -191,12 +197,11 @@
      * across device reboot, by simulating what happens on reboot without
      * actually rebooting the device.
      *
-     * @throws SecurityException if the caller does not have the
-     *            MANAGE_ROLLBACKS permission.
+     * @throws SecurityException if the caller does not have appropriate permissions.
      *
      * @hide
      */
-    @RequiresPermission(android.Manifest.permission.MANAGE_ROLLBACKS)
+    @RequiresPermission(android.Manifest.permission.TEST_MANAGE_ROLLBACKS)
     @TestApi
     public void reloadPersistedData() {
         try {
@@ -213,12 +218,11 @@
      * recently committed rollbacks that contain the given package.
      *
      * @param packageName the name of the package to expire data for.
-     * @throws SecurityException if the caller does not have the
-     *            MANAGE_ROLLBACKS permission.
+     * @throws SecurityException if the caller does not have appropriate permissions.
      *
      * @hide
      */
-    @RequiresPermission(android.Manifest.permission.MANAGE_ROLLBACKS)
+    @RequiresPermission(android.Manifest.permission.TEST_MANAGE_ROLLBACKS)
     @TestApi
     public void expireRollbackForPackage(@NonNull String packageName) {
         try {
diff --git a/core/res/AndroidManifest.xml b/core/res/AndroidManifest.xml
index c7417bf..efc6ad0 100644
--- a/core/res/AndroidManifest.xml
+++ b/core/res/AndroidManifest.xml
@@ -3929,6 +3929,10 @@
     <permission android:name="android.permission.MANAGE_ROLLBACKS"
         android:protectionLevel="signature|verifier" />
 
+    <!-- @SystemApi @TestApi @hide Allows testing apk level rollbacks. -->
+    <permission android:name="android.permission.TEST_MANAGE_ROLLBACKS"
+        android:protectionLevel="signature" />
+
     <!-- @SystemApi @hide Allows an application to mark other applications as harmful -->
     <permission android:name="android.permission.SET_HARMFUL_APP_WARNINGS"
         android:protectionLevel="signature|verifier" />
diff --git a/packages/Shell/AndroidManifest.xml b/packages/Shell/AndroidManifest.xml
index d639e5e..f0399c5 100644
--- a/packages/Shell/AndroidManifest.xml
+++ b/packages/Shell/AndroidManifest.xml
@@ -86,6 +86,7 @@
     <uses-permission android:name="android.permission.DELETE_CACHE_FILES" />
     <uses-permission android:name="android.permission.DELETE_PACKAGES" />
     <uses-permission android:name="android.permission.MANAGE_ROLLBACKS" />
+    <uses-permission android:name="android.permission.TEST_MANAGE_ROLLBACKS" />
     <uses-permission android:name="android.permission.ACCESS_SURFACE_FLINGER" />
     <uses-permission android:name="android.permission.READ_FRAME_BUFFER" />
     <uses-permission android:name="android.permission.DEVICE_POWER" />
diff --git a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
index 83d18a6..d93331e 100644
--- a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
+++ b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
@@ -16,6 +16,7 @@
 
 package com.android.server.rollback;
 
+import android.Manifest;
 import android.annotation.NonNull;
 import android.app.AppOpsManager;
 import android.content.BroadcastReceiver;
@@ -24,6 +25,7 @@
 import android.content.IntentFilter;
 import android.content.IntentSender;
 import android.content.pm.ApplicationInfo;
+import android.content.pm.ModuleInfo;
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageInstaller;
 import android.content.pm.PackageManager;
@@ -220,9 +222,7 @@
 
     @Override
     public ParceledListSlice getAvailableRollbacks() {
-        mContext.enforceCallingOrSelfPermission(
-                android.Manifest.permission.MANAGE_ROLLBACKS,
-                "getAvailableRollbacks");
+        enforceManageRollbacks("getAvailableRollbacks");
 
         synchronized (mLock) {
             ensureRollbackDataLoadedLocked();
@@ -239,9 +239,7 @@
 
     @Override
     public ParceledListSlice<RollbackInfo> getRecentlyExecutedRollbacks() {
-        mContext.enforceCallingOrSelfPermission(
-                android.Manifest.permission.MANAGE_ROLLBACKS,
-                "getRecentlyExecutedRollbacks");
+        enforceManageRollbacks("getRecentlyCommittedRollbacks");
 
         synchronized (mLock) {
             ensureRollbackDataLoadedLocked();
@@ -259,9 +257,7 @@
     @Override
     public void commitRollback(int rollbackId, ParceledListSlice causePackages,
             String callerPackageName, IntentSender statusReceiver) {
-        mContext.enforceCallingOrSelfPermission(
-                android.Manifest.permission.MANAGE_ROLLBACKS,
-                "executeRollback");
+        enforceManageRollbacks("executeRollback");
 
         final int callingUid = Binder.getCallingUid();
         AppOpsManager appOps = mContext.getSystemService(AppOpsManager.class);
@@ -462,7 +458,7 @@
                             // TODO: This call emits the warning "Calling a method in the
                             // system process without a qualified user". Fix that.
                             // TODO: Limit this to receivers holding the
-                            // MANAGE_ROLLBACKS permission?
+                            // MANAGE_ROLLBACKS or TEST_MANAGE_ROLLBACKS permissions?
                             mContext.sendBroadcast(broadcast);
                         });
                     }
@@ -484,7 +480,7 @@
     @Override
     public void reloadPersistedData() {
         mContext.enforceCallingOrSelfPermission(
-                android.Manifest.permission.MANAGE_ROLLBACKS,
+                Manifest.permission.TEST_MANAGE_ROLLBACKS,
                 "reloadPersistedData");
 
         synchronized (mLock) {
@@ -499,7 +495,7 @@
     @Override
     public void expireRollbackForPackage(String packageName) {
         mContext.enforceCallingOrSelfPermission(
-                android.Manifest.permission.MANAGE_ROLLBACKS,
+                Manifest.permission.TEST_MANAGE_ROLLBACKS,
                 "expireRollbackForPackage");
         synchronized (mLock) {
             ensureRollbackDataLoadedLocked();
@@ -894,12 +890,19 @@
         Log.i(TAG, "Enabling rollback for install of " + packageName
                 + ", session:" + session.sessionId);
 
+        String installerPackageName = session.getInstallerPackageName();
+        if (!enableRollbackAllowed(installerPackageName, packageName)) {
+            Log.e(TAG, "Installer " + installerPackageName
+                    + " is not allowed to enable rollback on " + packageName);
+            return false;
+        }
+
         VersionedPackage newVersion = new VersionedPackage(packageName, newPackage.versionCode);
         final boolean isApex = ((installFlags & PackageManager.INSTALL_APEX) != 0);
 
         // Get information about the currently installed package.
         PackageManager pm = mContext.getPackageManager();
-        PackageInfo pkgInfo = null;
+        final PackageInfo pkgInfo;
         try {
             pkgInfo = pm.getPackageInfo(packageName, isApex ? PackageManager.MATCH_APEX : 0);
         } catch (PackageManager.NameNotFoundException e) {
@@ -1086,6 +1089,44 @@
     }
 
     /**
+     * Returns true if the installer is allowed to enable rollback for the
+     * given named package, false otherwise.
+     */
+    private boolean enableRollbackAllowed(String installerPackageName, String packageName) {
+        if (installerPackageName == null) {
+            return false;
+        }
+
+        PackageManager pm = mContext.getPackageManager();
+        boolean manageRollbacksGranted = pm.checkPermission(
+                Manifest.permission.MANAGE_ROLLBACKS,
+                installerPackageName) == PackageManager.PERMISSION_GRANTED;
+
+        boolean testManageRollbacksGranted = pm.checkPermission(
+                Manifest.permission.TEST_MANAGE_ROLLBACKS,
+                installerPackageName) == PackageManager.PERMISSION_GRANTED;
+
+        // For now only allow rollbacks for modules or for testing.
+        return (isModule(packageName) && manageRollbacksGranted)
+            || testManageRollbacksGranted;
+    }
+
+    /**
+     * Returns true if the package name is the name of a module.
+     */
+    private boolean isModule(String packageName) {
+        PackageManager pm = mContext.getPackageManager();
+        final ModuleInfo moduleInfo;
+        try {
+            moduleInfo = pm.getModuleInfo(packageName, 0);
+        } catch (PackageManager.NameNotFoundException e) {
+            return false;
+        }
+
+        return moduleInfo != null;
+    }
+
+    /**
      * Gets the version of the package currently installed.
      * Returns null if the package is not currently installed.
      */
@@ -1311,4 +1352,15 @@
             }
         }
     }
+
+    private void enforceManageRollbacks(@NonNull String message) {
+        if ((PackageManager.PERMISSION_GRANTED != mContext.checkCallingOrSelfPermission(
+                        Manifest.permission.MANAGE_ROLLBACKS))
+                && (PackageManager.PERMISSION_GRANTED != mContext.checkCallingOrSelfPermission(
+                        Manifest.permission.TEST_MANAGE_ROLLBACKS))) {
+            throw new SecurityException(message + " requires "
+                    + Manifest.permission.MANAGE_ROLLBACKS + " or "
+                    + Manifest.permission.TEST_MANAGE_ROLLBACKS);
+        }
+    }
 }
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 7505230..679147a 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java
@@ -87,7 +87,7 @@
             RollbackTestUtils.adoptShellPermissionIdentity(
                     Manifest.permission.INSTALL_PACKAGES,
                     Manifest.permission.DELETE_PACKAGES,
-                    Manifest.permission.MANAGE_ROLLBACKS);
+                    Manifest.permission.TEST_MANAGE_ROLLBACKS);
 
             // Register a broadcast receiver for notification when the
             // rollback has been committed.
@@ -175,7 +175,7 @@
             RollbackTestUtils.adoptShellPermissionIdentity(
                     Manifest.permission.INSTALL_PACKAGES,
                     Manifest.permission.DELETE_PACKAGES,
-                    Manifest.permission.MANAGE_ROLLBACKS);
+                    Manifest.permission.TEST_MANAGE_ROLLBACKS);
 
             RollbackManager rm = RollbackTestUtils.getRollbackManager();
 
@@ -233,7 +233,7 @@
             RollbackTestUtils.adoptShellPermissionIdentity(
                     Manifest.permission.INSTALL_PACKAGES,
                     Manifest.permission.DELETE_PACKAGES,
-                    Manifest.permission.MANAGE_ROLLBACKS);
+                    Manifest.permission.TEST_MANAGE_ROLLBACKS);
 
             RollbackManager rm = RollbackTestUtils.getRollbackManager();
 
@@ -290,7 +290,7 @@
             RollbackTestUtils.adoptShellPermissionIdentity(
                     Manifest.permission.INSTALL_PACKAGES,
                     Manifest.permission.DELETE_PACKAGES,
-                    Manifest.permission.MANAGE_ROLLBACKS);
+                    Manifest.permission.TEST_MANAGE_ROLLBACKS);
 
             RollbackManager rm = RollbackTestUtils.getRollbackManager();
 
@@ -343,7 +343,7 @@
             RollbackTestUtils.adoptShellPermissionIdentity(
                     Manifest.permission.INSTALL_PACKAGES,
                     Manifest.permission.DELETE_PACKAGES,
-                    Manifest.permission.MANAGE_ROLLBACKS,
+                    Manifest.permission.TEST_MANAGE_ROLLBACKS,
                     Manifest.permission.WRITE_DEVICE_CONFIG);
 
             DeviceConfig.setProperty(DeviceConfig.Rollback.BOOT_NAMESPACE,
@@ -403,7 +403,7 @@
             RollbackTestUtils.adoptShellPermissionIdentity(
                     Manifest.permission.INSTALL_PACKAGES,
                     Manifest.permission.DELETE_PACKAGES,
-                    Manifest.permission.MANAGE_ROLLBACKS,
+                    Manifest.permission.TEST_MANAGE_ROLLBACKS,
                     Manifest.permission.WRITE_DEVICE_CONFIG,
                     Manifest.permission.SET_TIME);
 
@@ -475,7 +475,7 @@
             RollbackTestUtils.adoptShellPermissionIdentity(
                     Manifest.permission.INSTALL_PACKAGES,
                     Manifest.permission.DELETE_PACKAGES,
-                    Manifest.permission.MANAGE_ROLLBACKS);
+                    Manifest.permission.TEST_MANAGE_ROLLBACKS);
 
             RollbackManager rm = RollbackTestUtils.getRollbackManager();
             RollbackTestUtils.uninstall(TEST_APP_A);
@@ -512,7 +512,7 @@
             RollbackTestUtils.adoptShellPermissionIdentity(
                     Manifest.permission.INSTALL_PACKAGES,
                     Manifest.permission.DELETE_PACKAGES,
-                    Manifest.permission.MANAGE_ROLLBACKS);
+                    Manifest.permission.TEST_MANAGE_ROLLBACKS);
 
             RollbackTestUtils.uninstall(TEST_APP_A);
             RollbackTestUtils.install("RollbackTestAppAv1.apk", false);
@@ -540,7 +540,7 @@
             RollbackTestUtils.adoptShellPermissionIdentity(
                     Manifest.permission.INSTALL_PACKAGES,
                     Manifest.permission.DELETE_PACKAGES,
-                    Manifest.permission.MANAGE_ROLLBACKS);
+                    Manifest.permission.TEST_MANAGE_ROLLBACKS);
 
             RollbackTestUtils.uninstall(TEST_APP_A);
             RollbackTestUtils.installSplit(false,
@@ -598,7 +598,7 @@
             RollbackTestUtils.adoptShellPermissionIdentity(
                     Manifest.permission.INSTALL_PACKAGES,
                     Manifest.permission.DELETE_PACKAGES,
-                    Manifest.permission.MANAGE_ROLLBACKS);
+                    Manifest.permission.TEST_MANAGE_ROLLBACKS);
             RollbackManager rm = RollbackTestUtils.getRollbackManager();
 
             // Prep installation of the test apps.
@@ -693,6 +693,75 @@
     }
 
     /**
+     * Test that you cannot enable rollback for a package without the
+     * MANAGE_ROLLBACKS permission.
+     */
+    @Test
+    public void testEnableRollbackPermission() throws Exception {
+        try {
+            RollbackTestUtils.adoptShellPermissionIdentity(
+                    Manifest.permission.INSTALL_PACKAGES,
+                    Manifest.permission.DELETE_PACKAGES);
+
+            RollbackTestUtils.uninstall(TEST_APP_A);
+            RollbackTestUtils.install("RollbackTestAppAv1.apk", /* enableRollback */ false);
+            assertEquals(1, RollbackTestUtils.getInstalledVersion(TEST_APP_A));
+
+            RollbackTestUtils.install("RollbackTestAppAv2.apk", /* enableRollback */ true);
+
+            // We expect v2 of the app was installed, but rollback has not
+            // been enabled.
+            assertEquals(2, RollbackTestUtils.getInstalledVersion(TEST_APP_A));
+
+            // TODO: See if there is a way to remove this race condition
+            // between when the app is installed and when the rollback
+            // would be made available.
+            Thread.sleep(1000);
+
+            RollbackTestUtils.adoptShellPermissionIdentity(
+                    Manifest.permission.TEST_MANAGE_ROLLBACKS);
+            RollbackManager rm = RollbackTestUtils.getRollbackManager();
+            assertNull(getUniqueRollbackInfoForPackage(rm.getAvailableRollbacks(), TEST_APP_A));
+        } finally {
+            RollbackTestUtils.dropShellPermissionIdentity();
+        }
+    }
+
+    /**
+     * Test that you cannot enable rollback for a non-module package when
+     * holding the MANAGE_ROLLBACKS permission.
+     */
+    @Test
+    public void testNonModuleEnableRollback() throws Exception {
+        try {
+            RollbackTestUtils.adoptShellPermissionIdentity(
+                    Manifest.permission.INSTALL_PACKAGES,
+                    Manifest.permission.DELETE_PACKAGES,
+                    Manifest.permission.MANAGE_ROLLBACKS);
+
+            RollbackTestUtils.uninstall(TEST_APP_A);
+            RollbackTestUtils.install("RollbackTestAppAv1.apk", /* enableRollback */ false);
+            assertEquals(1, RollbackTestUtils.getInstalledVersion(TEST_APP_A));
+
+            RollbackTestUtils.install("RollbackTestAppAv2.apk", /* enableRollback */ true);
+
+            // We expect v2 of the app was installed, but rollback has not
+            // been enabled because the test app is not a module.
+            assertEquals(2, RollbackTestUtils.getInstalledVersion(TEST_APP_A));
+
+            // TODO: See if there is a way to remove this race condition
+            // between when the app is installed and when the rollback
+            // would be made available.
+            Thread.sleep(1000);
+
+            RollbackManager rm = RollbackTestUtils.getRollbackManager();
+            assertNull(getUniqueRollbackInfoForPackage(rm.getAvailableRollbacks(), TEST_APP_A));
+        } finally {
+            RollbackTestUtils.dropShellPermissionIdentity();
+        }
+    }
+
+    /**
      * Test rollback of multi-package installs is implemented.
      */
     @Test
@@ -701,7 +770,7 @@
             RollbackTestUtils.adoptShellPermissionIdentity(
                     Manifest.permission.INSTALL_PACKAGES,
                     Manifest.permission.DELETE_PACKAGES,
-                    Manifest.permission.MANAGE_ROLLBACKS);
+                    Manifest.permission.TEST_MANAGE_ROLLBACKS);
             RollbackManager rm = RollbackTestUtils.getRollbackManager();
 
             // Prep installation of the test apps.
@@ -760,7 +829,7 @@
             RollbackTestUtils.adoptShellPermissionIdentity(
                     Manifest.permission.INSTALL_PACKAGES,
                     Manifest.permission.DELETE_PACKAGES,
-                    Manifest.permission.MANAGE_ROLLBACKS,
+                    Manifest.permission.TEST_MANAGE_ROLLBACKS,
                     Manifest.permission.KILL_BACKGROUND_PROCESSES,
                     Manifest.permission.RESTART_PACKAGES);
             RollbackManager rm = RollbackTestUtils.getRollbackManager();
diff --git a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
index 047451b..7e711c2 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
@@ -63,7 +63,7 @@
         RollbackTestUtils.adoptShellPermissionIdentity(
                 Manifest.permission.INSTALL_PACKAGES,
                 Manifest.permission.DELETE_PACKAGES,
-                Manifest.permission.MANAGE_ROLLBACKS);
+                Manifest.permission.TEST_MANAGE_ROLLBACKS);
     }
 
     /**