From 5d959629c7553733ca45a12867db9ab4eb7a10fc Mon Sep 17 00:00:00 2001
From: JW Wang <wangchun@google.com>
Date: Thu, 25 Feb 2021 11:55:39 +0800
Subject: [PATCH 1/3] Do APK verification before committing sessions to
 StagingManager (1/n)

This allows us to share the verification flow between both staged and
non-staged sessions.

* Do APK verification before committing sessions to StagingManager.
* remove handlePreRebootVerification_Apk() and its friends which
  are dead code now. This fixes TODOs and reduces coupling between
  StagingManager and PackageInstallerSession.

Bug: 181092466
Test: atest StagedInstallTest StagedInstallInternalTest

Change-Id: I4f9f16471b8e321e1ded5ba6dc9d1cfb982e6bba
---
 .../server/pm/PackageInstallerSession.java    | 49 +++++++-----------
 .../com/android/server/pm/StagingManager.java | 51 ++-----------------
 2 files changed, 24 insertions(+), 76 deletions(-)

diff --git a/services/core/java/com/android/server/pm/PackageInstallerSession.java b/services/core/java/com/android/server/pm/PackageInstallerSession.java
index 903652ab76a59..fcebffaaad35e 100644
--- a/services/core/java/com/android/server/pm/PackageInstallerSession.java
+++ b/services/core/java/com/android/server/pm/PackageInstallerSession.java
@@ -699,28 +699,22 @@ public class PackageInstallerSession extends IPackageInstallerSession.Stub {
         }
 
         /**
-         * Notified by the staging manager that pre-reboot verification is about to start. The
-         * return value should be checked to decide whether it is OK to start pre-reboot
-         * verification. In the case of a destroyed session, {@code false} is returned and there is
-         * no need to start pre-reboot verification.
+         * Called when pre-reboot verification is about to start. This shouldn't be called
+         * on a destroyed session.
          */
-        @Override
-        public boolean notifyStartPreRebootVerification() {
+        private void notifyStartPreRebootVerification() {
             synchronized (mLock) {
+                Preconditions.checkState(!mDestroyed);
                 if (mInPreRebootVerification) {
                     throw new IllegalStateException("Pre-reboot verification has started");
                 }
-                if (mDestroyed) {
-                    return false;
-                }
                 mInPreRebootVerification = true;
-                return true;
             }
         }
 
         /**
-         * Notified by the staging manager that pre-reboot verification has ended. Now it is safe to
-         * clean up the session if {@link #abandon()} has been called previously.
+         * Notified by the staging manager or PIS that pre-reboot verification has ended.
+         * Now it is safe to clean up the session if {@link #abandon()} has been called previously.
          */
         @Override
         public void notifyEndPreRebootVerification() {
@@ -744,6 +738,7 @@ public class PackageInstallerSession extends IPackageInstallerSession.Stub {
             assertCallerIsOwnerOrRootOrSystemLocked();
             Preconditions.checkArgument(isCommitted());
             Preconditions.checkArgument(!mSessionApplied && !mSessionFailed);
+            notifyStartPreRebootVerification();
             verify();
         }
 
@@ -2037,9 +2032,7 @@ public class PackageInstallerSession extends IPackageInstallerSession.Stub {
         if (isStaged()) {
             mStagedSession.setSessionFailed(
                     SessionInfo.STAGED_SESSION_VERIFICATION_FAILED, msgWithErrorCode);
-            // TODO(b/136257624): Remove this once all verification logic has been transferred out
-            //  of StagingManager.
-            mStagingManager.notifyVerificationComplete(mStagedSession);
+            mStagedSession.notifyEndPreRebootVerification();
         } else {
             // Dispatch message to remove session from PackageInstallerService.
             dispatchSessionFinished(error, msg, null);
@@ -2159,24 +2152,19 @@ public class PackageInstallerSession extends IPackageInstallerSession.Stub {
                     .write();
         }
         if (params.isStaged) {
-            mStagingManager.commitSession(mStagedSession);
             // TODO(b/136257624): CTS test fails if we don't send session finished broadcast, even
             //  though ideally, we just need to send session committed broadcast.
             dispatchSessionFinished(INSTALL_SUCCEEDED, "Session staged", null);
-            return;
-        }
 
-        if (isApexSession()) {
-            destroyInternal();
-            dispatchSessionFinished(PackageManager.INSTALL_FAILED_INTERNAL_ERROR,
-                    "APEX packages can only be installed using staged sessions.", null);
-            return;
+            mStagedSession.verifySession();
+        } else {
+            verify();
         }
-        verify();
     }
 
     private void verify() {
         try {
+            //TODO: re-run verify() if reboot happens before staged sessions reach READY
             verifyNonStaged();
         } catch (PackageManagerException e) {
             final String completeMsg = ExceptionUtils.getCompleteMessage(e);
@@ -2451,16 +2439,17 @@ public class PackageInstallerSession extends IPackageInstallerSession.Stub {
     }
 
     private void onVerificationComplete() {
-        // Staged sessions will be installed later during boot
+        // APK verification is done. Continue the installation depending on whether it is a
+        // staged session or not. For a staged session, we will hand it over to the staging
+        // manager to complete the installation.
         if (isStaged()) {
-            // TODO(b/136257624): Remove this once all verification logic has been transferred out
-            //  of StagingManager.
-            mStagingManager.notifyPreRebootVerification_Apk_Complete(mStagedSession);
-            // TODO(b/136257624): We also need to destroy internals for verified staged session,
-            //  otherwise file descriptors are never closed for verified staged session until reboot
+            mStagingManager.commitSession(mStagedSession);
             return;
         }
 
+        // APEX sessions should be handled above
+        Preconditions.checkState(!isApexSession());
+
         install();
     }
 
diff --git a/services/core/java/com/android/server/pm/StagingManager.java b/services/core/java/com/android/server/pm/StagingManager.java
index 0a74032ab214a..2bbb9c5a6ef76 100644
--- a/services/core/java/com/android/server/pm/StagingManager.java
+++ b/services/core/java/com/android/server/pm/StagingManager.java
@@ -136,7 +136,6 @@ public class StagingManager {
         boolean hasParentSessionId();
         long getCommittedMillis();
         void abandon();
-        boolean notifyStartPreRebootVerification();
         void notifyEndPreRebootVerification();
         void verifySession();
     }
@@ -1083,20 +1082,6 @@ public class StagingManager {
         return session;
     }
 
-    // TODO(b/136257624): Temporary API to let PMS communicate with StagingManager. When all
-    //  verification logic is extracted out of StagingManager into PMS, we can remove
-    //  this.
-    void notifyVerificationComplete(StagedSession session) {
-        mPreRebootVerificationHandler.onPreRebootVerificationComplete(session);
-    }
-
-    // TODO(b/136257624): Temporary API to let PMS communicate with StagingManager. When all
-    //  verification logic is extracted out of StagingManager into PMS, we can remove
-    //  this.
-    void notifyPreRebootVerification_Apk_Complete(@NonNull StagedSession session) {
-        mPreRebootVerificationHandler.notifyPreRebootVerification_Apk_Complete(session);
-    }
-
     private final class PreRebootVerificationHandler extends Handler {
         // Hold sessions before handler gets ready to do the verification.
         private List<StagedSession> mPendingSessions;
@@ -1114,7 +1099,6 @@ public class StagingManager {
          * <p><ul>
          *     <li>MSG_PRE_REBOOT_VERIFICATION_START</li>
          *     <li>MSG_PRE_REBOOT_VERIFICATION_APEX</li>
-         *     <li>MSG_PRE_REBOOT_VERIFICATION_APK</li>
          *     <li>MSG_PRE_REBOOT_VERIFICATION_END</li>
          * </ul></p>
          *
@@ -1122,8 +1106,7 @@ public class StagingManager {
          */
         private static final int MSG_PRE_REBOOT_VERIFICATION_START = 1;
         private static final int MSG_PRE_REBOOT_VERIFICATION_APEX = 2;
-        private static final int MSG_PRE_REBOOT_VERIFICATION_APK = 3;
-        private static final int MSG_PRE_REBOOT_VERIFICATION_END = 4;
+        private static final int MSG_PRE_REBOOT_VERIFICATION_END = 3;
 
         @Override
         public void handleMessage(Message msg) {
@@ -1143,9 +1126,6 @@ public class StagingManager {
                     case MSG_PRE_REBOOT_VERIFICATION_APEX:
                         handlePreRebootVerification_Apex(session, rollbackId);
                         break;
-                    case MSG_PRE_REBOOT_VERIFICATION_APK:
-                        handlePreRebootVerification_Apk(session);
-                        break;
                     case MSG_PRE_REBOOT_VERIFICATION_END:
                         handlePreRebootVerification_End(session);
                         break;
@@ -1181,12 +1161,10 @@ public class StagingManager {
                 return;
             }
 
-            if (session.notifyStartPreRebootVerification()) {
-                int sessionId = session.sessionId();
-                Slog.d(TAG, "Starting preRebootVerification for session " + sessionId);
-                obtainMessage(MSG_PRE_REBOOT_VERIFICATION_START, sessionId, -1, session)
-                        .sendToTarget();
-            }
+            int sessionId = session.sessionId();
+            Slog.d(TAG, "Starting preRebootVerification for session " + sessionId);
+            obtainMessage(MSG_PRE_REBOOT_VERIFICATION_START, sessionId, -1, session)
+                    .sendToTarget();
         }
 
         private void onPreRebootVerificationFailure(StagedSession session,
@@ -1215,12 +1193,6 @@ public class StagingManager {
 
         private void notifyPreRebootVerification_Apex_Complete(
                 @NonNull StagedSession session) {
-            obtainMessage(MSG_PRE_REBOOT_VERIFICATION_APK, session.sessionId(), -1, session)
-                    .sendToTarget();
-        }
-
-        private void notifyPreRebootVerification_Apk_Complete(
-                @NonNull StagedSession session) {
             obtainMessage(MSG_PRE_REBOOT_VERIFICATION_END, session.sessionId(), -1, session)
                     .sendToTarget();
         }
@@ -1307,19 +1279,6 @@ public class StagingManager {
             notifyPreRebootVerification_Apex_Complete(session);
         }
 
-        /**
-         * Pre-reboot verification state for apk files. Session is sent to
-         * {@link PackageManagerService} for verification and it notifies back the result via
-         * {@link #notifyPreRebootVerification_Apk_Complete}
-         */
-        private void handlePreRebootVerification_Apk(@NonNull StagedSession session) {
-            if (!session.containsApkSession()) {
-                notifyPreRebootVerification_Apk_Complete(session);
-                return;
-            }
-            session.verifySession();
-        }
-
         /**
          * Pre-reboot verification state for wrapping up:
          * <p><ul>

From e2f2f264acc923d3e8fcf72ff22d89c9ee1c524f Mon Sep 17 00:00:00 2001
From: JW Wang <wangchun@google.com>
Date: Thu, 25 Feb 2021 14:45:43 +0800
Subject: [PATCH 2/3] Restart verification for committed sessions that are not
 READY (2/n)

If reboot happens before a staged session reaches READY, we need to
rerun both APK verificatioin and pre-reboot verification.

* Add a member mBootCompleted to ensure we run verification after boot
  completed since we can't send broadcasts during booting.
* Rewrite startPreRebootVerification() using mBootCompleted to simplify
  the code.
* Fix StagingManagerTest since verifySession() is now the entry point
  for starting verification.

Bug: 181092466
Test: atest StagedInstallTest StagingManagerTest
Change-Id: I7a2708dc6bbc1e758d7e3fa87264552f36e1f42d
---
 .../server/pm/PackageInstallerSession.java    |  1 -
 .../com/android/server/pm/StagingManager.java | 41 ++++++-------------
 .../android/server/pm/StagingManagerTest.java | 19 +++------
 3 files changed, 18 insertions(+), 43 deletions(-)

diff --git a/services/core/java/com/android/server/pm/PackageInstallerSession.java b/services/core/java/com/android/server/pm/PackageInstallerSession.java
index fcebffaaad35e..d70dcaee51e32 100644
--- a/services/core/java/com/android/server/pm/PackageInstallerSession.java
+++ b/services/core/java/com/android/server/pm/PackageInstallerSession.java
@@ -2164,7 +2164,6 @@ public class PackageInstallerSession extends IPackageInstallerSession.Stub {
 
     private void verify() {
         try {
-            //TODO: re-run verify() if reboot happens before staged sessions reach READY
             verifyNonStaged();
         } catch (PackageManagerException e) {
             final String completeMsg = ExceptionUtils.getCompleteMessage(e);
diff --git a/services/core/java/com/android/server/pm/StagingManager.java b/services/core/java/com/android/server/pm/StagingManager.java
index 2bbb9c5a6ef76..2ad16b860ad07 100644
--- a/services/core/java/com/android/server/pm/StagingManager.java
+++ b/services/core/java/com/android/server/pm/StagingManager.java
@@ -80,6 +80,7 @@ import java.io.FileWriter;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Predicate;
@@ -112,6 +113,8 @@ public class StagingManager {
     @GuardedBy("mSuccessfulStagedSessionIds")
     private final List<Integer> mSuccessfulStagedSessionIds = new ArrayList<>();
 
+    private final CompletableFuture<Void> mBootCompleted = new CompletableFuture<>();
+
     interface StagedSession {
         boolean isMultiPackage();
         boolean isApexSession();
@@ -861,7 +864,8 @@ public class StagingManager {
             } else if (!session.isSessionReady()) {
                 // The framework got restarted before the pre-reboot verification could complete,
                 // restart the verification.
-                mPreRebootVerificationHandler.startPreRebootVerification(session);
+                Slog.i(TAG, "Restart verification for session=" + session.sessionId());
+                mBootCompleted.thenRun(() -> session.verifySession());
                 StagedSession session2 = sessions.set(j - 1, session);
                 sessions.set(i, session2);
                 j--;
@@ -1041,7 +1045,7 @@ public class StagingManager {
 
     @VisibleForTesting
     void onBootCompletedBroadcastReceived() {
-        mPreRebootVerificationHandler.readyToStart();
+        mBootCompleted.complete(null);
         BackgroundThread.getExecutor().execute(() -> logFailedApexSessionsIfNecessary());
     }
 
@@ -1083,9 +1087,6 @@ public class StagingManager {
     }
 
     private final class PreRebootVerificationHandler extends Handler {
-        // Hold sessions before handler gets ready to do the verification.
-        private List<StagedSession> mPendingSessions;
-        private boolean mIsReady;
 
         PreRebootVerificationHandler(Looper looper) {
             super(looper);
@@ -1138,33 +1139,15 @@ public class StagingManager {
             }
         }
 
-        // Notify the handler that system is ready, and reschedule the pre-reboot verifications.
-        private synchronized void readyToStart() {
-            mIsReady = true;
-            if (mPendingSessions != null) {
-                for (int i = 0; i < mPendingSessions.size(); i++) {
-                    StagedSession session = mPendingSessions.get(i);
-                    startPreRebootVerification(session);
-                }
-                mPendingSessions = null;
-            }
-        }
-
         // Method for starting the pre-reboot verification
         private synchronized void startPreRebootVerification(
                 @NonNull StagedSession session) {
-            if (!mIsReady) {
-                if (mPendingSessions == null) {
-                    mPendingSessions = new ArrayList<>();
-                }
-                mPendingSessions.add(session);
-                return;
-            }
-
-            int sessionId = session.sessionId();
-            Slog.d(TAG, "Starting preRebootVerification for session " + sessionId);
-            obtainMessage(MSG_PRE_REBOOT_VERIFICATION_START, sessionId, -1, session)
-                    .sendToTarget();
+            mBootCompleted.thenRun(() -> {
+                int sessionId = session.sessionId();
+                Slog.d(TAG, "Starting preRebootVerification for session " + sessionId);
+                obtainMessage(MSG_PRE_REBOOT_VERIFICATION_START, sessionId, -1, session)
+                        .sendToTarget();
+            });
         }
 
         private void onPreRebootVerificationFailure(StagedSession session,
diff --git a/services/tests/mockingservicestests/src/com/android/server/pm/StagingManagerTest.java b/services/tests/mockingservicestests/src/com/android/server/pm/StagingManagerTest.java
index 195cc010c9a7e..65632df648e49 100644
--- a/services/tests/mockingservicestests/src/com/android/server/pm/StagingManagerTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/pm/StagingManagerTest.java
@@ -227,8 +227,8 @@ public class StagingManagerTest {
         assertThat(destroyedNonReadySession.isDestroyed()).isTrue();
 
         mStagingManager.onBootCompletedBroadcastReceived();
-        assertThat(nonReadyApkSession.hasPreRebootVerificationStarted()).isTrue();
-        assertThat(nonReadyApexSession.hasPreRebootVerificationStarted()).isTrue();
+        assertThat(nonReadyApkSession.hasVerificationStarted()).isTrue();
+        assertThat(nonReadyApexSession.hasVerificationStarted()).isTrue();
     }
 
     @Test
@@ -519,7 +519,7 @@ public class StagingManagerTest {
         private int mParentSessionId = -1;
         private String mPackageName;
         private boolean mIsAbandonded = false;
-        private boolean mPreRebootVerificationStarted = false;
+        private boolean mVerificationStarted = false;
         private final List<StagingManager.StagedSession> mChildSessions = new ArrayList<>();
 
         private FakeStagedSession(int sessionId) {
@@ -550,8 +550,8 @@ public class StagingManagerTest {
             return mIsAbandonded;
         }
 
-        private boolean hasPreRebootVerificationStarted() {
-            return mPreRebootVerificationStarted;
+        private boolean hasVerificationStarted() {
+            return mVerificationStarted;
         }
 
         private FakeStagedSession addChildSession(FakeStagedSession session) {
@@ -705,13 +705,6 @@ public class StagingManagerTest {
             mIsAbandonded = true;
         }
 
-        @Override
-        public boolean notifyStartPreRebootVerification() {
-            mPreRebootVerificationStarted = true;
-            // TODO(ioffe): change to true when tests for pre-reboot verification are added.
-            return false;
-        }
-
         @Override
         public void notifyEndPreRebootVerification() {
             throw new UnsupportedOperationException();
@@ -719,7 +712,7 @@ public class StagingManagerTest {
 
         @Override
         public void verifySession() {
-            throw new UnsupportedOperationException();
+            mVerificationStarted = true;
         }
     }
 }

From b26b25df3247cb2856a9d304950b6f991764cdc9 Mon Sep 17 00:00:00 2001
From: JW Wang <wangchun@google.com>
Date: Thu, 25 Feb 2021 16:57:17 +0800
Subject: [PATCH 3/3] Move version code checking out of StagingManager (3/n)

* Make verification work for APEX.
* Now we have version code checking for both APK and APEX in the same
  place (verifyReplacingVersionCode).
* This allows us to share more common flows between staged and
  non-staged sessions.

Note ag/13732406 will fix failed CTS for we move version code
checking into package verification.

Bug: 181092466
Test: atest StagedInstallTest StagedInstallInternalTest
Change-Id: Ia059b29139a5e1e0cda4a3ca9061ab3385a576ae
---
 .../server/pm/PackageInstallerSession.java    | 33 +++++++-----
 .../server/pm/PackageManagerService.java      | 51 ++++++++++++++++---
 .../com/android/server/pm/StagingManager.java | 44 ----------------
 3 files changed, 62 insertions(+), 66 deletions(-)

diff --git a/services/core/java/com/android/server/pm/PackageInstallerSession.java b/services/core/java/com/android/server/pm/PackageInstallerSession.java
index d70dcaee51e32..3ff2ef26d5c97 100644
--- a/services/core/java/com/android/server/pm/PackageInstallerSession.java
+++ b/services/core/java/com/android/server/pm/PackageInstallerSession.java
@@ -2324,16 +2324,14 @@ public class PackageInstallerSession extends IPackageInstallerSession.Stub {
                     "Session not sealed");
         }
 
-        // TODO(b/136257624): Some logic in this if block probably belongs in
-        //  makeInstallParams().
-        if (!isMultiPackage() && !isApexSession()) {
+        if (!isMultiPackage()) {
             Objects.requireNonNull(mPackageName);
             Objects.requireNonNull(mSigningDetails);
             Objects.requireNonNull(mResolvedBaseFile);
 
             // Inherit any packages and native libraries from existing install that
             // haven't been overridden.
-            if (params.mode == SessionParams.MODE_INHERIT_EXISTING) {
+            if (!isApexSession() && params.mode == SessionParams.MODE_INHERIT_EXISTING) {
                 try {
                     final List<File> fromFiles = mResolvedInheritedFiles;
                     final File toDir = stageDir;
@@ -2380,18 +2378,23 @@ public class PackageInstallerSession extends IPackageInstallerSession.Stub {
                             "Failed to inherit existing install", e);
                 }
             }
-            // For mode inherit existing, it would link/copy existing files to stage dir in the
-            // above block. Therefore, we need to parse the complete package in stage dir here.
-            // Besides, PackageLite may be null for staged sessions that don't complete pre-reboot
-            // verification.
-            mPackageLite = getOrParsePackageLiteLocked(stageDir, /* flags */ 0);
 
-            // TODO: surface more granular state from dexopt
-            mInternalProgress = 0.5f;
-            computeProgressLocked(true);
+            if (!isApexSession()) {
+                // For mode inherit existing, it would link/copy existing files to stage dir in the
+                // above block. Therefore, we need to parse the complete package in stage dir here.
+                // Besides, PackageLite may be null for staged sessions that don't complete
+                // pre-reboot verification.
+                mPackageLite = getOrParsePackageLiteLocked(stageDir, /* flags */ 0);
 
-            extractNativeLibraries(mPackageLite, stageDir, params.abiOverride,
-                    mayInheritNativeLibs());
+                // TODO: surface more granular state from dexopt
+                mInternalProgress = 0.5f;
+                computeProgressLocked(true);
+
+                extractNativeLibraries(mPackageLite, stageDir, params.abiOverride,
+                        mayInheritNativeLibs());
+            } else {
+                mPackageLite = getOrParsePackageLiteLocked(mResolvedBaseFile, /* flags */ 0);
+            }
         }
 
         final IPackageInstallObserver2 localObserver;
@@ -2681,6 +2684,8 @@ public class PackageInstallerSession extends IPackageInstallerSession.Stub {
             mPackageName = apk.getPackageName();
             mVersionCode = apk.getLongVersionCode();
         }
+
+        mSigningDetails = apk.getSigningDetails();
     }
 
     /**
diff --git a/services/core/java/com/android/server/pm/PackageManagerService.java b/services/core/java/com/android/server/pm/PackageManagerService.java
index bc991634fb078..1fa879b3796e9 100644
--- a/services/core/java/com/android/server/pm/PackageManagerService.java
+++ b/services/core/java/com/android/server/pm/PackageManagerService.java
@@ -17361,13 +17361,6 @@ public class PackageManagerService extends IPackageManager.Stub
         }
 
         public void handleStartCopy() {
-            if ((installFlags & PackageManager.INSTALL_APEX) != 0) {
-                // Apex packages get verified in StagingManager currently.
-                // TODO(b/136257624): Move apex verification logic out of StagingManager
-                mRet = INSTALL_SUCCEEDED;
-                return;
-            }
-
             PackageInfoLite pkgLite = PackageManagerServiceUtils.getMinimalPackageInfo(mContext,
                     mPackageLite, origin.resolvedPath, installFlags, packageAbiOverride);
 
@@ -17379,7 +17372,10 @@ public class PackageManagerService extends IPackageManager.Stub
             // Perform package verification and enable rollback (unless we are simply moving the
             // package).
             if (!origin.existing) {
-                sendApkVerificationRequest(pkgLite);
+                if ((installFlags & PackageManager.INSTALL_APEX) == 0) {
+                    // TODO(b/182426975): treat APEX as APK when APK verification is concerned
+                    sendApkVerificationRequest(pkgLite);
+                }
                 if ((installFlags & PackageManager.INSTALL_ENABLE_ROLLBACK) != 0) {
                     sendEnableRollbackRequest();
                 }
@@ -25835,6 +25831,11 @@ public class PackageManagerService extends IPackageManager.Stub
 
     private int verifyReplacingVersionCode(PackageInfoLite pkgLite,
             long requiredInstalledVersionCode, int installFlags) {
+        if ((installFlags & PackageManager.INSTALL_APEX) != 0) {
+            return verifyReplacingVersionCodeForApex(
+                    pkgLite, requiredInstalledVersionCode, installFlags);
+        }
+
         String packageName = pkgLite.packageName;
         synchronized (mLock) {
             // Package which currently owns the data that the new package will own if installed.
@@ -25881,6 +25882,40 @@ public class PackageManagerService extends IPackageManager.Stub
         return PackageManager.INSTALL_SUCCEEDED;
     }
 
+    private int verifyReplacingVersionCodeForApex(PackageInfoLite pkgLite,
+            long requiredInstalledVersionCode, int installFlags) {
+        String packageName = pkgLite.packageName;
+
+        final PackageInfo activePackage = mApexManager.getPackageInfo(packageName,
+                ApexManager.MATCH_ACTIVE_PACKAGE);
+        if (activePackage == null) {
+            Slog.w(TAG, "Attempting to install new APEX package " + packageName);
+            return PackageManager.INSTALL_FAILED_WRONG_INSTALLED_VERSION;
+        }
+
+        final long activeVersion = activePackage.getLongVersionCode();
+        if (requiredInstalledVersionCode != PackageManager.VERSION_CODE_HIGHEST
+                && activeVersion != requiredInstalledVersionCode) {
+            Slog.w(TAG, "Installed version of APEX package " + packageName
+                    + " does not match required. Active version: " + activeVersion
+                    + " required: " + requiredInstalledVersionCode);
+            return PackageManager.INSTALL_FAILED_WRONG_INSTALLED_VERSION;
+        }
+
+        final boolean isAppDebuggable = (activePackage.applicationInfo.flags
+                & ApplicationInfo.FLAG_DEBUGGABLE) != 0;
+        final long newVersionCode = pkgLite.getLongVersionCode();
+        if (!PackageManagerServiceUtils.isDowngradePermitted(installFlags, isAppDebuggable)
+                && newVersionCode < activeVersion) {
+            Slog.w(TAG, "Downgrade of APEX package " + packageName
+                    + " is not allowed. Active version: " + activeVersion
+                    + " attempted: " + newVersionCode);
+            return PackageManager.INSTALL_FAILED_VERSION_DOWNGRADE;
+        }
+
+        return PackageManager.INSTALL_SUCCEEDED;
+    }
+
     /**
      * Check and throw if the given before/after packages would be considered a
      * downgrade.
diff --git a/services/core/java/com/android/server/pm/StagingManager.java b/services/core/java/com/android/server/pm/StagingManager.java
index 2ad16b860ad07..dd4260cb039d6 100644
--- a/services/core/java/com/android/server/pm/StagingManager.java
+++ b/services/core/java/com/android/server/pm/StagingManager.java
@@ -28,7 +28,6 @@ import android.content.IIntentSender;
 import android.content.Intent;
 import android.content.IntentFilter;
 import android.content.IntentSender;
-import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageInstaller;
 import android.content.pm.PackageInstaller.SessionInfo;
@@ -295,15 +294,6 @@ public class StagingManager {
                 throw new PackageManagerException(SessionInfo.STAGED_SESSION_VERIFICATION_FAILED,
                         "Failed to parse APEX package " + apexInfo.modulePath, e);
             }
-            final PackageInfo activePackage = mApexManager.getPackageInfo(packageInfo.packageName,
-                    ApexManager.MATCH_ACTIVE_PACKAGE);
-            if (activePackage == null) {
-                Slog.w(TAG, "Attempting to install new APEX package " + packageInfo.packageName);
-                throw new PackageManagerException(SessionInfo.STAGED_SESSION_VERIFICATION_FAILED,
-                        "It is forbidden to install new APEX packages.");
-            }
-            checkRequiredVersionCode(session, activePackage);
-            checkDowngrade(session, activePackage, packageInfo);
             result.add(packageInfo);
             apexPackageNames.add(packageInfo.packageName);
         }
@@ -326,40 +316,6 @@ public class StagingManager {
                 "Could not find rollback id for commit session: " + sessionId);
     }
 
-    private void checkRequiredVersionCode(final StagedSession session,
-            final PackageInfo activePackage) throws PackageManagerException {
-        if (session.sessionParams().requiredInstalledVersionCode
-                == PackageManager.VERSION_CODE_HIGHEST) {
-            return;
-        }
-        final long activeVersion = activePackage.applicationInfo.longVersionCode;
-        if (activeVersion != session.sessionParams().requiredInstalledVersionCode) {
-            throw new PackageManagerException(
-                    SessionInfo.STAGED_SESSION_VERIFICATION_FAILED,
-                    "Installed version of APEX package " + activePackage.packageName
-                            + " does not match required. Active version: " + activeVersion
-                            + " required: " + session.sessionParams().requiredInstalledVersionCode);
-        }
-    }
-
-    private void checkDowngrade(final StagedSession session,
-            final PackageInfo activePackage, final PackageInfo newPackage)
-            throws PackageManagerException {
-        final long activeVersion = activePackage.applicationInfo.longVersionCode;
-        final long newVersionCode = newPackage.applicationInfo.longVersionCode;
-        final boolean isAppDebuggable = (activePackage.applicationInfo.flags
-                & ApplicationInfo.FLAG_DEBUGGABLE) != 0;
-        final boolean allowsDowngrade = PackageManagerServiceUtils.isDowngradePermitted(
-                session.sessionParams().installFlags, isAppDebuggable);
-        if (activeVersion > newVersionCode && !allowsDowngrade) {
-            throw new PackageManagerException(
-                    SessionInfo.STAGED_SESSION_VERIFICATION_FAILED,
-                    "Downgrade of APEX package " + newPackage.packageName
-                            + " is not allowed. Active version: " + activeVersion
-                            + " attempted: " + newVersionCode);
-        }
-    }
-
     // Reverts apex sessions and user data (if checkpoint is supported). Also reboots the device.
     private void abortCheckpoint(String failureReason, boolean supportsCheckpoint,
             boolean needsCheckpoint) {