Merge "Revert "Revert "Make RollbackPackageHealthObserver observe apk-in-apex"""
diff --git a/services/core/java/com/android/server/rollback/Rollback.java b/services/core/java/com/android/server/rollback/Rollback.java
index 7d0072a..5c0dd9a 100644
--- a/services/core/java/com/android/server/rollback/Rollback.java
+++ b/services/core/java/com/android/server/rollback/Rollback.java
@@ -360,7 +360,7 @@
      */
     boolean enableForPackageInApex(String packageName, long installedVersion,
             int rollbackDataPolicy) {
-        // TODO(b/142712057): Extract the new version number of apk-in-apex
+        // TODO(b/147666157): Extract the new version number of apk-in-apex
         // The new version for the apk-in-apex is set to 0 for now. If the package is then further
         // updated via non-staged install flow, then RollbackManagerServiceImpl#onPackageReplaced()
         // will be called and this rollback will be deleted. Other ways of package update have not
diff --git a/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java b/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java
index 7a8ddd4..6686de9 100644
--- a/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java
+++ b/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java
@@ -213,6 +213,23 @@
                 if (packageRollback.getVersionRolledBackFrom().equals(failedPackage)) {
                     return rollback;
                 }
+                // TODO(b/147666157): Extract version number of apk-in-apex so that we don't have
+                //  to rely on complicated reasoning as below
+
+                // Due to b/147666157, for apk in apex, we do not know the version we are rolling
+                // back from. But if a package X is embedded in apex A exclusively (not embedded in
+                // any other apex), which is not guaranteed, then it is sufficient to check only
+                // package names here, as the version of failedPackage and the PackageRollbackInfo
+                // can't be different. If failedPackage has a higher version, then it must have
+                // been updated somehow. There are two ways: it was updated by an update of apex A
+                // or updated directly as apk. In both cases, this rollback would have gotten
+                // expired when onPackageReplaced() was called. Since the rollback exists, it has
+                // same version as failedPackage.
+                if (packageRollback.isApkInApex()
+                        && packageRollback.getVersionRolledBackFrom().getPackageName()
+                        .equals(failedPackage.getPackageName())) {
+                    return rollback;
+                }
             }
         }
         return null;
diff --git a/tests/RollbackTest/Android.bp b/tests/RollbackTest/Android.bp
index 091edd4..98e7b4e 100644
--- a/tests/RollbackTest/Android.bp
+++ b/tests/RollbackTest/Android.bp
@@ -19,7 +19,10 @@
     static_libs: ["androidx.test.rules", "cts-rollback-lib", "cts-install-lib"],
     test_suites: ["general-tests"],
     test_config: "RollbackTest.xml",
-    java_resources: [":com.android.apex.apkrollback.test_v2"],
+    java_resources: [
+        ":com.android.apex.apkrollback.test_v2",
+        ":com.android.apex.apkrollback.test_v2Crashing"
+    ],
 }
 
 java_test_host {
@@ -79,4 +82,14 @@
   key: "com.android.apex.apkrollback.test.key",
   apps: ["TestAppAv2"],
   installable: false,
+}
+
+apex {
+  name: "com.android.apex.apkrollback.test_v2Crashing",
+  manifest: "testdata/manifest_v2.json",
+  androidManifest: "testdata/AndroidManifest.xml",
+  file_contexts: ":apex.test-file_contexts",
+  key: "com.android.apex.apkrollback.test.key",
+  apps: ["TestAppACrashingV2"],
+  installable: false,
 }
\ No newline at end of file
diff --git a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
index ce20311..80491cd 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
@@ -500,8 +500,9 @@
             APK_IN_APEX_TESTAPEX_NAME, 1, /*isApex*/true, APK_IN_APEX_TESTAPEX_NAME + "_v1.apex");
     private static final TestApp TEST_APEX_WITH_APK_V2 = new TestApp("TestApexWithApkV2",
             APK_IN_APEX_TESTAPEX_NAME, 2, /*isApex*/true, APK_IN_APEX_TESTAPEX_NAME + "_v2.apex");
-    private static final TestApp TEST_APP_A_V2_UNKNOWN = new TestApp("Av2Unknown", TestApp.A, 0,
-            /*isApex*/false, "TestAppAv2.apk");
+    private static final TestApp TEST_APEX_WITH_APK_V2_CRASHING = new TestApp(
+            "TestApexWithApkV2Crashing", APK_IN_APEX_TESTAPEX_NAME, 2, /*isApex*/true,
+            APK_IN_APEX_TESTAPEX_NAME + "_v2Crashing.apex");
 
     @Test
     public void testRollbackApexWithApk_Phase1() throws Exception {
@@ -523,7 +524,7 @@
         assertThat(available).isStaged();
         assertThat(available).packagesContainsExactly(
                 Rollback.from(TEST_APEX_WITH_APK_V2).to(TEST_APEX_WITH_APK_V1),
-                Rollback.from(TEST_APP_A_V2_UNKNOWN).to(TestApp.A1));
+                Rollback.from(TestApp.A, 0).to(TestApp.A1));
 
         RollbackUtils.rollback(available.getRollbackId(), TEST_APEX_WITH_APK_V2);
         RollbackInfo committed = RollbackUtils.getCommittedRollbackById(available.getRollbackId());
@@ -531,7 +532,7 @@
         assertThat(committed).isStaged();
         assertThat(committed).packagesContainsExactly(
                 Rollback.from(TEST_APEX_WITH_APK_V2).to(TEST_APEX_WITH_APK_V1),
-                Rollback.from(TEST_APP_A_V2_UNKNOWN).to(TestApp.A1));
+                Rollback.from(TestApp.A, 0).to(TestApp.A1));
         assertThat(committed).causePackagesContainsExactly(TEST_APEX_WITH_APK_V2);
         assertThat(committed.getCommittedSessionId()).isNotEqualTo(-1);
 
@@ -548,6 +549,41 @@
         InstallUtils.processUserData(TestApp.A);
     }
 
+    /**
+     * Installs an apex with an apk that can crash.
+     */
+    @Test
+    public void testRollbackApexWithApkCrashing_Phase1() throws Exception {
+        assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(1);
+        int sessionId = Install.single(TEST_APEX_WITH_APK_V2_CRASHING).setStaged()
+                .setEnableRollback().commit();
+        InstallUtils.waitForSessionReady(sessionId);
+    }
+
+    /**
+     * Verifies rollback has been enabled successfully. Then makes TestApp.A crash.
+     */
+    @Test
+    public void testRollbackApexWithApkCrashing_Phase2() throws Exception {
+        assertThat(InstallUtils.getInstalledVersion(APK_IN_APEX_TESTAPEX_NAME)).isEqualTo(2);
+        assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(2);
+
+        RollbackInfo available = RollbackUtils.getAvailableRollback(APK_IN_APEX_TESTAPEX_NAME);
+        assertThat(available).isStaged();
+        assertThat(available).packagesContainsExactly(
+                Rollback.from(TEST_APEX_WITH_APK_V2).to(TEST_APEX_WITH_APK_V1),
+                Rollback.from(TestApp.A, 0).to(TestApp.A1));
+
+        // Crash TestApp.A PackageWatchdog#TRIGGER_FAILURE_COUNT times to trigger rollback
+        RollbackUtils.sendCrashBroadcast(TestApp.A, 5);
+    }
+
+    @Test
+    public void testRollbackApexWithApkCrashing_Phase3() throws Exception {
+        assertThat(InstallUtils.getInstalledVersion(APK_IN_APEX_TESTAPEX_NAME)).isEqualTo(1);
+        assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(1);
+    }
+
     private static void runShellCommand(String cmd) {
         ParcelFileDescriptor pfd = InstrumentationRegistry.getInstrumentation().getUiAutomation()
                 .executeShellCommand(cmd);
diff --git a/tests/RollbackTest/StagedRollbackTest/src/com/android/tests/rollback/host/StagedRollbackTest.java b/tests/RollbackTest/StagedRollbackTest/src/com/android/tests/rollback/host/StagedRollbackTest.java
index 181e29a..672cbb0 100644
--- a/tests/RollbackTest/StagedRollbackTest/src/com/android/tests/rollback/host/StagedRollbackTest.java
+++ b/tests/RollbackTest/StagedRollbackTest/src/com/android/tests/rollback/host/StagedRollbackTest.java
@@ -261,6 +261,34 @@
         runPhase("testRollbackApexWithApk_Phase3");
     }
 
+    /**
+     * Tests that RollbackPackageHealthObserver is observing apk-in-apex.
+     */
+    @Test
+    public void testRollbackApexWithApkCrashing() throws Exception {
+        getDevice().uninstallPackage("com.android.cts.install.lib.testapp.A");
+        CompatibilityBuildHelper buildHelper = new CompatibilityBuildHelper(getBuild());
+        final String fileName = APK_IN_APEX_TESTAPEX_NAME + "_v1.apex";
+        final File apex = buildHelper.getTestFile(fileName);
+        if (!getDevice().isAdbRoot()) {
+            getDevice().enableAdbRoot();
+        }
+        getDevice().remountSystemWritable();
+        assertTrue(getDevice().pushFile(apex, "/system/apex/" + fileName));
+        getDevice().reboot();
+
+        // Install an apex with apk that crashes
+        runPhase("testRollbackApexWithApkCrashing_Phase1");
+        getDevice().reboot();
+        // Verify apex was installed and then crash the apk
+        runPhase("testRollbackApexWithApkCrashing_Phase2");
+        // Wait for crash to trigger rollback
+        assertTrue(getDevice().waitForDeviceNotAvailable(TimeUnit.MINUTES.toMillis(5)));
+        getDevice().waitForDeviceAvailable();
+        // Verify rollback occurred due to crash of apk-in-apex
+        runPhase("testRollbackApexWithApkCrashing_Phase3");
+    }
+
     private void crashProcess(String processName, int numberOfCrashes) throws Exception {
         String pid = "";
         String lastPid = "invalid";