Perform user data back up in post install stage

Currently the backup of user data is done in the enable rollback stage,
during which there is no guarantee that the package being backed up is
not currently running. Moving the backup to the post install stage will
guarantee that the package is not running.

Test: atest RollbackTest
Test: atest StagedRollbackTest
Bug: 124032231
Change-Id: I4b42a0c5ade1645585d1d6f698637df950d05c72

3 files changed, 46 insertions, 20 deletions

diff --git a/core/java/android/content/rollback/IRollbackManager.aidl b/core/java/android/content/rollback/IRollbackManager.aidl
index db9fcb6abdfa..1b84f2987a8f 100644
--- a/core/java/android/content/rollback/IRollbackManager.aidl
+++ b/core/java/android/content/rollback/IRollbackManager.aidl
@@ -30,9 +30,9 @@ interface IRollbackManager {
             String callerPackageName, in IntentSender statusReceiver);
 
     // Exposed for use from the system server only. Callback from the package
-    // manager during the install flow when user data can be restored for a given
+    // manager during the install flow when user data can be backed up and restored for a given
     // package.
-    void restoreUserData(String packageName, in int[] userIds, int appId, long ceDataInode,
+    void snapshotAndRestoreUserData(String packageName, in int[] userIds, int appId, long ceDataInode,
             String seInfo, int token);
 
     // Exposed for test purposes only.
diff --git a/services/core/java/com/android/server/pm/PackageManagerService.java b/services/core/java/com/android/server/pm/PackageManagerService.java
index 84470f9c2838..6ad32495c0ba 100644
--- a/services/core/java/com/android/server/pm/PackageManagerService.java
+++ b/services/core/java/com/android/server/pm/PackageManagerService.java
@@ -14748,7 +14748,7 @@ public class PackageManagerService extends IPackageManager.Stub
 
             if (ps != null) {
                 try {
-                    rm.restoreUserData(packageName, installedUsers, appId, ceDataInode,
+                    rm.snapshotAndRestoreUserData(packageName, installedUsers, appId, ceDataInode,
                             seInfo, token);
                 } catch (RemoteException re) {
                     // Cannot happen, the RollbackManager is hosted in the same process.
diff --git a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
index 9a2778dbf535..301f65061045 100644
--- a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
+++ b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
@@ -923,8 +923,8 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
         }
 
         if (rd != null) {
-            // This is the apk session for a staged session. We have already
-            // backed up the apks, we just need to do user data backup.
+            // This is the apk session for a staged session. We do not need to create a new rollback
+            // for this session.
             PackageParser.PackageLite newPackage = null;
             try {
                 newPackage = PackageParser.parsePackageLite(
@@ -937,8 +937,6 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
             for (PackageRollbackInfo info : rd.info.getPackages()) {
                 if (info.getPackageName().equals(packageName)) {
                     info.getInstalledUsers().addAll(IntArray.wrap(installedUsers));
-                    mAppDataRollbackHelper.snapshotAppData(rd.info.getRollbackId(), info);
-                    saveRollbackData(rd);
                     return true;
                 }
             }
@@ -959,8 +957,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
         }
         newRollback.addToken(token);
 
-        return enableRollbackForPackageSession(newRollback.data, packageSession,
-                installedUsers, /* snapshotUserData*/ true);
+        return enableRollbackForPackageSession(newRollback.data, packageSession, installedUsers);
     }
 
     /**
@@ -971,8 +968,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
      * @return true on success, false on failure.
      */
     private boolean enableRollbackForPackageSession(RollbackData data,
-            PackageInstaller.SessionInfo session, @NonNull int[] installedUsers,
-            boolean snapshotUserData) {
+            PackageInstaller.SessionInfo session, @NonNull int[] installedUsers) {
         // TODO: Don't attempt to enable rollback for split installs.
         final int installFlags = session.installFlags;
         if ((installFlags & PackageManager.INSTALL_ENABLE_ROLLBACK) == 0) {
@@ -1033,10 +1029,6 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
                 isApex, IntArray.wrap(installedUsers),
                 new SparseLongArray() /* ceSnapshotInodes */);
 
-        if (snapshotUserData && !isApex) {
-            mAppDataRollbackHelper.snapshotAppData(data.info.getRollbackId(), packageRollbackInfo);
-        }
-
         try {
             ApplicationInfo appInfo = pkgInfo.applicationInfo;
             RollbackStore.backupPackageCodePath(data, packageName, appInfo.sourceDir);
@@ -1057,13 +1049,15 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
     }
 
     @Override
-    public void restoreUserData(String packageName, int[] userIds, int appId, long ceDataInode,
-            String seInfo, int token) {
+    public void snapshotAndRestoreUserData(String packageName, int[] userIds, int appId,
+            long ceDataInode, String seInfo, int token) {
         if (Binder.getCallingUid() != Process.SYSTEM_UID) {
-            throw new SecurityException("restoreUserData may only be called by the system.");
+            throw new SecurityException(
+                    "snapshotAndRestoreUserData may only be called by the system.");
         }
 
         getHandler().post(() -> {
+            snapshotUserDataInternal(packageName);
             restoreUserDataInternal(packageName, userIds, appId, ceDataInode, seInfo, token);
             final PackageManagerInternal pmi = LocalServices.getService(
                     PackageManagerInternal.class);
@@ -1071,6 +1065,38 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
         });
     }
 
+    private void snapshotUserDataInternal(String packageName) {
+        synchronized (mLock) {
+            // staged installs
+            ensureRollbackDataLoadedLocked();
+            for (int i = 0; i < mRollbacks.size(); i++) {
+                RollbackData data = mRollbacks.get(i);
+                if (data.state != RollbackData.ROLLBACK_STATE_ENABLING) {
+                    continue;
+                }
+
+                for (PackageRollbackInfo info : data.info.getPackages()) {
+                    if (info.getPackageName().equals(packageName)) {
+                        mAppDataRollbackHelper.snapshotAppData(data.info.getRollbackId(), info);
+                        saveRollbackData(data);
+                        return;
+                    }
+                }
+            }
+            // non-staged installs
+            PackageRollbackInfo info;
+            for (NewRollback rollback : mNewRollbacks) {
+                info = getPackageRollbackInfo(rollback.data, packageName);
+                if (info != null) {
+                    mAppDataRollbackHelper.snapshotAppData(rollback.data.info.getRollbackId(),
+                            info);
+                    saveRollbackData(rollback.data);
+                    return;
+                }
+            }
+        }
+    }
+
     private void restoreUserDataInternal(String packageName, int[] userIds, int appId,
             long ceDataInode, String seInfo, int token) {
         PackageRollbackInfo info = null;
@@ -1130,7 +1156,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
 
             if (!session.isMultiPackage()) {
                 if (!enableRollbackForPackageSession(newRollback.data, session,
-                            new int[0], /* snapshotUserData */ false)) {
+                            new int[0])) {
                     Log.e(TAG, "Unable to enable rollback for session: " + sessionId);
                     result.offer(false);
                     return;
@@ -1145,7 +1171,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
                         return;
                     }
                     if (!enableRollbackForPackageSession(newRollback.data, childSession,
-                                new int[0], /* snapshotUserData */ false)) {
+                                new int[0])) {
                         Log.e(TAG, "Unable to enable rollback for session: " + sessionId);
                         result.offer(false);
                         return;