Add rollback parameters when submitting sessions to apexd.

This adds parameters to tell apexd if a session has rollback enabled,
or if it is a rollback. The rollback id is also included.

Previously in StagingManager, rollback was only enabled after the
session was submitted. This ordering is reversed in this change to
allow the rollback id to be retrieved.

For the 'is a rollback' case, a new InstallReason of INSTALL_REASON_ROLLBACK
is added, and the rollback id can be retrieved from RollbackManager.

This is an alternate approach to ag/9826517.

Bug: 141148175
Test: atest CtsRollbackManagerHostTestCases
Test: atest CtsStagedInstallHostTestCases
Change-Id: I560ee4ed6dd82277892a511909378d5a5a8502ff

7 files changed, 105 insertions, 49 deletions

diff --git a/core/java/android/content/pm/PackageManager.java b/core/java/android/content/pm/PackageManager.java
index 6e890ba0d827..94af5416aa8d 100644
--- a/core/java/android/content/pm/PackageManager.java
+++ b/core/java/android/content/pm/PackageManager.java
@@ -946,7 +946,8 @@ public abstract class PackageManager {
             INSTALL_REASON_POLICY,
             INSTALL_REASON_DEVICE_RESTORE,
             INSTALL_REASON_DEVICE_SETUP,
-            INSTALL_REASON_USER
+            INSTALL_REASON_USER,
+            INSTALL_REASON_ROLLBACK
     })
     @Retention(RetentionPolicy.SOURCE)
     public @interface InstallReason {}
@@ -977,6 +978,13 @@ public abstract class PackageManager {
     public static final int INSTALL_REASON_USER = 4;
 
     /**
+     * Code indicating that the package installation was a rollback initiated by RollbackManager.
+     *
+     * @hide
+     */
+    public static final int INSTALL_REASON_ROLLBACK = 5;
+
+    /**
      * @hide
      */
     public static final int INSTALL_UNKNOWN = 0;
diff --git a/core/java/android/content/rollback/IRollbackManager.aidl b/core/java/android/content/rollback/IRollbackManager.aidl
index 8c2a65faf8ed..cda0e9865054 100644
--- a/core/java/android/content/rollback/IRollbackManager.aidl
+++ b/core/java/android/content/rollback/IRollbackManager.aidl
@@ -44,9 +44,10 @@ interface IRollbackManager {
     // Used by the staging manager to notify the RollbackManager that a session is
     // being staged. In the case of multi-package sessions, the specified sessionId
     // is that of the parent session.
+    // Returns the rollback id if rollback was enabled successfully, or -1 if not.
     //
     // NOTE: This call is synchronous.
-    boolean notifyStagedSession(int sessionId);
+    int notifyStagedSession(int sessionId);
 
     // Used by the staging manager to notify the RollbackManager of the apk
     // session for a staged session.
diff --git a/services/core/java/com/android/server/pm/ApexManager.java b/services/core/java/com/android/server/pm/ApexManager.java
index 2b4b409f329a..5ae8c5881cb1 100644
--- a/services/core/java/com/android/server/pm/ApexManager.java
+++ b/services/core/java/com/android/server/pm/ApexManager.java
@@ -17,11 +17,11 @@
 package com.android.server.pm;
 
 import android.annotation.IntDef;
-import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.apex.ApexInfo;
 import android.apex.ApexInfoList;
 import android.apex.ApexSessionInfo;
+import android.apex.ApexSessionParams;
 import android.apex.IApexService;
 import android.content.BroadcastReceiver;
 import android.content.Context;
@@ -176,13 +176,9 @@ abstract class ApexManager {
      * enough for it to be activated at the next boot, the caller needs to call
      * {@link #markStagedSessionReady(int)}.
      *
-     * @param sessionId the identifier of the {@link PackageInstallerSession} being submitted.
-     * @param childSessionIds if {@code sessionId} is a multi-package session, this should contain
-     *                        an array of identifiers of all the child sessions. Otherwise it should
-     *                        be an empty array.
      * @throws PackageManagerException if call to apexd fails
      */
-    abstract ApexInfoList submitStagedSession(int sessionId, @NonNull int[] childSessionIds)
+    abstract ApexInfoList submitStagedSession(ApexSessionParams params)
             throws PackageManagerException;
 
     /**
@@ -450,11 +446,10 @@ abstract class ApexManager {
         }
 
         @Override
-        ApexInfoList submitStagedSession(int sessionId, @NonNull int[] childSessionIds)
-                throws PackageManagerException {
+        ApexInfoList submitStagedSession(ApexSessionParams params) throws PackageManagerException {
             try {
                 final ApexInfoList apexInfoList = new ApexInfoList();
-                mApexService.submitStagedSession(sessionId, childSessionIds, apexInfoList);
+                mApexService.submitStagedSession(params, apexInfoList);
                 return apexInfoList;
             } catch (RemoteException re) {
                 Slog.e(TAG, "Unable to contact apexservice", re);
@@ -686,7 +681,7 @@ abstract class ApexManager {
         }
 
         @Override
-        ApexInfoList submitStagedSession(int sessionId, int[] childSessionIds)
+        ApexInfoList submitStagedSession(ApexSessionParams params)
                 throws PackageManagerException {
             throw new PackageManagerException(PackageManager.INSTALL_FAILED_INTERNAL_ERROR,
                     "Device doesn't support updating APEX");
diff --git a/services/core/java/com/android/server/pm/StagingManager.java b/services/core/java/com/android/server/pm/StagingManager.java
index 7ea4e984b426..688c34fed44f 100644
--- a/services/core/java/com/android/server/pm/StagingManager.java
+++ b/services/core/java/com/android/server/pm/StagingManager.java
@@ -21,6 +21,7 @@ import android.annotation.Nullable;
 import android.apex.ApexInfo;
 import android.apex.ApexInfoList;
 import android.apex.ApexSessionInfo;
+import android.apex.ApexSessionParams;
 import android.content.Context;
 import android.content.IIntentReceiver;
 import android.content.IIntentSender;
@@ -36,6 +37,8 @@ import android.content.pm.PackageParser.SigningDetails;
 import android.content.pm.PackageParser.SigningDetails.SignatureSchemeVersion;
 import android.content.pm.ParceledListSlice;
 import android.content.rollback.IRollbackManager;
+import android.content.rollback.RollbackInfo;
+import android.content.rollback.RollbackManager;
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.IBinder;
@@ -49,6 +52,7 @@ import android.os.storage.StorageManager;
 import android.util.IntArray;
 import android.util.Slog;
 import android.util.SparseArray;
+import android.util.SparseIntArray;
 import android.util.apk.ApkSignatureVerifier;
 
 import com.android.internal.annotations.GuardedBy;
@@ -82,6 +86,9 @@ public class StagingManager {
     @GuardedBy("mStagedSessions")
     private final SparseArray<PackageInstallerSession> mStagedSessions = new SparseArray<>();
 
+    @GuardedBy("mStagedSessions")
+    private final SparseIntArray mSessionRollbackIds = new SparseIntArray();
+
     StagingManager(PackageInstallerService pi, ApexManager am, Context context) {
         mPi = pi;
         mApexManager = am;
@@ -166,18 +173,32 @@ public class StagingManager {
 
     private List<PackageInfo> submitSessionToApexService(
             @NonNull PackageInstallerSession session) throws PackageManagerException {
-        final IntArray childSessionsIds = new IntArray();
+        final IntArray childSessionIds = new IntArray();
         if (session.isMultiPackage()) {
             for (int id : session.getChildSessionIds()) {
                 if (isApexSession(mStagedSessions.get(id))) {
-                    childSessionsIds.add(id);
+                    childSessionIds.add(id);
+                }
+            }
+        }
+        ApexSessionParams apexSessionParams = new ApexSessionParams();
+        apexSessionParams.sessionId = session.sessionId;
+        apexSessionParams.childSessionIds = childSessionIds.toArray();
+        if (session.params.installReason == PackageManager.INSTALL_REASON_ROLLBACK) {
+            apexSessionParams.isRollback = true;
+            apexSessionParams.rollbackId = retrieveRollbackIdForCommitSession(session.sessionId);
+        } else {
+            synchronized (mStagedSessions) {
+                int rollbackId = mSessionRollbackIds.get(session.sessionId, -1);
+                if (rollbackId != -1) {
+                    apexSessionParams.hasRollbackEnabled = true;
+                    apexSessionParams.rollbackId = rollbackId;
                 }
             }
         }
         // submitStagedSession will throw a PackageManagerException if apexd verification fails,
         // which will be propagated to populate stagedSessionErrorMessage of this session.
-        final ApexInfoList apexInfoList = mApexManager.submitStagedSession(session.sessionId,
-                childSessionsIds.toArray());
+        final ApexInfoList apexInfoList = mApexManager.submitStagedSession(apexSessionParams);
         final List<PackageInfo> result = new ArrayList<>();
         for (ApexInfo apexInfo : apexInfoList.apexInfos) {
             final PackageInfo packageInfo;
@@ -208,6 +229,19 @@ public class StagingManager {
         return result;
     }
 
+    private int retrieveRollbackIdForCommitSession(int sessionId) throws PackageManagerException {
+        RollbackManager rm = mContext.getSystemService(RollbackManager.class);
+
+        List<RollbackInfo> rollbacks = rm.getRecentlyCommittedRollbacks();
+        for (RollbackInfo rollback : rollbacks) {
+            if (rollback.getCommittedSessionId() == sessionId) {
+                return rollback.getRollbackId();
+            }
+        }
+        throw new PackageManagerException(
+                "Could not find rollback id for commit session: " + sessionId);
+    }
+
     private void checkRequiredVersionCode(final PackageInstallerSession session,
             final PackageInfo activePackage) throws PackageManagerException {
         if (session.params.requiredInstalledVersionCode == PackageManager.VERSION_CODE_HIGHEST) {
@@ -633,6 +667,7 @@ public class StagingManager {
     void abortSession(@NonNull PackageInstallerSession session) {
         synchronized (mStagedSessions) {
             mStagedSessions.remove(session.sessionId);
+            mSessionRollbackIds.delete(session.sessionId);
         }
     }
 
@@ -865,6 +900,28 @@ public class StagingManager {
          */
         private void handlePreRebootVerification_Start(@NonNull PackageInstallerSession session) {
             Slog.d(TAG, "Starting preRebootVerification for session " + session.sessionId);
+
+            if ((session.params.installFlags & PackageManager.INSTALL_ENABLE_ROLLBACK) != 0) {
+                // If rollback is enabled for this session, we call through to the RollbackManager
+                // with the list of sessions it must enable rollback for. Note that
+                // notifyStagedSession is a synchronous operation.
+                final IRollbackManager rm = IRollbackManager.Stub.asInterface(
+                        ServiceManager.getService(Context.ROLLBACK_SERVICE));
+                try {
+                    // NOTE: To stay consistent with the non-staged install flow, we don't fail the
+                    // entire install if rollbacks can't be enabled.
+                    int rollbackId = rm.notifyStagedSession(session.sessionId);
+                    if (rollbackId != -1) {
+                        synchronized (mStagedSessions) {
+                            mSessionRollbackIds.put(session.sessionId, rollbackId);
+                        }
+                    }
+                } catch (RemoteException re) {
+                    Slog.e(TAG, "Failed to notifyStagedSession for session: "
+                            + session.sessionId, re);
+                }
+            }
+
             notifyPreRebootVerification_Start_Complete(session.sessionId);
         }
 
@@ -929,25 +986,6 @@ public class StagingManager {
          * </ul></p>
          */
         private void handlePreRebootVerification_End(@NonNull PackageInstallerSession session) {
-            if ((session.params.installFlags & PackageManager.INSTALL_ENABLE_ROLLBACK) != 0) {
-                // If rollback is enabled for this session, we call through to the RollbackManager
-                // with the list of sessions it must enable rollback for. Note that
-                // notifyStagedSession is a synchronous operation.
-                final IRollbackManager rm = IRollbackManager.Stub.asInterface(
-                        ServiceManager.getService(Context.ROLLBACK_SERVICE));
-                try {
-                    // NOTE: To stay consistent with the non-staged install flow, we don't fail the
-                    // entire install if rollbacks can't be enabled.
-                    if (!rm.notifyStagedSession(session.sessionId)) {
-                        Slog.e(TAG, "Unable to enable rollback for session: "
-                                + session.sessionId);
-                    }
-                } catch (RemoteException re) {
-                    Slog.e(TAG, "Failed to notifyStagedSession for session: "
-                            + session.sessionId, re);
-                }
-            }
-
             // Proactively mark session as ready before calling apexd. Although this call order
             // looks counter-intuitive, this is the easiest way to ensure that session won't end up
             // in the inconsistent state:
diff --git a/services/core/java/com/android/server/rollback/Rollback.java b/services/core/java/com/android/server/rollback/Rollback.java
index 8b79c3ff66c6..6898e1cd1f83 100644
--- a/services/core/java/com/android/server/rollback/Rollback.java
+++ b/services/core/java/com/android/server/rollback/Rollback.java
@@ -418,6 +418,7 @@ class Rollback {
                 if (isStaged()) {
                     parentParams.setStaged();
                 }
+                parentParams.setInstallReason(PackageManager.INSTALL_REASON_ROLLBACK);
 
                 int parentSessionId = packageInstaller.createSession(parentParams);
                 PackageInstaller.Session parentSession = packageInstaller.openSession(
@@ -484,6 +485,7 @@ class Rollback {
                                 synchronized (mLock) {
                                     mState = ROLLBACK_STATE_AVAILABLE;
                                     mRestoreUserDataInProgress = false;
+                                    info.setCommittedSessionId(-1);
                                 }
                                 sendFailure(context, statusReceiver,
                                         RollbackManager.STATUS_FAILURE_INSTALL,
@@ -500,7 +502,6 @@ class Rollback {
                                     mRestoreUserDataInProgress = false;
                                 }
 
-                                info.setCommittedSessionId(parentSessionId);
                                 info.getCausePackages().addAll(causePackages);
                                 RollbackStore.deletePackageCodePaths(this);
                                 RollbackStore.saveRollback(this);
@@ -528,6 +529,7 @@ class Rollback {
                 );
 
                 mState = ROLLBACK_STATE_COMMITTED;
+                info.setCommittedSessionId(parentSessionId);
                 mRestoreUserDataInProgress = true;
                 parentSession.commit(receiver.getIntentSender());
             } catch (IOException e) {
diff --git a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
index 3f5e2a447d28..9324870904d9 100644
--- a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
+++ b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
@@ -897,11 +897,11 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
     }
 
     @Override
-    public boolean notifyStagedSession(int sessionId) {
+    public int notifyStagedSession(int sessionId) {
         if (Binder.getCallingUid() != Process.SYSTEM_UID) {
             throw new SecurityException("notifyStagedSession may only be called by the system.");
         }
-        final LinkedBlockingQueue<Boolean> result = new LinkedBlockingQueue<>();
+        final LinkedBlockingQueue<Integer> result = new LinkedBlockingQueue<>();
 
         // NOTE: We post this runnable on the RollbackManager's binder thread because we'd prefer
         // to preserve the invariant that all operations that modify state happen there.
@@ -911,7 +911,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
             final PackageInstaller.SessionInfo session = installer.getSessionInfo(sessionId);
             if (session == null) {
                 Slog.e(TAG, "No matching install session for: " + sessionId);
-                result.offer(false);
+                result.offer(-1);
                 return;
             }
 
@@ -923,7 +923,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
             if (!session.isMultiPackage()) {
                 if (!enableRollbackForPackageSession(newRollback.rollback, session)) {
                     Slog.e(TAG, "Unable to enable rollback for session: " + sessionId);
-                    result.offer(false);
+                    result.offer(-1);
                     return;
                 }
             } else {
@@ -932,25 +932,30 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
                             installer.getSessionInfo(childSessionId);
                     if (childSession == null) {
                         Slog.e(TAG, "No matching child install session for: " + childSessionId);
-                        result.offer(false);
+                        result.offer(-1);
                         return;
                     }
                     if (!enableRollbackForPackageSession(newRollback.rollback, childSession)) {
                         Slog.e(TAG, "Unable to enable rollback for session: " + sessionId);
-                        result.offer(false);
+                        result.offer(-1);
                         return;
                     }
                 }
             }
 
-            result.offer(completeEnableRollback(newRollback, true) != null);
+            Rollback rollback = completeEnableRollback(newRollback, true);
+            if (rollback == null) {
+                result.offer(-1);
+            } else {
+                result.offer(rollback.info.getRollbackId());
+            }
         });
 
         try {
             return result.take();
         } catch (InterruptedException ie) {
             Slog.e(TAG, "Interrupted while waiting for notifyStagedSession response");
-            return false;
+            return -1;
         }
     }
 
diff --git a/services/tests/servicestests/src/com/android/server/pm/ApexManagerTest.java b/services/tests/servicestests/src/com/android/server/pm/ApexManagerTest.java
index fb9c68a5b70d..40ada2aedd59 100644
--- a/services/tests/servicestests/src/com/android/server/pm/ApexManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/pm/ApexManagerTest.java
@@ -30,6 +30,7 @@ import static org.testng.Assert.assertThrows;
 
 import android.apex.ApexInfo;
 import android.apex.ApexSessionInfo;
+import android.apex.ApexSessionParams;
 import android.apex.IApexService;
 import android.content.Context;
 import android.content.pm.PackageInfo;
@@ -183,19 +184,18 @@ public class ApexManagerTest {
     public void testSubmitStagedSession_throwPackageManagerException() throws RemoteException {
         doAnswer(invocation -> {
             throw new Exception();
-        }).when(mApexService).submitStagedSession(anyInt(), any(), any());
+        }).when(mApexService).submitStagedSession(any(), any());
 
         assertThrows(PackageManagerException.class,
-                () -> mApexManager.submitStagedSession(TEST_SESSION_ID, TEST_CHILD_SESSION_ID));
+                () -> mApexManager.submitStagedSession(testParamsWithChildren()));
     }
 
     @Test
     public void testSubmitStagedSession_throwRunTimeException() throws RemoteException {
-        doThrow(RemoteException.class).when(mApexService).submitStagedSession(anyInt(), any(),
-                any());
+        doThrow(RemoteException.class).when(mApexService).submitStagedSession(any(), any());
 
         assertThrows(RuntimeException.class,
-                () -> mApexManager.submitStagedSession(TEST_SESSION_ID, TEST_CHILD_SESSION_ID));
+                () -> mApexManager.submitStagedSession(testParamsWithChildren()));
     }
 
     @Test
@@ -272,6 +272,13 @@ public class ApexManagerTest {
         return stagedSessionInfo;
     }
 
+    private static ApexSessionParams testParamsWithChildren() {
+        ApexSessionParams params = new ApexSessionParams();
+        params.sessionId = TEST_SESSION_ID;
+        params.childSessionIds = TEST_CHILD_SESSION_ID;
+        return params;
+    }
+
     /**
      * Copies a specified {@code resourceId} to a temp file. Returns a non-null file if the copy
      * succeeded