RMS: Delete user data snapshots when deleting a rollback

- PackageRollbackInfo now also stores list of users the package is
  installed for and the inodes of backup CE data snapshot.
- Changed snapshotAppData to return a pair of
  <pendingBackups, ceSnapshotInodes>.
- RollbackData is updated with CE user data snapshot inode after user
  unlocks the phone and snapshotAppData call is made.
- Added javadoc to some methods in Installer.java to make lint happy.

Not sure how to automatically test:
- multi-user flow.
- the fact that user data snapshots were deleted. I can't access folder
  with snapshot data in RollbackTest, because it is owned by installd.

Bug: 112431924
Test: AppDataRollbackHelperTest, RollbackTest
Change-Id: Ife15f1aed9cf1d2b15be238bdaf766c51b85cab6

6 files changed, 297 insertions, 45 deletions

diff --git a/core/java/android/content/rollback/PackageRollbackInfo.java b/core/java/android/content/rollback/PackageRollbackInfo.java
index 0ec40183c319..1d0ab5ad2679 100644
--- a/core/java/android/content/rollback/PackageRollbackInfo.java
+++ b/core/java/android/content/rollback/PackageRollbackInfo.java
@@ -22,6 +22,7 @@ import android.content.pm.VersionedPackage;
 import android.os.Parcel;
 import android.os.Parcelable;
 import android.util.IntArray;
+import android.util.SparseLongArray;
 
 import java.util.ArrayList;
 
@@ -73,6 +74,18 @@ public final class PackageRollbackInfo implements Parcelable {
      */
     private final boolean mIsApex;
 
+    /*
+     * The list of users the package is installed for.
+     */
+    // NOTE: Not a part of the Parcelable representation of this object.
+    private final IntArray mInstalledUsers;
+
+    /**
+     * A mapping between user and an inode of theirs CE data snapshot.
+     */
+    // NOTE: Not a part of the Parcelable representation of this object.
+    private final SparseLongArray mCeSnapshotInodes;
+
     /**
      * Returns the name of the package to roll back from.
      */
@@ -126,15 +139,33 @@ public final class PackageRollbackInfo implements Parcelable {
     }
 
     /** @hide */
+    public IntArray getInstalledUsers() {
+        return mInstalledUsers;
+    }
+
+    /** @hide */
+    public SparseLongArray getCeSnapshotInodes() {
+        return mCeSnapshotInodes;
+    }
+
+    /** @hide */
+    public void putCeSnapshotInode(int userId, long ceSnapshotInode) {
+        mCeSnapshotInodes.put(userId, ceSnapshotInode);
+    }
+
+    /** @hide */
     public PackageRollbackInfo(VersionedPackage packageRolledBackFrom,
             VersionedPackage packageRolledBackTo,
             @NonNull IntArray pendingBackups, @NonNull ArrayList<RestoreInfo> pendingRestores,
-            boolean isApex) {
+            boolean isApex, @NonNull IntArray installedUsers,
+            @NonNull SparseLongArray ceSnapshotInodes) {
         this.mVersionRolledBackFrom = packageRolledBackFrom;
         this.mVersionRolledBackTo = packageRolledBackTo;
         this.mPendingBackups = pendingBackups;
         this.mPendingRestores = pendingRestores;
         this.mIsApex = isApex;
+        this.mInstalledUsers = installedUsers;
+        this.mCeSnapshotInodes = ceSnapshotInodes;
     }
 
     private PackageRollbackInfo(Parcel in) {
@@ -143,6 +174,8 @@ public final class PackageRollbackInfo implements Parcelable {
         this.mIsApex = in.readBoolean();
         this.mPendingRestores = null;
         this.mPendingBackups = null;
+        this.mInstalledUsers = null;
+        this.mCeSnapshotInodes = null;
     }
 
     @Override
diff --git a/services/core/java/com/android/server/pm/Installer.java b/services/core/java/com/android/server/pm/Installer.java
index efafdfaf2b54..c2a75aba28b9 100644
--- a/services/core/java/com/android/server/pm/Installer.java
+++ b/services/core/java/com/android/server/pm/Installer.java
@@ -611,18 +611,43 @@ public class Installer extends SystemService {
         }
     }
 
-    public boolean snapshotAppData(String pkg, @UserIdInt int userId, int storageFlags)
+    /**
+     * Snapshots user data of the given package.
+     *
+     * @param pkg name of the package to snapshot user data for.
+     * @param userId id of the user whose data to snapshot.
+     * @param storageFlags flags controlling which data (CE or DE) to snapshot.
+     *
+     * @return inode of the snapshot of users CE package data, or {@code 0} if a remote calls
+     *  shouldn't be continued. See {@link #checkBeforeRemote}.
+     *
+     * @throws InstallerException if failed to snapshot user data.
+     */
+    public long snapshotAppData(String pkg, @UserIdInt int userId, int storageFlags)
             throws InstallerException {
-        if (!checkBeforeRemote()) return false;
+        if (!checkBeforeRemote()) return 0;
 
         try {
-            mInstalld.snapshotAppData(null, pkg, userId, storageFlags);
-            return true;
+            return mInstalld.snapshotAppData(null, pkg, userId, storageFlags);
         } catch (Exception e) {
             throw InstallerException.from(e);
         }
     }
 
+    /**
+     * Restores user data snapshot of the given package.
+     *
+     * @param pkg name of the package to restore user data for.
+     * @param appId id of the package to restore user data for.
+     * @param ceDataInode inode of CE user data folder of this app.
+     * @param userId id of the user whose data to restore.
+     * @param storageFlags flags controlling which data (CE or DE) to restore.
+     *
+     * @return {@code true} if user data restore was successful, or {@code false} if a remote call
+     *  shouldn't be continued. See {@link #checkBeforeRemote}.
+     *
+     * @throws InstallerException if failed to restore user data.
+     */
     public boolean restoreAppDataSnapshot(String pkg, @AppIdInt  int appId, long ceDataInode,
             String seInfo, @UserIdInt int userId, int storageFlags) throws InstallerException {
         if (!checkBeforeRemote()) return false;
@@ -636,6 +661,31 @@ public class Installer extends SystemService {
         }
     }
 
+    /**
+     * Deletes user data snapshot of the given package.
+     *
+     * @param pkg name of the package to delete user data snapshot for.
+     * @param userId id of the user whose user data snapshot to delete.
+     * @param ceSnapshotInode inode of CE user data snapshot.
+     * @param storageFlags flags controlling which user data snapshot (CE or DE) to delete.
+     *
+     * @return {@code true} if user data snapshot was successfully deleted, or {@code false} if a
+     *  remote call shouldn't be continued. See {@link #checkBeforeRemote}.
+     *
+     * @throws InstallerException if failed to delete user data snapshot.
+     */
+    public boolean destroyAppDataSnapshot(String pkg, @UserIdInt int userId, long ceSnapshotInode,
+            int storageFlags) throws InstallerException {
+        if (!checkBeforeRemote()) return false;
+
+        try {
+            mInstalld.destroyAppDataSnapshot(null, pkg, userId, ceSnapshotInode, storageFlags);
+            return true;
+        } catch (Exception e) {
+            throw InstallerException.from(e);
+        }
+    }
+
     private static void assertValidInstructionSet(String instructionSet)
             throws InstallerException {
         for (String abi : Build.SUPPORTED_ABIS) {
diff --git a/services/core/java/com/android/server/rollback/AppDataRollbackHelper.java b/services/core/java/com/android/server/rollback/AppDataRollbackHelper.java
index 8dd076028b43..f3b838560ebd 100644
--- a/services/core/java/com/android/server/rollback/AppDataRollbackHelper.java
+++ b/services/core/java/com/android/server/rollback/AppDataRollbackHelper.java
@@ -22,6 +22,7 @@ import android.content.rollback.RollbackInfo;
 import android.os.storage.StorageManager;
 import android.util.IntArray;
 import android.util.Log;
+import android.util.SparseLongArray;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.pm.Installer;
@@ -51,11 +52,13 @@ public class AppDataRollbackHelper {
      * Creates an app data snapshot for a specified {@code packageName} for {@code installedUsers},
      * a specified set of users for whom the package is installed.
      *
-     * @return a list of users for which the snapshot is pending, usually because data for one or
-     *         more users is still credential locked.
+     * @return a {@link SnapshotAppDataResult}/
+     * @see SnapshotAppDataResult
      */
-    public IntArray snapshotAppData(String packageName, int[] installedUsers) {
+    public SnapshotAppDataResult snapshotAppData(String packageName, int[] installedUsers) {
         final IntArray pendingBackups = new IntArray();
+        final SparseLongArray ceSnapshotInodes = new SparseLongArray();
+
         for (int user : installedUsers) {
             final int storageFlags;
             if (isUserCredentialLocked(user)) {
@@ -69,14 +72,17 @@ public class AppDataRollbackHelper {
             }
 
             try {
-                mInstaller.snapshotAppData(packageName, user, storageFlags);
+                long ceSnapshotInode = mInstaller.snapshotAppData(packageName, user, storageFlags);
+                if ((storageFlags & Installer.FLAG_STORAGE_CE) != 0) {
+                    ceSnapshotInodes.put(user, ceSnapshotInode);
+                }
             } catch (InstallerException ie) {
                 Log.e(TAG, "Unable to create app data snapshot for: " + packageName
                         + ", userId: " + user, ie);
             }
         }
 
-        return pendingBackups;
+        return new SnapshotAppDataResult(pendingBackups, ceSnapshotInodes);
     }
 
     /**
@@ -138,6 +144,22 @@ public class AppDataRollbackHelper {
     }
 
     /**
+     * Deletes an app data data snapshot for a specified package {@code packageName} for a
+     * given {@code user}.
+     */
+    public void destroyAppDataSnapshot(String packageName, int user, long ceSnapshotInode) {
+        int storageFlags = Installer.FLAG_STORAGE_DE;
+        if (ceSnapshotInode > 0) {
+            storageFlags |= Installer.FLAG_STORAGE_CE;
+        }
+        try {
+            mInstaller.destroyAppDataSnapshot(packageName, user, ceSnapshotInode, storageFlags);
+        } catch (InstallerException ie) {
+            Log.e(TAG, "Unable to delete app data snapshot for " + packageName, ie);
+        }
+    }
+
+    /**
      * Computes the list of pending backups and restores for {@code userId} given lists of
      * available and recent rollbacks. Packages pending backup for the given user are added
      * to {@code pendingBackups} and packages pending restore are added to {@code pendingRestores}
@@ -191,16 +213,28 @@ public class AppDataRollbackHelper {
     }
 
     /**
-     * Commits the list of pending backups and restores for a given {@code userId}.
+     * Commits the list of pending backups and restores for a given {@code userId}. For the pending
+     * backups updates corresponding {@code changedRollbackData} with a mapping from {@code userId}
+     * to a inode of theirs CE user data snapshot.
      */
     public void commitPendingBackupAndRestoreForUser(int userId,
-            ArrayList<String> pendingBackups, Map<String, RestoreInfo> pendingRestores) {
+            ArrayList<String> pendingBackups, Map<String, RestoreInfo> pendingRestores,
+            List<RollbackData> changedRollbackData) {
         if (!pendingBackups.isEmpty()) {
             for (String packageName : pendingBackups) {
                 try {
-                    mInstaller.snapshotAppData(packageName, userId, Installer.FLAG_STORAGE_CE);
+                    long ceSnapshotInode = mInstaller.snapshotAppData(packageName, userId,
+                            Installer.FLAG_STORAGE_CE);
+                    for (RollbackData data : changedRollbackData) {
+                        for (PackageRollbackInfo info : data.packages) {
+                            if (info.getPackageName().equals(packageName)) {
+                                info.putCeSnapshotInode(userId, ceSnapshotInode);
+                            }
+                        }
+                    }
                 } catch (InstallerException ie) {
-                    Log.e(TAG, "Unable to create app data snapshot for: " + packageName, ie);
+                    Log.e(TAG, "Unable to create app data snapshot for: " + packageName
+                            + ", userId: " + userId, ie);
                 }
             }
         }
@@ -233,4 +267,26 @@ public class AppDataRollbackHelper {
         return StorageManager.isFileEncryptedNativeOrEmulated()
                 && !StorageManager.isUserKeyUnlocked(userId);
     }
+
+    /**
+     * Encapsulates a result of {@link #snapshotAppData} method.
+     */
+    public static final class SnapshotAppDataResult {
+
+        /**
+         * A list of users for which the snapshot is pending, usually because data for one or more
+         * users is still credential locked.
+         */
+        public final IntArray pendingBackups;
+
+        /**
+         * A mapping between user and an inode of theirs CE data snapshot.
+         */
+        public final SparseLongArray ceSnapshotInodes;
+
+        public SnapshotAppDataResult(IntArray pendingBackups, SparseLongArray ceSnapshotInodes) {
+            this.pendingBackups = pendingBackups;
+            this.ceSnapshotInodes = ceSnapshotInodes;
+        }
+    }
 }
diff --git a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
index 27f7bcf36ddb..24d5bd14bc44 100644
--- a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
+++ b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
@@ -43,6 +43,7 @@ import android.os.Process;
 import android.util.IntArray;
 import android.util.Log;
 import android.util.SparseBooleanArray;
+import android.util.SparseLongArray;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.server.LocalServices;
@@ -110,7 +111,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
     private final HandlerThread mHandlerThread;
     private final Installer mInstaller;
     private final RollbackPackageHealthObserver mPackageHealthObserver;
-    private final AppDataRollbackHelper mUserdataHelper;
+    private final AppDataRollbackHelper mAppDataRollbackHelper;
 
     RollbackManagerServiceImpl(Context context) {
         mContext = context;
@@ -124,7 +125,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
         mRollbackStore = new RollbackStore(new File(Environment.getDataDirectory(), "rollback"));
 
         mPackageHealthObserver = new RollbackPackageHealthObserver(mContext);
-        mUserdataHelper = new AppDataRollbackHelper(mInstaller);
+        mAppDataRollbackHelper = new AppDataRollbackHelper(mInstaller);
 
         // Kick off loading of the rollback data from strorage in a background
         // thread.
@@ -449,7 +450,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
                 for (PackageRollbackInfo info : data.packages) {
                     if (info.getPackageName().equals(packageName)) {
                         iter.remove();
-                        mRollbackStore.deleteAvailableRollback(data);
+                        deleteRollback(data);
                         break;
                     }
                 }
@@ -464,13 +465,13 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
             final List<RollbackData> changed;
             synchronized (mLock) {
                 ensureRollbackDataLoadedLocked();
-                changed = mUserdataHelper.computePendingBackupsAndRestores(userId,
+                changed = mAppDataRollbackHelper.computePendingBackupsAndRestores(userId,
                         pendingBackupPackages, pendingRestorePackages, mAvailableRollbacks,
                         mRecentlyExecutedRollbacks);
             }
 
-            mUserdataHelper.commitPendingBackupAndRestoreForUser(userId,
-                    pendingBackupPackages, pendingRestorePackages);
+            mAppDataRollbackHelper.commitPendingBackupAndRestoreForUser(userId,
+                    pendingBackupPackages, pendingRestorePackages, changed);
 
             for (RollbackData rd : changed) {
                 try {
@@ -520,7 +521,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
                         // mAvailableRollbacks, or is it okay to leave as
                         // unavailable until the next reboot when it will go
                         // away on its own?
-                        mRollbackStore.deleteAvailableRollback(data);
+                        deleteRollback(data);
                     }
                 }
             }
@@ -592,7 +593,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
                                         info.getVersionRolledBackFrom(),
                                         installedVersion)) {
                         iter.remove();
-                        mRollbackStore.deleteAvailableRollback(data);
+                        deleteRollback(data);
                         break;
                     }
                 }
@@ -705,7 +706,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
 
                 if (!now.isBefore(data.timestamp.plusMillis(ROLLBACK_LIFETIME_DURATION_MILLIS))) {
                     iter.remove();
-                    mRollbackStore.deleteAvailableRollback(data);
+                    deleteRollback(data);
                 } else if (oldest == null || oldest.isAfter(data.timestamp)) {
                     oldest = data.timestamp;
                 }
@@ -821,9 +822,13 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
             String packageName = newPackage.packageName;
             for (PackageRollbackInfo info : rd.packages) {
                 if (info.getPackageName().equals(packageName)) {
-                    IntArray pendingBackups = mUserdataHelper.snapshotAppData(
-                            packageName, installedUsers);
-                    info.getPendingBackups().addAll(pendingBackups);
+                    AppDataRollbackHelper.SnapshotAppDataResult rs =
+                            mAppDataRollbackHelper.snapshotAppData(packageName, installedUsers);
+                    info.getPendingBackups().addAll(rs.pendingBackups);
+                    for (int i = 0; i < rs.ceSnapshotInodes.size(); i++) {
+                        info.putCeSnapshotInode(rs.ceSnapshotInodes.keyAt(i),
+                                rs.ceSnapshotInodes.valueAt(i));
+                    }
                     try {
                         mRollbackStore.saveAvailableRollback(rd);
                     } catch (IOException ioe) {
@@ -892,13 +897,18 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
         VersionedPackage installedVersion = new VersionedPackage(packageName,
                 pkgInfo.getLongVersionCode());
 
-        IntArray pendingBackups = IntArray.wrap(new int[0]);
+        final AppDataRollbackHelper.SnapshotAppDataResult result;
         if (snapshotUserData && !isApex) {
-            pendingBackups = mUserdataHelper.snapshotAppData(packageName, installedUsers);
+            result = mAppDataRollbackHelper.snapshotAppData(packageName, installedUsers);
+        } else {
+            result = new AppDataRollbackHelper.SnapshotAppDataResult(IntArray.wrap(new int[0]),
+                new SparseLongArray());
         }
 
         PackageRollbackInfo info = new PackageRollbackInfo(newVersion, installedVersion,
-                pendingBackups, new ArrayList<>(), isApex);
+                result.pendingBackups, new ArrayList<>(), isApex, IntArray.wrap(installedUsers),
+                result.ceSnapshotInodes);
+
         RollbackData data;
         try {
             int childSessionId = session.getSessionId();
@@ -948,9 +958,8 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
         getHandler().post(() -> {
             final RollbackData rollbackData = getRollbackForPackage(packageName);
             for (int userId : userIds) {
-                final boolean changedRollbackData = mUserdataHelper.restoreAppData(packageName,
-                        rollbackData, userId, appId, ceDataInode, seInfo);
-
+                final boolean changedRollbackData = mAppDataRollbackHelper.restoreAppData(
+                        packageName, rollbackData, userId, appId, ceDataInode, seInfo);
                 // We've updated metadata about this rollback, so save it to flash.
                 if (changedRollbackData) {
                     try {
@@ -1142,12 +1151,12 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
                     scheduleExpiration(ROLLBACK_LIFETIME_DURATION_MILLIS);
                 } catch (IOException e) {
                     Log.e(TAG, "Unable to enable rollback", e);
-                    mRollbackStore.deleteAvailableRollback(data);
+                    deleteRollback(data);
                 }
             } else {
                 // The install session was aborted, clean up the pending
                 // install.
-                mRollbackStore.deleteAvailableRollback(data);
+                deleteRollback(data);
             }
         }
     }
@@ -1246,4 +1255,17 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
 
         throw new IOException("Failed to allocate rollback ID");
     }
+
+    private void deleteRollback(RollbackData rollbackData) {
+        for (PackageRollbackInfo info : rollbackData.packages) {
+            IntArray installedUsers = info.getInstalledUsers();
+            SparseLongArray ceSnapshotInodes = info.getCeSnapshotInodes();
+            for (int i = 0; i < installedUsers.size(); i++) {
+                int userId = installedUsers.get(i);
+                mAppDataRollbackHelper.destroyAppDataSnapshot(info.getPackageName(), userId,
+                        ceSnapshotInodes.get(userId, 0));
+            }
+        }
+        mRollbackStore.deleteAvailableRollback(rollbackData);
+    }
 }
diff --git a/services/core/java/com/android/server/rollback/RollbackStore.java b/services/core/java/com/android/server/rollback/RollbackStore.java
index d17ebaef0b6c..be904eacebff 100644
--- a/services/core/java/com/android/server/rollback/RollbackStore.java
+++ b/services/core/java/com/android/server/rollback/RollbackStore.java
@@ -23,6 +23,7 @@ import android.content.rollback.PackageRollbackInfo.RestoreInfo;
 import android.content.rollback.RollbackInfo;
 import android.util.IntArray;
 import android.util.Log;
+import android.util.SparseLongArray;
 
 import libcore.io.IoUtils;
 
@@ -160,6 +161,28 @@ class RollbackStore {
         return restoreInfos;
     }
 
+    private static @NonNull JSONArray ceSnapshotInodesToJson(
+            @NonNull SparseLongArray ceSnapshotInodes) throws JSONException {
+        JSONArray array = new JSONArray();
+        for (int i = 0; i < ceSnapshotInodes.size(); i++) {
+            JSONObject entryJson = new JSONObject();
+            entryJson.put("userId", ceSnapshotInodes.keyAt(i));
+            entryJson.put("ceSnapshotInode", ceSnapshotInodes.valueAt(i));
+            array.put(entryJson);
+        }
+        return array;
+    }
+
+    private static @NonNull SparseLongArray ceSnapshotInodesFromJson(JSONArray json)
+            throws JSONException {
+        SparseLongArray ceSnapshotInodes = new SparseLongArray(json.length());
+        for (int i = 0; i < json.length(); i++) {
+            JSONObject entry = json.getJSONObject(i);
+            ceSnapshotInodes.append(entry.getInt("userId"), entry.getLong("ceSnapshotInode"));
+        }
+        return ceSnapshotInodes;
+    }
+
     /**
      * Reads the list of recently executed rollbacks from persistent storage.
      */
@@ -263,8 +286,6 @@ class RollbackStore {
      * rollback.
      */
     void deleteAvailableRollback(RollbackData data) {
-        // TODO(narayan): Make sure we delete the userdata snapshot along with the backup of the
-        // actual app.
         removeFile(data.backupDir);
     }
 
@@ -341,11 +362,15 @@ class RollbackStore {
 
         IntArray pendingBackups = info.getPendingBackups();
         List<RestoreInfo> pendingRestores = info.getPendingRestores();
+        IntArray installedUsers = info.getInstalledUsers();
         json.put("pendingBackups", convertToJsonArray(pendingBackups));
         json.put("pendingRestores", convertToJsonArray(pendingRestores));
 
         json.put("isApex", info.isApex());
 
+        json.put("installedUsers", convertToJsonArray(installedUsers));
+        json.put("ceSnapshotInodes", ceSnapshotInodesToJson(info.getCeSnapshotInodes()));
+
         return json;
     }
 
@@ -362,8 +387,12 @@ class RollbackStore {
 
         final boolean isApex = json.getBoolean("isApex");
 
+        final IntArray installedUsers = convertToIntArray(json.getJSONArray("installedUsers"));
+        final SparseLongArray ceSnapshotInodes = ceSnapshotInodesFromJson(
+                json.getJSONArray("ceSnapshotInodes"));
+
         return new PackageRollbackInfo(versionRolledBackFrom, versionRolledBackTo,
-                pendingBackups, pendingRestores, isApex);
+                pendingBackups, pendingRestores, isApex, installedUsers, ceSnapshotInodes);
     }
 
     private 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 43e3eb0193b1..50dbaf570e9e 100644
--- a/services/tests/servicestests/src/com/android/server/rollback/AppDataRollbackHelperTest.java
+++ b/services/tests/servicestests/src/com/android/server/rollback/AppDataRollbackHelperTest.java
@@ -18,15 +18,19 @@ package com.android.server.rollback;
 
 import static org.junit.Assert.*;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Matchers.anyInt;
+import static org.mockito.Matchers.anyString;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verifyZeroInteractions;
+import static org.mockito.Mockito.when;
 
 import android.content.pm.VersionedPackage;
 import android.content.rollback.PackageRollbackInfo;
 import android.content.rollback.PackageRollbackInfo.RestoreInfo;
 import android.util.IntArray;
+import android.util.SparseLongArray;
 
 import com.android.server.pm.Installer;
 
@@ -38,6 +42,7 @@ import org.mockito.Mockito;
 
 import java.io.File;
 import java.util.ArrayList;
+import java.util.HashMap;
 
 @RunWith(JUnit4.class)
 public class AppDataRollbackHelperTest {
@@ -50,10 +55,14 @@ public class AppDataRollbackHelperTest {
         // 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));
-        IntArray pending = helper.snapshotAppData("com.foo.bar", new int[]{10, 11});
-        assertEquals(2, pending.size());
-        assertEquals(10, pending.get(0));
-        assertEquals(11, pending.get(1));
+        AppDataRollbackHelper.SnapshotAppDataResult result = helper.snapshotAppData("com.foo.bar",
+                new int[]{10, 11});
+
+        assertEquals(2, result.pendingBackups.size());
+        assertEquals(10, result.pendingBackups.get(0));
+        assertEquals(11, result.pendingBackups.get(1));
+
+        assertEquals(0, result.ceSnapshotInodes.size());
 
         InOrder inOrder = Mockito.inOrder(installer);
         inOrder.verify(installer).snapshotAppData(
@@ -65,10 +74,14 @@ public class AppDataRollbackHelperTest {
         // One of the users is unlocked but the other isn't
         doReturn(false).when(helper).isUserCredentialLocked(eq(10));
         doReturn(true).when(helper).isUserCredentialLocked(eq(11));
+        when(installer.snapshotAppData(anyString(), anyInt(), anyInt())).thenReturn(239L);
 
-        pending = helper.snapshotAppData("com.foo.bar", new int[]{10, 11});
-        assertEquals(1, pending.size());
-        assertEquals(11, pending.get(0));
+        result = helper.snapshotAppData("com.foo.bar", new int[]{10, 11});
+        assertEquals(1, result.pendingBackups.size());
+        assertEquals(11, result.pendingBackups.get(0));
+
+        assertEquals(1, result.ceSnapshotInodes.size());
+        assertEquals(239L, result.ceSnapshotInodes.get(10));
 
         inOrder = Mockito.inOrder(installer);
         inOrder.verify(installer).snapshotAppData(
@@ -83,7 +96,7 @@ public class AppDataRollbackHelperTest {
         RollbackData data = new RollbackData(1, new File("/does/not/exist"), -1, true);
         data.packages.add(new PackageRollbackInfo(
                 new VersionedPackage(packageName, 1), new VersionedPackage(packageName, 1),
-                new IntArray(), new ArrayList<>(), false));
+                new IntArray(), new ArrayList<>(), false, new IntArray(), new SparseLongArray()));
         data.inProgress = true;
 
         return data;
@@ -173,4 +186,53 @@ public class AppDataRollbackHelperTest {
         ArrayList<RestoreInfo> pendingRestores = rd.packages.get(0).getPendingRestores();
         assertEquals(0, pendingRestores.size());
     }
+
+    @Test
+    public void destroyAppData() throws Exception {
+        Installer installer = mock(Installer.class);
+        AppDataRollbackHelper helper = new AppDataRollbackHelper(installer);
+        SparseLongArray ceSnapshotInodes = new SparseLongArray();
+        ceSnapshotInodes.put(11, 239L);
+
+        helper.destroyAppDataSnapshot("com.foo.bar", 10, 0L);
+        helper.destroyAppDataSnapshot("com.foo.bar", 11, 239L);
+
+        InOrder inOrder = Mockito.inOrder(installer);
+        inOrder.verify(installer).destroyAppDataSnapshot(
+                eq("com.foo.bar"), eq(10), eq(0L),
+                eq(Installer.FLAG_STORAGE_DE));
+        inOrder.verify(installer).destroyAppDataSnapshot(
+                eq("com.foo.bar"), eq(11), eq(239L),
+                eq(Installer.FLAG_STORAGE_DE | Installer.FLAG_STORAGE_CE));
+        inOrder.verifyNoMoreInteractions();
+    }
+
+    @Test
+    public void commitPendingBackupAndRestoreForUser_updatesRollbackData() throws Exception {
+        Installer installer = mock(Installer.class);
+        AppDataRollbackHelper helper = new AppDataRollbackHelper(installer);
+
+        ArrayList<RollbackData> changedRollbackData = new ArrayList<>();
+        changedRollbackData.add(createInProgressRollbackData("com.foo.bar"));
+
+        when(installer.snapshotAppData(anyString(), anyInt(), anyInt())).thenReturn(239L);
+
+        ArrayList<String> pendingBackups = new ArrayList<>();
+        pendingBackups.add("com.foo.bar");
+
+        helper.commitPendingBackupAndRestoreForUser(11, pendingBackups,
+                new HashMap<>() /* pendingRestores */, changedRollbackData);
+
+        assertEquals(1, changedRollbackData.size());
+        assertEquals(1, changedRollbackData.get(0).packages.size());
+        PackageRollbackInfo info = changedRollbackData.get(0).packages.get(0);
+
+        assertEquals(1, info.getCeSnapshotInodes().size());
+        assertEquals(239L, info.getCeSnapshotInodes().get(11));
+
+        InOrder inOrder = Mockito.inOrder(installer);
+        inOrder.verify(installer).snapshotAppData("com.foo.bar", 11 /* userId */,
+                Installer.FLAG_STORAGE_CE);
+        inOrder.verifyNoMoreInteractions();
+    }
 }