Expire rollback when apex is updated.

RollbackManager relies on package changed broadcasts to expire rollbacks
when an application is updated, but we don't receive package changed
broadcasts for apex.

This change adds an extra check when starting RollbackManager to see if
any apex versions have changed and expires rollbacks as appropriate.

Adds a test case to cover the scenario and refactors the test code
to properly set up the test apex in all cases with minimal reboots.

Bug: 126358044
Test: atest StagedRollbackTest, with new test added.
Change-Id: I7ea4953e4aff8d1c7560d6c61e6be5e4e8e1f194

5 files changed, 122 insertions, 36 deletions

diff --git a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
index 615472661f9e..db2c742b8d4e 100644
--- a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
+++ b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
@@ -70,6 +70,7 @@ import java.time.Instant;
 import java.time.temporal.ChronoUnit;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -588,6 +589,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
             // rollback sessions been applied.
             List<RollbackData> enabling = new ArrayList<>();
             List<RollbackData> restoreInProgress = new ArrayList<>();
+            Set<String> apexPackageNames = new HashSet<>();
             synchronized (mLock) {
                 ensureRollbackDataLoadedLocked();
                 for (RollbackData data : mRollbacks) {
@@ -597,6 +599,12 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
                         } else if (data.restoreUserDataInProgress) {
                             restoreInProgress.add(data);
                         }
+
+                        for (PackageRollbackInfo info : data.info.getPackages()) {
+                            if (info.isApex()) {
+                                apexPackageNames.add(info.getPackageName());
+                            }
+                        }
                     }
                 }
             }
@@ -634,6 +642,14 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
                 }
             }
 
+            for (String apexPackageName : apexPackageNames) {
+                // We will not recieve notifications when an apex is updated,
+                // so check now in case any rollbacks ought to be expired. The
+                // onPackagedReplace function is safe to call if the package
+                // hasn't actually been updated.
+                onPackageReplaced(apexPackageName);
+            }
+
             mPackageHealthObserver.onBootCompleted();
         });
     }
diff --git a/tests/RollbackTest/Android.bp b/tests/RollbackTest/Android.bp
index dfc3b6e15b3e..e556b0acb1a3 100644
--- a/tests/RollbackTest/Android.bp
+++ b/tests/RollbackTest/Android.bp
@@ -88,6 +88,15 @@ apex {
     installable: false,
 }
 
+apex {
+    name: "com.android.tests.rollback.testapex.RollbackTestApexV3",
+    manifest: "TestApex/RollbackTestApexV3.json",
+    file_contexts: "apex.test",
+    prebuilts: ["RollbackTestApex.prebuilt.txt"],
+    key: "RollbackTestApex.key",
+    installable: false,
+}
+
 apex_key {
     name: "RollbackTestApex.key",
     public_key: "TestApex/com.android.tests.rollback.testapex.avbpubkey",
@@ -116,6 +125,7 @@ android_test {
         ":RollbackTestAppASplitV2",
         ":com.android.tests.rollback.testapex.RollbackTestApexV1",
         ":com.android.tests.rollback.testapex.RollbackTestApexV2",
+        ":com.android.tests.rollback.testapex.RollbackTestApexV3",
     ],
     test_config: "RollbackTest.xml",
     sdk_version: "test_current",
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 7e711c290e5a..3b0e2a56db3c 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
@@ -26,6 +26,7 @@ import android.content.rollback.RollbackManager;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
 import org.junit.After;
@@ -54,6 +55,8 @@ public class StagedRollbackTest {
             "com.android.tests.rollback.testapex.RollbackTestApexV1.apex";
     private static final String TEST_APEX_V2 =
             "com.android.tests.rollback.testapex.RollbackTestApexV2.apex";
+    private static final String TEST_APEX_V3 =
+            "com.android.tests.rollback.testapex.RollbackTestApexV3.apex";
 
     /**
      * Adopts common shell permissions needed for rollback tests.
@@ -145,26 +148,13 @@ public class StagedRollbackTest {
 
     /**
      * Test rollbacks of staged installs an apk and an apex.
-     * Prepare apex (and apk) phase.
-     */
-    @Test
-    public void testApkAndApexPrepare() throws Exception {
-        RollbackTestUtils.uninstall(TEST_APP_A);
-        assertEquals(-1, RollbackTestUtils.getInstalledVersion(TEST_APP_A));
-
-        // Note: can't uninstall the apex. See note in #testApexOnlyPrepareApex().
-        RollbackTestUtils.installStaged(false, TEST_APP_A_V1, TEST_APEX_V1);
-
-        // At this point, the host test driver will reboot the device and run
-        // testApkAndApexEnableRollback().
-    }
-
-    /**
-     * Test rollbacks of staged installs an apk and an apex.
      * Enable rollback phase.
      */
     @Test
     public void testApkAndApexEnableRollback() throws Exception {
+        RollbackTestUtils.uninstall(TEST_APP_A);
+        RollbackTestUtils.install(TEST_APP_A_V1, false);
+
         assertEquals(1, RollbackTestUtils.getInstalledVersion(TEST_APEX_PKG));
         assertEquals(1, RollbackTestUtils.getInstalledVersion(TEST_APP_A));
 
@@ -225,22 +215,6 @@ public class StagedRollbackTest {
 
     /**
      * Test rollbacks of staged installs involving only apex.
-     * Prepare apex phase.
-     */
-    @Test
-    public void testApexOnlyPrepareApex() throws Exception {
-        // Note: We can't uninstall the apex if it is already on device,
-        // because that isn't supported yet (b/123667725). As long as nothing
-        // is failing, this should be fine because we don't expect the tests
-        // to leave the device with v2 of the apex installed.
-        RollbackTestUtils.installStaged(false, TEST_APEX_V1);
-
-        // At this point, the host test driver will reboot the device and run
-        // testApexOnlyEnableRollback().
-    }
-
-    /**
-     * Test rollbacks of staged installs involving only apex.
      * Enable rollback phase.
      */
     @Test
@@ -291,4 +265,51 @@ public class StagedRollbackTest {
     public void testApexOnlyConfirmRollback() throws Exception {
         assertEquals(1, RollbackTestUtils.getInstalledVersion(TEST_APEX_PKG));
     }
+
+    /**
+     * Tests that apex update expires existing rollbacks for that apex.
+     * Enable rollback phase.
+     */
+    @Test
+    public void testApexRollbackExpirationEnableRollback() throws Exception {
+        assertEquals(1, RollbackTestUtils.getInstalledVersion(TEST_APEX_PKG));
+        RollbackTestUtils.installStaged(true, TEST_APEX_V2);
+
+        // At this point, the host test driver will reboot the device and run
+        // testApexRollbackExpirationUpdateApex().
+    }
+
+    /**
+     * Tests that apex update expires existing rollbacks for that apex.
+     * Update apex phase.
+     */
+    @Test
+    public void testApexRollbackExpirationUpdateApex() throws Exception {
+        assertEquals(2, RollbackTestUtils.getInstalledVersion(TEST_APEX_PKG));
+        RollbackTestUtils.installStaged(false, TEST_APEX_V3);
+
+        // At this point, the host test driver will reboot the device and run
+        // testApexRollbackExpirationConfirmExpiration().
+    }
+
+    /**
+     * Tests that apex update expires existing rollbacks for that apex.
+     * Confirm expiration phase.
+     */
+    @Test
+    public void testApexRollbackExpirationConfirmExpiration() throws Exception {
+        assertEquals(3, RollbackTestUtils.getInstalledVersion(TEST_APEX_PKG));
+
+        RollbackManager rm = RollbackTestUtils.getRollbackManager();
+        assertNull(getUniqueRollbackInfoForPackage(rm.getAvailableRollbacks(), TEST_APEX_PKG));
+    }
+
+    /**
+     * Helper function called by the host test to install v1 of the test apex,
+     * assuming the test apex is not installed.
+     */
+    @Test
+    public void installTestApexV1() throws Exception {
+        RollbackTestUtils.installStaged(false, TEST_APEX_V1);
+    }
 }
diff --git a/tests/RollbackTest/StagedRollbackTest/src/com/android/tests/rollback/host/StagedRollbackTest.java b/tests/RollbackTest/StagedRollbackTest/src/com/android/tests/rollback/host/StagedRollbackTest.java
index ac7f634d51f1..1f87ed863034 100644
--- a/tests/RollbackTest/StagedRollbackTest/src/com/android/tests/rollback/host/StagedRollbackTest.java
+++ b/tests/RollbackTest/StagedRollbackTest/src/com/android/tests/rollback/host/StagedRollbackTest.java
@@ -18,6 +18,7 @@ package com.android.tests.rollback.host;
 
 import static org.junit.Assert.assertTrue;
 
+import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
 import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test;
 
@@ -30,6 +31,8 @@ import org.junit.runner.RunWith;
 @RunWith(DeviceJUnit4ClassRunner.class)
 public class StagedRollbackTest extends BaseHostJUnit4Test {
 
+    private static final String TEST_APEX_PKG = "com.android.tests.rollback.testapex";
+
     /**
      * Runs the given phase of a test by calling into the device.
      * Throws an exception if the test phase fails.
@@ -59,8 +62,7 @@ public class StagedRollbackTest extends BaseHostJUnit4Test {
      */
     @Test
     public void testApexOnly() throws Exception {
-        runPhase("testApexOnlyPrepareApex");
-        getDevice().reboot();
+        installTestApexV1();
         runPhase("testApexOnlyEnableRollback");
         getDevice().reboot();
         runPhase("testApexOnlyCommitRollback");
@@ -73,12 +75,45 @@ public class StagedRollbackTest extends BaseHostJUnit4Test {
      */
     @Test
     public void testApkAndApex() throws Exception {
-        runPhase("testApkAndApexPrepare");
-        getDevice().reboot();
+        installTestApexV1();
         runPhase("testApkAndApexEnableRollback");
         getDevice().reboot();
         runPhase("testApkAndApexCommitRollback");
         getDevice().reboot();
         runPhase("testApkAndApexConfirmRollback");
     }
+
+    /**
+     * Tests that apex update expires existing rollbacks for that apex.
+     */
+    @Test
+    public void testApexRollbackExpiration() throws Exception {
+        installTestApexV1();
+        runPhase("testApexRollbackExpirationEnableRollback");
+        getDevice().reboot();
+        runPhase("testApexRollbackExpirationUpdateApex");
+        getDevice().reboot();
+        runPhase("testApexRollbackExpirationConfirmExpiration");
+    }
+
+    /**
+     * Do whatever is necessary to get version 1 of the test apex installed on
+     * the device. Try to do so without extra reboots where possible to keep
+     * the test execution time down.
+     */
+    private void installTestApexV1() throws Exception {
+        for (ITestDevice.ApexInfo apexInfo : getDevice().getActiveApexes()) {
+            if (TEST_APEX_PKG.equals(apexInfo.name)) {
+                if (apexInfo.versionCode == 1) {
+                    return;
+                }
+                getDevice().uninstallPackage(TEST_APEX_PKG);
+                getDevice().reboot();
+                break;
+            }
+        }
+
+        runPhase("installTestApexV1");
+        getDevice().reboot();
+    }
 }
diff --git a/tests/RollbackTest/TestApex/RollbackTestApexV3.json b/tests/RollbackTest/TestApex/RollbackTestApexV3.json
new file mode 100644
index 000000000000..87a2c9dbcb6a
--- /dev/null
+++ b/tests/RollbackTest/TestApex/RollbackTestApexV3.json
@@ -0,0 +1,4 @@
+{
+    "name": "com.android.tests.rollback.testapex",
+    "version": 3
+}