Multi-user support for app data snapshot / rollback.

This change implements support in the RollbackManagerService for
keeping track of pending backups and restores for users that are
credential unlocked, and subsequently performing after user
credentials are available.

All code pertaining to app data snapshot / rollback has been factored
out in AppDataRollbackHelper to allow for better mockability and
testing, given that integration tests for these classes are hard
to write. Some of the unit-tests for AppDataRollbackHelper contravene
good testing practices, but in combination with RollbackTest, they
exercise most code related to multi user scenarios.

Test: RollbackTest, AppDataRollbackHelperTest
Test: manual

Change-Id: I8debc24c9318cfd22e8c278dde7b5aee6688f3e6

6 files changed, 644 insertions, 58 deletions

diff --git a/core/java/android/content/rollback/PackageRollbackInfo.java b/core/java/android/content/rollback/PackageRollbackInfo.java
index 4644a83de462..d4ed35a9cb71 100644
--- a/core/java/android/content/rollback/PackageRollbackInfo.java
+++ b/core/java/android/content/rollback/PackageRollbackInfo.java
@@ -16,10 +16,14 @@
 
 package android.content.rollback;
 
+import android.annotation.NonNull;
 import android.annotation.SystemApi;
 import android.content.pm.VersionedPackage;
 import android.os.Parcel;
 import android.os.Parcelable;
+import android.util.IntArray;
+
+import java.util.ArrayList;
 
 /**
  * Information about a rollback available for a particular package.
@@ -33,6 +37,38 @@ public final class PackageRollbackInfo implements Parcelable {
     private final VersionedPackage mVersionRolledBackTo;
 
     /**
+     * Encapsulates information required to restore a snapshot of an app's userdata.
+     *
+     * @hide
+     */
+    public static class RestoreInfo {
+        public final int userId;
+        public final int appId;
+        public final String seInfo;
+
+        public RestoreInfo(int userId, int appId, String seInfo) {
+            this.userId = userId;
+            this.appId = appId;
+            this.seInfo = seInfo;
+        }
+    }
+
+    /*
+     * The list of users for which we need to backup userdata for this package. Backups of
+     * credential encrypted data are listed as pending if the user hasn't unlocked their device
+     * with credentials yet.
+     */
+    // NOTE: Not a part of the Parcelable representation of this object.
+    private final IntArray mPendingBackups;
+
+    /**
+     * The list of users for which we need to restore userdata for this package. This field is
+     * non-null only after a rollback for this package has been committed.
+     */
+    // NOTE: Not a part of the Parcelable representation of this object.
+    private final ArrayList<RestoreInfo> mPendingRestores;
+
+    /**
      * Returns the name of the package to roll back from.
      */
     public String getPackageName() {
@@ -54,15 +90,46 @@ public final class PackageRollbackInfo implements Parcelable {
     }
 
     /** @hide */
+    public IntArray getPendingBackups() {
+        return mPendingBackups;
+    }
+
+    /** @hide */
+    public ArrayList<RestoreInfo> getPendingRestores() {
+        return mPendingRestores;
+    }
+
+    /** @hide */
+    public RestoreInfo getRestoreInfo(int userId) {
+        for (RestoreInfo ri : mPendingRestores) {
+            if (ri.userId == userId) {
+                return ri;
+            }
+        }
+
+        return null;
+    }
+
+    /** @hide */
+    public void removeRestoreInfo(RestoreInfo ri) {
+        mPendingRestores.remove(ri);
+    }
+
+    /** @hide */
     public PackageRollbackInfo(VersionedPackage packageRolledBackFrom,
-            VersionedPackage packageRolledBackTo) {
+            VersionedPackage packageRolledBackTo,
+            @NonNull IntArray pendingBackups, @NonNull ArrayList<RestoreInfo> pendingRestores) {
         this.mVersionRolledBackFrom = packageRolledBackFrom;
         this.mVersionRolledBackTo = packageRolledBackTo;
+        this.mPendingBackups = pendingBackups;
+        this.mPendingRestores = pendingRestores;
     }
 
     private PackageRollbackInfo(Parcel in) {
         this.mVersionRolledBackFrom = VersionedPackage.CREATOR.createFromParcel(in);
         this.mVersionRolledBackTo = VersionedPackage.CREATOR.createFromParcel(in);
+        this.mPendingRestores = null;
+        this.mPendingBackups = null;
     }
 
     @Override
diff --git a/services/core/java/com/android/server/rollback/AppDataRollbackHelper.java b/services/core/java/com/android/server/rollback/AppDataRollbackHelper.java
new file mode 100644
index 000000000000..8dd076028b43
--- /dev/null
+++ b/services/core/java/com/android/server/rollback/AppDataRollbackHelper.java
@@ -0,0 +1,236 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.rollback;
+
+import android.content.rollback.PackageRollbackInfo;
+import android.content.rollback.PackageRollbackInfo.RestoreInfo;
+import android.content.rollback.RollbackInfo;
+import android.os.storage.StorageManager;
+import android.util.IntArray;
+import android.util.Log;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.server.pm.Installer;
+import com.android.server.pm.Installer.InstallerException;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Encapsulates the logic for initiating userdata snapshots and rollbacks via installd.
+ */
+@VisibleForTesting
+// TODO(narayan): Reason about the failure scenarios that involve one or more IPCs to installd
+// failing. We need to decide what course of action to take if calls to snapshotAppData or
+// restoreAppDataSnapshot fail.
+public class AppDataRollbackHelper {
+    private static final String TAG = "RollbackManager";
+
+    private final Installer mInstaller;
+
+    public AppDataRollbackHelper(Installer installer) {
+        mInstaller = installer;
+    }
+
+    /**
+     * 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.
+     */
+    public IntArray snapshotAppData(String packageName, int[] installedUsers) {
+        final IntArray pendingBackups = new IntArray();
+        for (int user : installedUsers) {
+            final int storageFlags;
+            if (isUserCredentialLocked(user)) {
+                // We've encountered a user that hasn't unlocked on a FBE device, so we can't copy
+                // across app user data until the user unlocks their device.
+                Log.v(TAG, "User: " + user + " isn't unlocked, skipping CE userdata backup.");
+                storageFlags = Installer.FLAG_STORAGE_DE;
+                pendingBackups.add(user);
+            } else {
+                storageFlags = Installer.FLAG_STORAGE_CE | Installer.FLAG_STORAGE_DE;
+            }
+
+            try {
+                mInstaller.snapshotAppData(packageName, user, storageFlags);
+            } catch (InstallerException ie) {
+                Log.e(TAG, "Unable to create app data snapshot for: " + packageName
+                        + ", userId: " + user, ie);
+            }
+        }
+
+        return pendingBackups;
+    }
+
+    /**
+     * Restores an app data snapshot for a specified package ({@code packageName},
+     * {@code rollbackData}) for a specified {@code userId}.
+     *
+     * @return {@code true} iff. a change to the {@code rollbackData} has been made. Changes to
+     *         {@code rollbackData} are restricted to the removal or addition of {@code userId} to
+     *         the list of pending backups or restores.
+     */
+    public boolean restoreAppData(String packageName, RollbackData rollbackData,
+            int userId, int appId, long ceDataInode, String seInfo) {
+        if (rollbackData == null) {
+            return false;
+        }
+
+        if (!rollbackData.inProgress) {
+            Log.e(TAG, "Request to restore userData for: " + packageName
+                    + ", but no rollback in progress.");
+            return false;
+        }
+
+        PackageRollbackInfo packageInfo = RollbackManagerServiceImpl.getPackageRollbackInfo(
+                rollbackData, packageName);
+        int storageFlags = Installer.FLAG_STORAGE_DE;
+
+        final IntArray pendingBackups = packageInfo.getPendingBackups();
+        final List<RestoreInfo> pendingRestores = packageInfo.getPendingRestores();
+        boolean changedRollbackData = false;
+
+        // If we still have a userdata backup pending for this user, it implies that the user
+        // hasn't unlocked their device between the point of backup and the point of restore,
+        // so the data cannot have changed. We simply skip restoring CE data in this case.
+        if (pendingBackups != null && pendingBackups.indexOf(userId) != -1) {
+            pendingBackups.remove(pendingBackups.indexOf(userId));
+            changedRollbackData = true;
+        } else {
+            // There's no pending CE backup for this user, which means that we successfully
+            // managed to backup data for the user, which means we seek to restore it
+            if (isUserCredentialLocked(userId)) {
+                // We've encountered a user that hasn't unlocked on a FBE device, so we can't
+                // copy across app user data until the user unlocks their device.
+                pendingRestores.add(new RestoreInfo(userId, appId, seInfo));
+                changedRollbackData = true;
+            } else {
+                // This user has unlocked, we can proceed to restore both CE and DE data.
+                storageFlags = storageFlags | Installer.FLAG_STORAGE_CE;
+            }
+        }
+
+        try {
+            mInstaller.restoreAppDataSnapshot(packageName, appId, ceDataInode,
+                    seInfo, userId, storageFlags);
+        } catch (InstallerException ie) {
+            Log.e(TAG, "Unable to restore app data snapshot: " + packageName, ie);
+        }
+
+        return changedRollbackData;
+    }
+
+    /**
+     * 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}
+     * along with their corresponding {@code RestoreInfo}.
+     *
+     * @return the list of {@code RollbackData} that have been modified during this computation.
+     */
+    public List<RollbackData> computePendingBackupsAndRestores(int userId,
+            ArrayList<String> pendingBackupPackages, Map<String, RestoreInfo> pendingRestores,
+            List<RollbackData> availableRollbacks, List<RollbackInfo> recentRollbacks) {
+        List<RollbackData> rd = new ArrayList<>();
+        // First check with the list of available rollbacks to see whether there are any
+        // pending backup operations that we've not managed to execute.
+        for (RollbackData data : availableRollbacks) {
+            for (PackageRollbackInfo info : data.packages) {
+                final IntArray pendingBackupUsers = info.getPendingBackups();
+                if (pendingBackupUsers != null) {
+                    final int idx = pendingBackupUsers.indexOf(userId);
+                    if (idx != -1) {
+                        pendingBackupPackages.add(info.getPackageName());
+                        pendingBackupUsers.remove(idx);
+                        if (rd.indexOf(data) == -1) {
+                            rd.add(data);
+                        }
+                    }
+                }
+            }
+        }
+
+        // Then check with the list of recently executed rollbacks to see whether there are
+        // any rollback operations
+        for (RollbackInfo data : recentRollbacks) {
+            for (PackageRollbackInfo info : data.getPackages()) {
+                final RestoreInfo ri = info.getRestoreInfo(userId);
+                if (ri != null) {
+                    if (pendingBackupPackages.contains(info.getPackageName())) {
+                        // This implies that the user hasn't unlocked their device between
+                        // the request to backup data for this user and the request to restore
+                        // it, so we do nothing here.
+                        pendingBackupPackages.remove(info.getPackageName());
+                    } else {
+                        pendingRestores.put(info.getPackageName(), ri);
+                    }
+
+                    info.removeRestoreInfo(ri);
+                }
+            }
+        }
+
+        return rd;
+    }
+
+    /**
+     * Commits the list of pending backups and restores for a given {@code userId}.
+     */
+    public void commitPendingBackupAndRestoreForUser(int userId,
+            ArrayList<String> pendingBackups, Map<String, RestoreInfo> pendingRestores) {
+        if (!pendingBackups.isEmpty()) {
+            for (String packageName : pendingBackups) {
+                try {
+                    mInstaller.snapshotAppData(packageName, userId, Installer.FLAG_STORAGE_CE);
+                } catch (InstallerException ie) {
+                    Log.e(TAG, "Unable to create app data snapshot for: " + packageName, ie);
+                }
+            }
+        }
+
+        // TODO(narayan): Should we perform the restore before the backup for packages that have
+        // both backups and restores pending ? We could get into this case if we have a pending
+        // restore from a rollback + a snapshot request from a new restore.
+        if (!pendingRestores.isEmpty()) {
+            for (String packageName : pendingRestores.keySet()) {
+                try {
+                    final RestoreInfo ri = pendingRestores.get(packageName);
+
+                    // TODO(narayan): Verify that the user of "0" for ceDataInode is accurate
+                    // here. We know that the user has unlocked (and that their CE data is
+                    // available) so we shouldn't need to resort to the fallback path.
+                    mInstaller.restoreAppDataSnapshot(packageName, ri.appId,
+                            0 /* ceDataInode */, ri.seInfo, userId, Installer.FLAG_STORAGE_CE);
+                } catch (InstallerException ie) {
+                    Log.e(TAG, "Unable to restore app data snapshot for: " + packageName, ie);
+                }
+            }
+        }
+    }
+
+    /**
+     * @return {@code true} iff. {@code userId} is locked on an FBE device.
+     */
+    @VisibleForTesting
+    public boolean isUserCredentialLocked(int userId) {
+        return StorageManager.isFileEncryptedNativeOrEmulated()
+                && !StorageManager.isUserKeyUnlocked(userId);
+    }
+}
diff --git a/services/core/java/com/android/server/rollback/RollbackManagerService.java b/services/core/java/com/android/server/rollback/RollbackManagerService.java
index 4b5e76497f62..ba6cdddb4cc8 100644
--- a/services/core/java/com/android/server/rollback/RollbackManagerService.java
+++ b/services/core/java/com/android/server/rollback/RollbackManagerService.java
@@ -39,4 +39,9 @@ public final class RollbackManagerService extends SystemService {
         mService = new RollbackManagerServiceImpl(getContext());
         publishBinderService(Context.ROLLBACK_SERVICE, mService);
     }
+
+    @Override
+    public void onUnlockUser(int user) {
+        mService.onUnlockUser(user);
+    }
 }
diff --git a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
index 289618eb133c..5eb137b4c50c 100644
--- a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
+++ b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
@@ -31,6 +31,7 @@ import android.content.pm.ParceledListSlice;
 import android.content.pm.VersionedPackage;
 import android.content.rollback.IRollbackManager;
 import android.content.rollback.PackageRollbackInfo;
+import android.content.rollback.PackageRollbackInfo.RestoreInfo;
 import android.content.rollback.RollbackInfo;
 import android.content.rollback.RollbackManager;
 import android.os.Binder;
@@ -39,14 +40,13 @@ import android.os.Handler;
 import android.os.HandlerThread;
 import android.os.ParcelFileDescriptor;
 import android.os.Process;
-import android.os.storage.StorageManager;
+import android.util.IntArray;
 import android.util.Log;
 import android.util.SparseBooleanArray;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.server.LocalServices;
 import com.android.server.pm.Installer;
-import com.android.server.pm.Installer.InstallerException;
 import com.android.server.pm.PackageManagerServiceUtils;
 
 import java.io.File;
@@ -111,6 +111,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
     private final HandlerThread mHandlerThread;
     private final Installer mInstaller;
     private final RollbackPackageHealthObserver mPackageHealthObserver;
+    private final AppDataRollbackHelper mUserdataHelper;
 
     RollbackManagerServiceImpl(Context context) {
         mContext = context;
@@ -124,6 +125,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
         mRollbackStore = new RollbackStore(new File(Environment.getDataDirectory(), "rollback"));
 
         mPackageHealthObserver = new RollbackPackageHealthObserver(mContext);
+        mUserdataHelper = new AppDataRollbackHelper(mInstaller);
 
         // Kick off loading of the rollback data from strorage in a background
         // thread.
@@ -424,6 +426,35 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
         }
     }
 
+    void onUnlockUser(int userId) {
+        getHandler().post(() -> {
+            final ArrayList<String> pendingBackupPackages = new ArrayList<>();
+            final Map<String, RestoreInfo> pendingRestorePackages = new HashMap<>();
+            final List<RollbackData> changed;
+            synchronized (mLock) {
+                ensureRollbackDataLoadedLocked();
+                changed = mUserdataHelper.computePendingBackupsAndRestores(userId,
+                        pendingBackupPackages, pendingRestorePackages, mAvailableRollbacks,
+                        mRecentlyExecutedRollbacks);
+            }
+
+            mUserdataHelper.commitPendingBackupAndRestoreForUser(userId,
+                    pendingBackupPackages, pendingRestorePackages);
+
+            for (RollbackData rd : changed) {
+                try {
+                    mRollbackStore.saveAvailableRollback(rd);
+                } catch (IOException ioe) {
+                    Log.e(TAG, "Unable to save rollback info for : " + rd.rollbackId, ioe);
+                }
+            }
+
+            synchronized (mLock) {
+                mRollbackStore.saveRecentlyExecutedRollbacks(mRecentlyExecutedRollbacks);
+            }
+        });
+    }
+
     /**
      * Load rollback data from storage if it has not already been loaded.
      * After calling this funciton, mAvailableRollbacks and
@@ -533,6 +564,20 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
         // that are necessary to keep track of.
         synchronized (mLock) {
             ensureRollbackDataLoadedLocked();
+
+            // This should never happen because we can't have any pending backups left after
+            // a rollback has been executed. See AppDataRollbackHelper#restoreAppData where we
+            // clear all pending backups at the point of restore because they're guaranteed to be
+            // no-ops.
+            //
+            // We may, however, have one or more pending restores left to handle.
+            for (PackageRollbackInfo target : rollback.getPackages()) {
+                if (target.getPendingBackups().size() > 0) {
+                    Log.e(TAG, "No backups allowed to be pending for: " + target);
+                    target.getPendingBackups().clear();
+                }
+            }
+
             mRecentlyExecutedRollbacks.add(rollback);
             mRollbackStore.saveRecentlyExecutedRollbacks(mRecentlyExecutedRollbacks);
         }
@@ -701,27 +746,12 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
         VersionedPackage installedVersion = new VersionedPackage(packageName,
                 installedPackage.getLongVersionCode());
 
-        for (int user : installedUsers) {
-            final int storageFlags;
-            if (StorageManager.isFileEncryptedNativeOrEmulated()
-                    && !StorageManager.isUserKeyUnlocked(user)) {
-                // We've encountered a user that hasn't unlocked on a FBE device, so we can't copy
-                // across app user data until the user unlocks their device.
-                Log.e(TAG, "User: " + user + " isn't unlocked, skipping CE userdata backup.");
-                storageFlags = Installer.FLAG_STORAGE_DE;
-            } else {
-                storageFlags = Installer.FLAG_STORAGE_CE | Installer.FLAG_STORAGE_DE;
-            }
-
-            try {
-                mInstaller.snapshotAppData(packageName, user, storageFlags);
-            } catch (InstallerException ie) {
-                Log.e(TAG, "Unable to create app data snapshot for: " + packageName, ie);
-            }
-        }
 
-        PackageRollbackInfo info = new PackageRollbackInfo(newVersion, installedVersion);
+        final IntArray pendingBackups = mUserdataHelper.snapshotAppData(packageName,
+                installedUsers);
 
+        PackageRollbackInfo info = new PackageRollbackInfo(newVersion, installedVersion,
+                pendingBackups, new ArrayList<>());
         RollbackData data;
         try {
             synchronized (mLock) {
@@ -760,40 +790,24 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
         }
 
         getHandler().post(() -> {
-            PackageManagerInternal pmi = LocalServices.getService(PackageManagerInternal.class);
             final RollbackData rollbackData = getRollbackForPackage(packageName);
-            if (rollbackData == null) {
-                pmi.finishPackageInstall(token, false);
-                return;
-            }
-
-            if (!rollbackData.inProgress) {
-                Log.e(TAG, "Request to restore userData for: " + packageName
-                        + ", but no rollback in progress.");
-                pmi.finishPackageInstall(token, false);
-                return;
-            }
-
-            final int storageFlags;
-            if (StorageManager.isFileEncryptedNativeOrEmulated()
-                    && !StorageManager.isUserKeyUnlocked(userId)) {
-                // We've encountered a user that hasn't unlocked on a FBE device, so we can't copy
-                // across app user data until the user unlocks their device.
-                Log.e(TAG, "User: " + userId + " isn't unlocked, skipping CE userdata restore.");
-
-                storageFlags = Installer.FLAG_STORAGE_DE;
-            } else {
-                storageFlags = Installer.FLAG_STORAGE_CE | Installer.FLAG_STORAGE_DE;
-            }
+            final boolean changedRollbackData = mUserdataHelper.restoreAppData(packageName,
+                    rollbackData, userId, appId, ceDataInode, seInfo);
+            final PackageManagerInternal pmi = LocalServices.getService(
+                    PackageManagerInternal.class);
+            pmi.finishPackageInstall(token, false);
 
-            try {
-                mInstaller.restoreAppDataSnapshot(packageName, appId, ceDataInode,
-                        seInfo, userId, storageFlags);
-            } catch (InstallerException ie) {
-                Log.e(TAG, "Unable to restore app data snapshot: " + packageName, ie);
+            // We've updated metadata about this rollback, so save it to flash.
+            if (changedRollbackData) {
+                try {
+                    mRollbackStore.saveAvailableRollback(rollbackData);
+                } catch (IOException ioe) {
+                    // TODO(narayan): What is the right thing to do here ? This isn't a fatal error,
+                    // since it will only result in us trying to restore data again, which will be
+                    // a no-op if there's no data available.
+                    Log.e(TAG, "Unable to save available rollback: " + packageName, ioe);
+                }
             }
-
-            pmi.finishPackageInstall(token, false);
         });
     }
 
@@ -900,10 +914,8 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
             ensureRollbackDataLoadedLocked();
             for (int i = 0; i < mAvailableRollbacks.size(); ++i) {
                 RollbackData data = mAvailableRollbacks.get(i);
-                for (PackageRollbackInfo info : data.packages) {
-                    if (info.getPackageName().equals(packageName)) {
-                        return data;
-                    }
+                if (getPackageRollbackInfo(data, packageName) != null) {
+                    return data;
                 }
             }
         }
@@ -926,6 +938,22 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
                 }
             }
         }
+
+        return null;
+    }
+
+    /**
+     * Returns the {@code PackageRollbackInfo} associated with {@code packageName} from
+     * a specified {@code RollbackData}.
+     */
+    static PackageRollbackInfo getPackageRollbackInfo(RollbackData data,
+            String packageName) {
+        for (PackageRollbackInfo info : data.packages) {
+            if (info.getPackageName().equals(packageName)) {
+                return info;
+            }
+        }
+
         return null;
     }
 
diff --git a/services/core/java/com/android/server/rollback/RollbackStore.java b/services/core/java/com/android/server/rollback/RollbackStore.java
index 98ebb0943f43..c70f47dc5ed3 100644
--- a/services/core/java/com/android/server/rollback/RollbackStore.java
+++ b/services/core/java/com/android/server/rollback/RollbackStore.java
@@ -16,9 +16,12 @@
 
 package com.android.server.rollback;
 
+import android.annotation.NonNull;
 import android.content.pm.VersionedPackage;
 import android.content.rollback.PackageRollbackInfo;
+import android.content.rollback.PackageRollbackInfo.RestoreInfo;
 import android.content.rollback.RollbackInfo;
+import android.util.IntArray;
 import android.util.Log;
 
 import libcore.io.IoUtils;
@@ -99,6 +102,64 @@ class RollbackStore {
     }
 
     /**
+     * Converts an {@code JSONArray} of integers to an {@code IntArray}.
+     */
+    private static @NonNull IntArray convertToIntArray(@NonNull JSONArray jsonArray)
+            throws JSONException {
+        if (jsonArray.length() == 0) {
+            return new IntArray();
+        }
+
+        final int[] ret = new int[jsonArray.length()];
+        for (int i = 0; i < ret.length; ++i) {
+            ret[i] = jsonArray.getInt(i);
+        }
+
+        return IntArray.wrap(ret);
+    }
+
+    /**
+     * Converts an {@code IntArray} into an {@code JSONArray} of integers.
+     */
+    private static @NonNull JSONArray convertToJsonArray(@NonNull IntArray intArray) {
+        JSONArray jsonArray = new JSONArray();
+        for (int i = 0; i < intArray.size(); ++i) {
+            jsonArray.put(intArray.get(i));
+        }
+
+        return jsonArray;
+    }
+
+    private static @NonNull JSONArray convertToJsonArray(@NonNull List<RestoreInfo> list)
+            throws JSONException {
+        JSONArray jsonArray = new JSONArray();
+        for (RestoreInfo ri : list) {
+            JSONObject jo = new JSONObject();
+            jo.put("userId", ri.userId);
+            jo.put("appId", ri.appId);
+            jo.put("seInfo", ri.seInfo);
+            jsonArray.put(jo);
+        }
+
+        return jsonArray;
+    }
+
+    private static @NonNull ArrayList<RestoreInfo> convertToRestoreInfoArray(
+            @NonNull JSONArray array) throws JSONException {
+        ArrayList<RestoreInfo> restoreInfos = new ArrayList<>();
+
+        for (int i = 0; i < array.length(); ++i) {
+            JSONObject jo = array.getJSONObject(i);
+            restoreInfos.add(new RestoreInfo(
+                    jo.getInt("userId"),
+                    jo.getInt("appId"),
+                    jo.getString("seInfo")));
+        }
+
+        return restoreInfos;
+    }
+
+    /**
      * Reads the list of recently executed rollbacks from persistent storage.
      */
     List<RollbackInfo> loadRecentlyExecutedRollbacks() {
@@ -239,6 +300,12 @@ class RollbackStore {
         JSONObject json = new JSONObject();
         json.put("versionRolledBackFrom", toJson(info.getVersionRolledBackFrom()));
         json.put("versionRolledBackTo", toJson(info.getVersionRolledBackTo()));
+
+        IntArray pendingBackups = info.getPendingBackups();
+        List<RestoreInfo> pendingRestores = info.getPendingRestores();
+        json.put("pendingBackups", convertToJsonArray(pendingBackups));
+        json.put("pendingRestores", convertToJsonArray(pendingRestores));
+
         return json;
     }
 
@@ -247,7 +314,14 @@ class RollbackStore {
                 json.getJSONObject("versionRolledBackFrom"));
         VersionedPackage versionRolledBackTo = versionedPackageFromJson(
                 json.getJSONObject("versionRolledBackTo"));
-        return new PackageRollbackInfo(versionRolledBackFrom, versionRolledBackTo);
+
+        final IntArray pendingBackups = convertToIntArray(
+                json.getJSONArray("pendingBackups"));
+        final ArrayList<RestoreInfo> pendingRestores = convertToRestoreInfoArray(
+                json.getJSONArray("pendingRestores"));
+
+        return new PackageRollbackInfo(versionRolledBackFrom, versionRolledBackTo,
+                pendingBackups, pendingRestores);
     }
 
     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
new file mode 100644
index 000000000000..33cbf7ae8ba6
--- /dev/null
+++ b/services/tests/servicestests/src/com/android/server/rollback/AppDataRollbackHelperTest.java
@@ -0,0 +1,176 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.rollback;
+
+import static org.junit.Assert.*;
+import static org.mockito.ArgumentMatchers.eq;
+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 android.content.pm.VersionedPackage;
+import android.content.rollback.PackageRollbackInfo;
+import android.content.rollback.PackageRollbackInfo.RestoreInfo;
+import android.util.IntArray;
+
+import com.android.server.pm.Installer;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.InOrder;
+import org.mockito.Mockito;
+
+import java.io.File;
+import java.util.ArrayList;
+
+@RunWith(JUnit4.class)
+public class AppDataRollbackHelperTest {
+
+    @Test
+    public void testSnapshotAppData() throws Exception {
+        Installer installer = mock(Installer.class);
+        AppDataRollbackHelper helper = spy(new AppDataRollbackHelper(installer));
+
+        // 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));
+
+        InOrder inOrder = Mockito.inOrder(installer);
+        inOrder.verify(installer).snapshotAppData(
+                eq("com.foo.bar"), eq(10), eq(Installer.FLAG_STORAGE_DE));
+        inOrder.verify(installer).snapshotAppData(
+                eq("com.foo.bar"), eq(11), eq(Installer.FLAG_STORAGE_DE));
+        inOrder.verifyNoMoreInteractions();
+
+        // 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));
+
+        pending = helper.snapshotAppData("com.foo.bar", new int[]{10, 11});
+        assertEquals(1, pending.size());
+        assertEquals(11, pending.get(0));
+
+        inOrder = Mockito.inOrder(installer);
+        inOrder.verify(installer).snapshotAppData(
+                eq("com.foo.bar"), eq(10),
+                eq(Installer.FLAG_STORAGE_CE | Installer.FLAG_STORAGE_DE));
+        inOrder.verify(installer).snapshotAppData(
+                eq("com.foo.bar"), eq(11), eq(Installer.FLAG_STORAGE_DE));
+        inOrder.verifyNoMoreInteractions();
+    }
+
+    private static RollbackData createInProgressRollbackData(String packageName) {
+        RollbackData data = new RollbackData(1, new File("/does/not/exist"));
+        data.packages.add(new PackageRollbackInfo(
+                new VersionedPackage(packageName, 1), new VersionedPackage(packageName, 1),
+                new IntArray(), new ArrayList<>()));
+        data.inProgress = true;
+
+        return data;
+    }
+
+    @Test
+    public void testRestoreAppDataSnapshot_noRollbackData() throws Exception {
+        Installer installer = mock(Installer.class);
+        AppDataRollbackHelper helper = spy(new AppDataRollbackHelper(installer));
+
+        assertFalse(helper.restoreAppData("com.foo", null, 0, 0, 0, "seinfo"));
+        verifyZeroInteractions(installer);
+    }
+
+    @Test
+    public void testRestoreAppDataSnapshot_noRollbackInProgress() throws Exception {
+        Installer installer = mock(Installer.class);
+        AppDataRollbackHelper helper = spy(new AppDataRollbackHelper(installer));
+
+        RollbackData rd = createInProgressRollbackData("com.foo");
+        // Override the in progress flag.
+        rd.inProgress = false;
+        assertFalse(helper.restoreAppData("com.foo", rd, 0, 0, 0, "seinfo"));
+        verifyZeroInteractions(installer);
+    }
+
+    @Test
+    public void testRestoreAppDataSnapshot_pendingBackupForUser() throws Exception {
+        Installer installer = mock(Installer.class);
+        AppDataRollbackHelper helper = spy(new AppDataRollbackHelper(installer));
+
+        RollbackData rd = createInProgressRollbackData("com.foo");
+        IntArray pendingBackups = rd.packages.get(0).getPendingBackups();
+        pendingBackups.add(10);
+        pendingBackups.add(11);
+
+        assertTrue(helper.restoreAppData("com.foo", rd, 10 /* userId */, 1, 2, "seinfo"));
+
+        // Should only require FLAG_STORAGE_DE here because we have a pending backup that we
+        // didn't manage to execute.
+        InOrder inOrder = Mockito.inOrder(installer);
+        inOrder.verify(installer).restoreAppDataSnapshot(
+                eq("com.foo"), eq(1), eq(2L), eq("seinfo"), eq(10), eq(Installer.FLAG_STORAGE_DE));
+        inOrder.verifyNoMoreInteractions();
+
+        assertEquals(1, pendingBackups.size());
+        assertEquals(11, pendingBackups.get(0));
+    }
+
+    @Test
+    public void testRestoreAppDataSnapshot_availableBackupForLockedUser() throws Exception {
+        Installer installer = mock(Installer.class);
+        AppDataRollbackHelper helper = spy(new AppDataRollbackHelper(installer));
+        doReturn(true).when(helper).isUserCredentialLocked(eq(10));
+
+        RollbackData rd = createInProgressRollbackData("com.foo");
+
+        assertTrue(helper.restoreAppData("com.foo", rd, 10 /* userId */, 1, 2, "seinfo"));
+
+        InOrder inOrder = Mockito.inOrder(installer);
+        inOrder.verify(installer).restoreAppDataSnapshot(
+                eq("com.foo"), eq(1), eq(2L), eq("seinfo"), eq(10), eq(Installer.FLAG_STORAGE_DE));
+        inOrder.verifyNoMoreInteractions();
+
+        ArrayList<RestoreInfo> pendingRestores = rd.packages.get(0).getPendingRestores();
+        assertEquals(1, pendingRestores.size());
+        assertEquals(10, pendingRestores.get(0).userId);
+        assertEquals(1, pendingRestores.get(0).appId);
+        assertEquals("seinfo", pendingRestores.get(0).seInfo);
+    }
+
+    @Test
+    public void testRestoreAppDataSnapshot_availableBackupForUnockedUser() throws Exception {
+        Installer installer = mock(Installer.class);
+        AppDataRollbackHelper helper = spy(new AppDataRollbackHelper(installer));
+        doReturn(false).when(helper).isUserCredentialLocked(eq(10));
+
+        RollbackData rd = createInProgressRollbackData("com.foo");
+        assertFalse(helper.restoreAppData("com.foo", rd, 10 /* userId */, 1, 2, "seinfo"));
+
+        InOrder inOrder = Mockito.inOrder(installer);
+        inOrder.verify(installer).restoreAppDataSnapshot(
+                eq("com.foo"), eq(1), eq(2L), eq("seinfo"), eq(10),
+                eq(Installer.FLAG_STORAGE_DE | Installer.FLAG_STORAGE_CE));
+        inOrder.verifyNoMoreInteractions();
+
+        ArrayList<RestoreInfo> pendingRestores = rd.packages.get(0).getPendingRestores();
+        assertEquals(0, pendingRestores.size());
+    }
+}