Fix RollbackPackageHealthObserver for rollback-all

Making RollbackPackageHealthObserver should start monitoring all
packages if any rollback is available.
Handling the cases when a package without rollbacks fails and
mitigationCount equals 1.
Setting 'sys.attempting_reboot' property when we reboot after
committing the rollback, which would prevent factory reset to
execute synchronously.

Test: atest RollbackPackageHealthObserverTest
Bug: b/264997660
Change-Id: I12780acf8dcedb2b139d19339bc5acc632129ea0
diff --git a/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java b/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java
index e437be8..2007079 100644
--- a/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java
+++ b/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java
@@ -17,10 +17,12 @@
 package com.android.server.rollback;
 
 import android.annotation.AnyThread;
+import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.WorkerThread;
 import android.content.Context;
 import android.content.Intent;
+import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
 import android.content.pm.VersionedPackage;
 import android.content.rollback.PackageRollbackInfo;
@@ -68,6 +70,9 @@
 final class RollbackPackageHealthObserver implements PackageHealthObserver {
     private static final String TAG = "RollbackPackageHealthObserver";
     private static final String NAME = "rollback-observer";
+    private static final String PROP_ATTEMPTING_REBOOT = "sys.attempting_reboot";
+    private static final int PERSISTENT_MASK = ApplicationInfo.FLAG_PERSISTENT
+            | ApplicationInfo.FLAG_SYSTEM;
 
     private final Context mContext;
     private final Handler mHandler;
@@ -114,10 +119,10 @@
             // For native crashes, we will directly roll back any available rollbacks
             // Note: For non-native crashes the rollback-all step has higher impact
             impact = PackageHealthObserverImpact.USER_IMPACT_LEVEL_30;
-        } else if (mitigationCount == 1 && getAvailableRollback(failedPackage) != null) {
+        } else if (getAvailableRollback(failedPackage) != null) {
             // Rollback is available, we may get a callback into #execute
             impact = PackageHealthObserverImpact.USER_IMPACT_LEVEL_30;
-        } else if (mitigationCount > 1 && anyRollbackAvailable) {
+        } else if (anyRollbackAvailable) {
             // If any rollbacks are available, we will commit them
             impact = PackageHealthObserverImpact.USER_IMPACT_LEVEL_70;
         }
@@ -133,14 +138,10 @@
             return true;
         }
 
-        if (mitigationCount == 1) {
-            RollbackInfo rollback = getAvailableRollback(failedPackage);
-            if (rollback == null) {
-                Slog.w(TAG, "Expected rollback but no valid rollback found for " + failedPackage);
-                return false;
-            }
+        RollbackInfo rollback = getAvailableRollback(failedPackage);
+        if (rollback != null) {
             mHandler.post(() -> rollbackPackage(rollback, failedPackage, rollbackReason));
-        } else if (mitigationCount > 1) {
+        } else {
             mHandler.post(() -> rollbackAll(rollbackReason));
         }
 
@@ -153,6 +154,30 @@
         return NAME;
     }
 
+    @Override
+    public boolean isPersistent() {
+        return true;
+    }
+
+    @Override
+    public boolean mayObservePackage(String packageName) {
+        if (mContext.getSystemService(RollbackManager.class)
+                .getAvailableRollbacks().isEmpty()) {
+            return false;
+        }
+        return isPersistentSystemApp(packageName);
+    }
+
+    private boolean isPersistentSystemApp(@NonNull String packageName) {
+        PackageManager pm = mContext.getPackageManager();
+        try {
+            ApplicationInfo info = pm.getApplicationInfo(packageName, 0);
+            return (info.flags & PERSISTENT_MASK) == PERSISTENT_MASK;
+        } catch (PackageManager.NameNotFoundException e) {
+            return false;
+        }
+    }
+
     private void assertInWorkerThread() {
         Preconditions.checkState(mHandler.getLooper().isCurrentThread());
     }
@@ -425,6 +450,7 @@
                 markStagedSessionHandled(rollback.getRollbackId());
                 // Wait for all pending staged sessions to get handled before rebooting.
                 if (isPendingStagedSessionsEmpty()) {
+                    SystemProperties.set(PROP_ATTEMPTING_REBOOT, "true");
                     mContext.getSystemService(PowerManager.class).reboot("Rollback staged install");
                 }
             }
diff --git a/services/tests/mockingservicestests/src/com/android/server/rollback/RollbackPackageHealthObserverTest.java b/services/tests/mockingservicestests/src/com/android/server/rollback/RollbackPackageHealthObserverTest.java
index 541b077..a140730 100644
--- a/services/tests/mockingservicestests/src/com/android/server/rollback/RollbackPackageHealthObserverTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/rollback/RollbackPackageHealthObserverTest.java
@@ -21,10 +21,14 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.when;
 
 import android.content.Context;
+import android.content.pm.ApplicationInfo;
+import android.content.pm.PackageManager;
 import android.content.pm.VersionedPackage;
 import android.content.rollback.PackageRollbackInfo;
 import android.content.rollback.RollbackInfo;
@@ -71,9 +75,12 @@
     RollbackInfo mRollbackInfo;
     @Mock
     PackageRollbackInfo mPackageRollbackInfo;
+    @Mock
+    PackageManager mMockPackageManager;
 
     private MockitoSession mSession;
     private static final String APP_A = "com.package.a";
+    private static final String APP_B = "com.package.b";
     private static final long VERSION_CODE = 1L;
     private static final String LOG_TAG = "RollbackPackageHealthObserverTest";
 
@@ -116,7 +123,7 @@
         RollbackPackageHealthObserver observer =
                 spy(new RollbackPackageHealthObserver(mMockContext));
         VersionedPackage testFailedPackage = new VersionedPackage(APP_A, VERSION_CODE);
-
+        VersionedPackage secondFailedPackage = new VersionedPackage(APP_B, VERSION_CODE);
 
         when(mMockContext.getSystemService(RollbackManager.class)).thenReturn(mRollbackManager);
 
@@ -137,13 +144,16 @@
         assertEquals(PackageWatchdog.PackageHealthObserverImpact.USER_IMPACT_LEVEL_30,
                 observer.onHealthCheckFailed(null,
                         PackageWatchdog.FAILURE_REASON_NATIVE_CRASH, 1));
-        // non-native crash
+        // non-native crash for the package
         assertEquals(PackageWatchdog.PackageHealthObserverImpact.USER_IMPACT_LEVEL_30,
                 observer.onHealthCheckFailed(testFailedPackage,
                         PackageWatchdog.FAILURE_REASON_APP_CRASH, 1));
-        // Second non-native crash again
+        // non-native crash for a different package
         assertEquals(PackageWatchdog.PackageHealthObserverImpact.USER_IMPACT_LEVEL_70,
-                observer.onHealthCheckFailed(testFailedPackage,
+                observer.onHealthCheckFailed(secondFailedPackage,
+                        PackageWatchdog.FAILURE_REASON_APP_CRASH, 1));
+        assertEquals(PackageWatchdog.PackageHealthObserverImpact.USER_IMPACT_LEVEL_70,
+                observer.onHealthCheckFailed(secondFailedPackage,
                         PackageWatchdog.FAILURE_REASON_APP_CRASH, 2));
         // Subsequent crashes when rollbacks have completed
         when(mRollbackManager.getAvailableRollbacks()).thenReturn(List.of());
@@ -152,6 +162,51 @@
                         PackageWatchdog.FAILURE_REASON_APP_CRASH, 3));
     }
 
+    @Test
+    public void testIsPersistent() {
+        RollbackPackageHealthObserver observer =
+                spy(new RollbackPackageHealthObserver(mMockContext));
+        assertTrue(observer.isPersistent());
+    }
+
+    @Test
+    public void testMayObservePackage_withoutAnyRollback() {
+        RollbackPackageHealthObserver observer =
+                spy(new RollbackPackageHealthObserver(mMockContext));
+        when(mMockContext.getSystemService(RollbackManager.class)).thenReturn(mRollbackManager);
+        when(mRollbackManager.getAvailableRollbacks()).thenReturn(List.of());
+        assertFalse(observer.mayObservePackage(APP_A));
+    }
+
+    @Test
+    public void testMayObservePackage_forPersistentApp()
+            throws PackageManager.NameNotFoundException {
+        RollbackPackageHealthObserver observer =
+                spy(new RollbackPackageHealthObserver(mMockContext));
+        ApplicationInfo info = new ApplicationInfo();
+        info.flags = ApplicationInfo.FLAG_PERSISTENT | ApplicationInfo.FLAG_SYSTEM;
+        when(mMockContext.getSystemService(RollbackManager.class)).thenReturn(mRollbackManager);
+        when(mRollbackManager.getAvailableRollbacks()).thenReturn(List.of(mRollbackInfo));
+        when(mRollbackInfo.getPackages()).thenReturn(List.of(mPackageRollbackInfo));
+        when(mMockContext.getPackageManager()).thenReturn(mMockPackageManager);
+        when(mMockPackageManager.getApplicationInfo(APP_A, 0)).thenReturn(info);
+        assertTrue(observer.mayObservePackage(APP_A));
+    }
+
+    @Test
+    public void testMayObservePackage_forNonPersistentApp()
+            throws PackageManager.NameNotFoundException {
+        RollbackPackageHealthObserver observer =
+                spy(new RollbackPackageHealthObserver(mMockContext));
+        when(mMockContext.getSystemService(RollbackManager.class)).thenReturn(mRollbackManager);
+        when(mRollbackManager.getAvailableRollbacks()).thenReturn(List.of(mRollbackInfo));
+        when(mRollbackInfo.getPackages()).thenReturn(List.of(mPackageRollbackInfo));
+        when(mMockContext.getPackageManager()).thenReturn(mMockPackageManager);
+        when(mMockPackageManager.getApplicationInfo(APP_A, 0))
+                .thenThrow(new PackageManager.NameNotFoundException());
+        assertFalse(observer.mayObservePackage(APP_A));
+    }
+
     /**
      * Test that isAutomaticRollbackDenied works correctly when packages that are not
      * denied are sent.