Delete APEX CE snapshots when a rollbck is deleted (2/n)

Now we delete APEX CE snapshots immediately when a rollback is deleted.
If the user is locked, the deletion will fail, but that's OK. The
orphaned snapshots will be cleaned when the user is unlocked.

Bug: 158712947
Test: atest StagedRollbackTest RollbackUnitTest
Change-Id: I3e3eef4c52f281d28a6a6ac6735c27ef8b83ddc2

services/core/java/com/android/server/pm/ApexManager.java

@@ -341,9 +341,18 @@ public abstract class ApexManager {
      */
     public abstract boolean destroyDeSnapshots(int rollbackId);
 
+    /**
+     *  Deletes snapshots of the credential encrypted apex data directories for the specified user,
+     *  for the given rollback id as long as the user is credential unlocked.
+     *
+     * @return boolean true if the delete was successful
+     */
+    public abstract boolean destroyCeSnapshots(int userId, int rollbackId);
+
     /**
      * Deletes snapshots of the credential encrypted apex data directories for the specified user,
-     * where the rollback id is not included in {@code retainRollbackIds}.
+     * where the rollback id is not included in {@code retainRollbackIds} as long as the user is
+     * credential unlocked.
      *
      * @return boolean true if the delete was successful
      */
@@ -874,6 +883,17 @@ public abstract class ApexManager {
             }
         }
 
+        @Override
+        public boolean destroyCeSnapshots(int userId, int rollbackId) {
+            try {
+                waitForApexService().destroyCeSnapshots(userId, rollbackId);
+                return true;
+            } catch (Exception e) {
+                Slog.e(TAG, e.getMessage(), e);
+                return false;
+            }
+        }
+
         @Override
         public boolean destroyCeSnapshotsNotSpecified(int userId, int[] retainRollbackIds) {
             try {
@@ -1135,6 +1155,11 @@ public abstract class ApexManager {
             throw new UnsupportedOperationException();
         }
 
+        @Override
+        public boolean destroyCeSnapshots(int userId, int rollbackId) {
+            return true;
+        }
+
         @Override
         public boolean destroyCeSnapshotsNotSpecified(int userId, int[] retainRollbackIds) {
             return true;

services/core/java/com/android/server/rollback/AppDataRollbackHelper.java

@@ -205,6 +205,16 @@ public class AppDataRollbackHelper {
         mApexManager.destroyDeSnapshots(rollbackId);
     }
 
+    /**
+     * Deletes snapshots of the credential encrypted apex data directories for the specified user,
+     * for the given rollback id. This method will be a no-op if the user is not unlocked.
+     */
+    public void destroyApexCeSnapshots(int userId, int rollbackId) {
+        if (!isUserCredentialLocked(userId)) {
+            mApexManager.destroyCeSnapshots(userId, rollbackId);
+        }
+    }
+
     /**
      * Commits the pending backups and restores for a given {@code userId} and {@code rollback}. If
      * the rollback has a pending backup, it is updated with a mapping from {@code userId} to inode

services/core/java/com/android/server/rollback/Rollback.java

@@ -42,6 +42,7 @@ import android.os.UserHandle;
 import android.os.UserManager;
 import android.os.ext.SdkExtensions;
 import android.text.TextUtils;
+import android.util.ArraySet;
 import android.util.Slog;
 import android.util.SparseIntArray;
 
@@ -61,9 +62,9 @@ import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.Set;
 import java.util.function.Consumer;
 
-
 /**
  * Information about a rollback available for a set of atomically installed packages.
  *
@@ -699,11 +700,13 @@ class Rollback {
     void delete(AppDataRollbackHelper dataHelper) {
         assertInWorkerThread();
         boolean containsApex = false;
+        Set<Integer> apexUsers = new ArraySet<>();
         for (PackageRollbackInfo pkgInfo : info.getPackages()) {
+            List<Integer> snapshottedUsers = pkgInfo.getSnapshottedUsers();
             if (pkgInfo.isApex()) {
                 containsApex = true;
+                apexUsers.addAll(snapshottedUsers);
             } else {
-                List<Integer> snapshottedUsers = pkgInfo.getSnapshottedUsers();
                 for (int i = 0; i < snapshottedUsers.size(); i++) {
                     // Destroy app data snapshot.
                     int userId = snapshottedUsers.get(i);
@@ -714,6 +717,9 @@ class Rollback {
         }
         if (containsApex) {
             dataHelper.destroyApexDeSnapshots(info.getRollbackId());
+            for (int user : apexUsers) {
+                dataHelper.destroyApexCeSnapshots(user, info.getRollbackId());
+            }
         }
 
         RollbackStore.deleteRollback(this);

services/tests/servicestests/src/com/android/server/rollback/RollbackUnitTest.java

@@ -252,6 +252,7 @@ public class RollbackUnitTest {
         verify(mMockDataHelper).destroyAppDataSnapshot(eq(123), pkgRollbackInfoFor(PKG_2), eq(111));
         verify(mMockDataHelper).destroyAppDataSnapshot(eq(123), pkgRollbackInfoFor(PKG_2), eq(222));
         verify(mMockDataHelper, never()).destroyApexDeSnapshots(anyInt());
+        verify(mMockDataHelper, never()).destroyApexCeSnapshots(anyInt(), anyInt());
 
         assertThat(rollback.isDeleted()).isTrue();
     }
@@ -273,6 +274,8 @@ public class RollbackUnitTest {
         verify(mMockDataHelper, never())
                 .destroyAppDataSnapshot(anyInt(), pkgRollbackInfoFor(PKG_2), anyInt());
         verify(mMockDataHelper).destroyApexDeSnapshots(123);
+        verify(mMockDataHelper).destroyApexCeSnapshots(111, 123);
+        verify(mMockDataHelper).destroyApexCeSnapshots(222, 123);
 
         assertThat(rollback.isDeleted()).isTrue();
     }

tests/RollbackTest/StagedRollbackTest/src/com/android/tests/rollback/host/StagedRollbackTest.java

@@ -453,6 +453,32 @@ public class StagedRollbackTest extends BaseHostJUnit4Test {
         after.forEach(dir -> assertDirectoryIsEmpty(dir));
     }
 
+    @Test
+    public void testExpireApexRollback() throws Exception {
+        List<String> before = getSnapshotDirectories("/data/misc_ce/0/apexrollback");
+        pushTestApex();
+
+        // Push files to apex data directory
+        String oldFilePath1 = apexDataDirCe(APK_IN_APEX_TESTAPEX_NAME, 0) + "/" + TEST_FILENAME_1;
+        String oldFilePath2 =
+                apexDataDirCe(APK_IN_APEX_TESTAPEX_NAME, 0) + TEST_SUBDIR + TEST_FILENAME_2;
+        assertTrue(getDevice().pushString(TEST_STRING_1, oldFilePath1));
+        assertTrue(getDevice().pushString(TEST_STRING_2, oldFilePath2));
+
+        // Install new version of the APEX with rollback enabled
+        runPhase("testRollbackApexDataDirectories_Phase1");
+        getDevice().reboot();
+
+        List<String> after = getSnapshotDirectories("/data/misc_ce/0/apexrollback");
+        // Only check directories newly created during the test
+        after.removeAll(before);
+        // Expire all rollbacks and check CE snapshot directories are deleted
+        runPhase("testCleanUp");
+        for (String dir : after) {
+            assertNull(getDevice().getFileEntry(dir));
+        }
+    }
+
     private void pushTestApex() throws Exception {
         CompatibilityBuildHelper buildHelper = new CompatibilityBuildHelper(getBuild());
         final String fileName = APK_IN_APEX_TESTAPEX_NAME + "_v1.apex";