Fix alarm processing of QUERY_PACKAGE_RESTART

- Code wasn't checking for the correct user which meant that it could
  return true incorrectly
- Code wasn't checking in the ancillary pending alarm lists, which means
  it could return false incorrectly.

Test: atest FrameworksMockingServicesTests:AlarmManagerServiceTest

Bug: 164215561
Change-Id: Ie595b9e208a1a1aa0fdcf52af737b239d9ed69ce
diff --git a/apex/jobscheduler/service/java/com/android/server/alarm/AlarmManagerService.java b/apex/jobscheduler/service/java/com/android/server/alarm/AlarmManagerService.java
index c76a43f..af99536 100644
--- a/apex/jobscheduler/service/java/com/android/server/alarm/AlarmManagerService.java
+++ b/apex/jobscheduler/service/java/com/android/server/alarm/AlarmManagerService.java
@@ -4285,10 +4285,25 @@
         }
     }
 
-    boolean lookForPackageLocked(String packageName) {
-        final ArrayList<Alarm> allAlarms = mAlarmStore.asList();
-        for (final Alarm alarm : allAlarms) {
-            if (alarm.matches(packageName)) {
+    @GuardedBy("mLock")
+    boolean lookForPackageLocked(String packageName, int uid) {
+        // This is called extremely rarely, e.g. when the user opens the force-stop page in settings
+        // so the loops using an iterator should be fine.
+        for (final Alarm alarm : mAlarmStore.asList()) {
+            if (alarm.matches(packageName) && alarm.creatorUid == uid) {
+                return true;
+            }
+        }
+        final ArrayList<Alarm> alarmsForUid = mPendingBackgroundAlarms.get(uid);
+        if (alarmsForUid != null) {
+            for (final Alarm alarm : alarmsForUid) {
+                if (alarm.matches(packageName)) {
+                    return true;
+                }
+            }
+        }
+        for (final Alarm alarm : mPendingNonWakeupAlarms) {
+            if (alarm.matches(packageName) && alarm.creatorUid == uid) {
                 return true;
             }
         }
@@ -5269,7 +5284,7 @@
                     case Intent.ACTION_QUERY_PACKAGE_RESTART:
                         pkgList = intent.getStringArrayExtra(Intent.EXTRA_PACKAGES);
                         for (String packageName : pkgList) {
-                            if (lookForPackageLocked(packageName)) {
+                            if (lookForPackageLocked(packageName, uid)) {
                                 setResultCode(Activity.RESULT_OK);
                                 return;
                             }
diff --git a/services/tests/mockingservicestests/src/com/android/server/alarm/AlarmManagerServiceTest.java b/services/tests/mockingservicestests/src/com/android/server/alarm/AlarmManagerServiceTest.java
index f9f5325..9788c83 100644
--- a/services/tests/mockingservicestests/src/com/android/server/alarm/AlarmManagerServiceTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/alarm/AlarmManagerServiceTest.java
@@ -3856,4 +3856,52 @@
         assertAndHandleMessageSync(REMOVE_EXACT_LISTENER_ALARMS_ON_CACHED);
         assertEquals(2, mService.mAlarmsPerUid.get(TEST_CALLING_UID_2));
     }
+
+    @Test
+    public void lookForPackageLocked() throws Exception {
+        final String package2 = "test.package.2";
+        final int uid2 = 359712;
+        setTestAlarm(ELAPSED_REALTIME, mNowElapsedTest + 10, getNewMockPendingIntent());
+        setTestAlarm(ELAPSED_REALTIME_WAKEUP, mNowElapsedTest + 15,
+                getNewMockPendingIntent(uid2, package2));
+
+        doReturn(true).when(mService).checkAllowNonWakeupDelayLocked(anyLong());
+
+        assertTrue(mService.lookForPackageLocked(TEST_CALLING_PACKAGE, TEST_CALLING_UID));
+        assertTrue(mService.lookForPackageLocked(package2, uid2));
+
+        mNowElapsedTest += 10;  // Advance time past the first alarm only.
+        mTestTimer.expire();
+
+        assertTrue(mService.lookForPackageLocked(TEST_CALLING_PACKAGE, TEST_CALLING_UID));
+        assertTrue(mService.lookForPackageLocked(package2, uid2));
+
+        // The non-wakeup alarm is sent on interactive state change: false -> true.
+        mService.interactiveStateChangedLocked(false);
+        mService.interactiveStateChangedLocked(true);
+
+        assertFalse(mService.lookForPackageLocked(TEST_CALLING_PACKAGE, TEST_CALLING_UID));
+        assertTrue(mService.lookForPackageLocked(package2, uid2));
+
+        mNowElapsedTest += 10; // Advance time past the second alarm.
+        mTestTimer.expire();
+
+        assertFalse(mService.lookForPackageLocked(TEST_CALLING_PACKAGE, TEST_CALLING_UID));
+        assertFalse(mService.lookForPackageLocked(package2, uid2));
+    }
+
+    @Test
+    public void onQueryPackageRestart() {
+        final String[] packages = {"p1", "p2", "p3"};
+        final int uid = 5421;
+        final Intent packageAdded = new Intent(Intent.ACTION_QUERY_PACKAGE_RESTART)
+                .setData(Uri.fromParts("package", packages[0], null))
+                .putExtra(Intent.EXTRA_PACKAGES, packages)
+                .putExtra(Intent.EXTRA_UID, uid);
+        mPackageChangesReceiver.onReceive(mMockContext, packageAdded);
+
+        for (String p : packages) {
+            verify(mService).lookForPackageLocked(p, uid);
+        }
+    }
 }