Fix an issue where we might load bubbles with an invalid user id

The existing code that handles removing bubbles for removed users
is dependent on the volatile repository having data. The volatile
repository will only have data if the bubble overflow had previously
been loaded.

This means if you removed an account that has bubbles previously
saved in XML, but hadn't received any bubbles before that
(i.e. never loaded bubble overflow, volatile repository
is not populated), next time you receive a bubble, we'd still
try to load that bubble for an account no longer on the device. This
would happen because we'd try to remove the bubble but the volatile
repository wouldn't have any data so we'd never update the XML.

This CL fixes this issue by passing active accounts to the overflow
loading method & when it reads from disc, it'll go through the list
of bubbles loaded and remove any that don't have a valid user. Then
it persists that new list to disk (if anything changed).

Test: atest BubbleDataRepositoryTest
Test: manual - have a user with a work profile
             - get bubbles on the work profile, dismiss them
             - reboot the device
             - delete the work profile
             - get a bubble on the main user
             => observe that there is no crash
Bug: 285409137
Change-Id: Ide5ca24e88e5e81a0a3b24cd48e2905ea1e5b28e

libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java

@@ -1282,7 +1282,9 @@ public class BubbleController implements ConfigurationChangeListener,
             return;
         }
         mOverflowDataLoadNeeded = false;
-        mDataRepository.loadBubbles(mCurrentUserId, (bubbles) -> {
+        List<UserInfo> users = mUserManager.getAliveUsers();
+        List<Integer> userIds = users.stream().map(userInfo -> userInfo.id).toList();
+        mDataRepository.loadBubbles(mCurrentUserId, userIds, (bubbles) -> {
             bubbles.forEach(bubble -> {
                 if (mBubbleData.hasAnyBubbleWithKey(bubble.getKey())) {
                     // if the bubble is already active, there's no need to push it to overflow

libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleDataRepository.kt

@@ -17,7 +17,6 @@ package com.android.wm.shell.bubbles
 
 import android.annotation.SuppressLint
 import android.annotation.UserIdInt
-import android.content.Context
 import android.content.pm.LauncherApps
 import android.content.pm.LauncherApps.ShortcutQuery.FLAG_MATCH_CACHED
 import android.content.pm.LauncherApps.ShortcutQuery.FLAG_MATCH_DYNAMIC
@@ -25,6 +24,8 @@ import android.content.pm.LauncherApps.ShortcutQuery.FLAG_MATCH_PINNED_BY_ANY_LA
 import android.content.pm.UserInfo
 import android.os.UserHandle
 import android.util.Log
+import android.util.SparseArray
+import com.android.internal.annotations.VisibleForTesting
 import com.android.wm.shell.bubbles.Bubbles.BubbleMetadataFlagListener
 import com.android.wm.shell.bubbles.storage.BubbleEntity
 import com.android.wm.shell.bubbles.storage.BubblePersistentRepository
@@ -38,13 +39,12 @@ import kotlinx.coroutines.cancelAndJoin
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.yield
 
-internal class BubbleDataRepository(
-    context: Context,
+class BubbleDataRepository(
     private val launcherApps: LauncherApps,
-    private val mainExecutor: ShellExecutor
+    private val mainExecutor: ShellExecutor,
+    private val persistentRepository: BubblePersistentRepository,
 ) {
     private val volatileRepository = BubbleVolatileRepository(launcherApps)
-    private val persistentRepository = BubblePersistentRepository(context)
 
     private val coroutineScope = CoroutineScope(Dispatchers.IO + SupervisorJob())
     private var job: Job? = null
@@ -99,6 +99,43 @@ internal class BubbleDataRepository(
         if (volatileRepository.sanitizeBubbles(userIds)) persistToDisk()
     }
 
+    /**
+     * Removes all entities that don't have a user in the activeUsers list, if any entities were
+     * removed it persists the new list to disk.
+     */
+    private fun filterForActiveUsersAndPersist(
+            activeUsers: List<Int>,
+            entitiesByUser: SparseArray<List<BubbleEntity>>
+    ): SparseArray<List<BubbleEntity>> {
+        val validEntitiesByUser = SparseArray<List<BubbleEntity>>()
+        var entitiesChanged = false
+        for (i in 0 until entitiesByUser.size()) {
+            val parentUserId = entitiesByUser.keyAt(i)
+            if (activeUsers.contains(parentUserId)) {
+                val validEntities = mutableListOf<BubbleEntity>()
+                // Check if each of the bubbles in the top-level user still has a valid user
+                // as it could belong to a profile and have a different id from the parent.
+                for (entity in entitiesByUser.get(parentUserId)) {
+                    if (activeUsers.contains(entity.userId)) {
+                        validEntities.add(entity)
+                    } else {
+                        entitiesChanged = true
+                    }
+                }
+                if (validEntities.isNotEmpty()) {
+                    validEntitiesByUser.put(parentUserId, validEntities)
+                }
+            } else {
+                entitiesChanged = true
+            }
+        }
+        if (entitiesChanged) {
+            persistToDisk(validEntitiesByUser)
+            return validEntitiesByUser
+        }
+        return entitiesByUser
+    }
+
     private fun transform(bubbles: List<Bubble>): List<BubbleEntity> {
         return bubbles.mapNotNull { b ->
             BubbleEntity(
@@ -130,7 +167,9 @@ internal class BubbleDataRepository(
      * Job C resumes and reaches yield() and is then cancelled
      * Job D resumes and performs another blocking I/O
      */
-    private fun persistToDisk() {
+    private fun persistToDisk(
+            entitiesByUser: SparseArray<List<BubbleEntity>> = volatileRepository.bubbles
+    ) {
         val prev = job
         job = coroutineScope.launch {
             // if there was an ongoing disk I/O operation, they can be cancelled
@@ -138,7 +177,7 @@ internal class BubbleDataRepository(
             // check for cancellation before disk I/O
             yield()
             // save to disk
-            persistentRepository.persistsToDisk(volatileRepository.bubbles)
+            persistentRepository.persistsToDisk(entitiesByUser)
         }
     }
 
@@ -149,7 +188,12 @@ internal class BubbleDataRepository(
      *           bubbles.
      */
     @SuppressLint("WrongConstant")
-    fun loadBubbles(userId: Int, cb: (List<Bubble>) -> Unit) = coroutineScope.launch {
+    @VisibleForTesting
+    fun loadBubbles(
+            userId: Int,
+            currentUsers: List<Int>,
+            cb: (List<Bubble>) -> Unit
+    ) = coroutineScope.launch {
         /**
          * Load BubbleEntity from disk.
          * e.g.
@@ -160,7 +204,12 @@ internal class BubbleDataRepository(
          * ]
          */
         val entitiesByUser = persistentRepository.readFromDisk()
-        val entities = entitiesByUser.get(userId) ?: return@launch
+
+        // Before doing anything, validate that the entities we loaded are valid & have an existing
+        // user.
+        val validEntitiesByUser = filterForActiveUsersAndPersist(currentUsers, entitiesByUser)
+
+        val entities = validEntitiesByUser.get(userId) ?: return@launch
         volatileRepository.addBubbles(userId, entities)
         /**
          * Extract userId/packageName from these entities.

libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java

@@ -37,6 +37,7 @@ import com.android.wm.shell.bubbles.BubbleDataRepository;
 import com.android.wm.shell.bubbles.BubbleLogger;
 import com.android.wm.shell.bubbles.BubblePositioner;
 import com.android.wm.shell.bubbles.properties.ProdBubbleProperties;
+import com.android.wm.shell.bubbles.storage.BubblePersistentRepository;
 import com.android.wm.shell.common.DisplayController;
 import com.android.wm.shell.common.DisplayImeController;
 import com.android.wm.shell.common.DisplayInsetsController;
@@ -184,7 +185,8 @@ public abstract class WMShellModule {
             IWindowManager wmService) {
         return new BubbleController(context, shellInit, shellCommandHandler, shellController, data,
                 null /* synchronizer */, floatingContentCoordinator,
-                new BubbleDataRepository(context, launcherApps, mainExecutor),
+                new BubbleDataRepository(launcherApps, mainExecutor,
+                        new BubblePersistentRepository(context)),
                 statusBarService, windowManager, windowManagerShellWrapper, userManager,
                 launcherApps, logger, taskStackListener, organizer, positioner, displayController,
                 oneHandedOptional, dragAndDropController, mainExecutor, mainHandler, bgExecutor,

libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataRepositoryTest.kt

@@ -0,0 +1,190 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.bubbles
+
+import android.app.ActivityTaskManager
+import android.content.pm.LauncherApps
+import android.content.pm.ShortcutInfo
+import android.util.SparseArray
+import com.android.wm.shell.ShellTestCase
+import com.android.wm.shell.bubbles.storage.BubbleEntity
+import com.android.wm.shell.bubbles.storage.BubblePersistentRepository
+import com.android.wm.shell.common.ShellExecutor
+import com.google.common.truth.Truth.assertThat
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.mockito.Mockito.mock
+import org.mockito.Mockito.spy
+
+class BubbleDataRepositoryTest : ShellTestCase() {
+
+    private val user0BubbleEntities = listOf(
+        BubbleEntity(
+            userId = 0,
+            packageName = "com.example.messenger",
+            shortcutId = "shortcut-1",
+            key = "0k1",
+            desiredHeight = 120,
+            desiredHeightResId = 0,
+            title = null,
+            taskId = 1,
+            locus = null,
+            isDismissable = true
+        ),
+        BubbleEntity(
+            userId = 10,
+            packageName = "com.example.chat",
+            shortcutId = "alice and bob",
+            key = "0k2",
+            desiredHeight = 0,
+            desiredHeightResId = 16537428,
+            title = "title",
+            taskId = 2,
+            locus = null
+        ),
+        BubbleEntity(
+            userId = 0,
+            packageName = "com.example.messenger",
+            shortcutId = "shortcut-2",
+            key = "0k3",
+            desiredHeight = 120,
+            desiredHeightResId = 0,
+            title = null,
+            taskId = ActivityTaskManager.INVALID_TASK_ID,
+            locus = null
+        )
+    )
+
+    private val user1BubbleEntities = listOf(
+        BubbleEntity(
+            userId = 1,
+            packageName = "com.example.messenger",
+            shortcutId = "shortcut-1",
+            key = "1k1",
+            desiredHeight = 120,
+            desiredHeightResId = 0,
+            title = null,
+            taskId = 3,
+            locus = null,
+            isDismissable = true
+        ),
+        BubbleEntity(
+            userId = 12,
+            packageName = "com.example.chat",
+            shortcutId = "alice and bob",
+            key = "1k2",
+            desiredHeight = 0,
+            desiredHeightResId = 16537428,
+            title = "title",
+            taskId = 4,
+            locus = null
+        ),
+        BubbleEntity(
+            userId = 1,
+            packageName = "com.example.messenger",
+            shortcutId = "shortcut-2",
+            key = "1k3",
+            desiredHeight = 120,
+            desiredHeightResId = 0,
+            title = null,
+            taskId = ActivityTaskManager.INVALID_TASK_ID,
+            locus = null
+        ),
+        BubbleEntity(
+            userId = 12,
+            packageName = "com.example.chat",
+            shortcutId = "alice",
+            key = "1k4",
+            desiredHeight = 0,
+            desiredHeightResId = 16537428,
+            title = "title",
+            taskId = 5,
+            locus = null
+        )
+    )
+
+    private val mainExecutor = mock(ShellExecutor::class.java)
+    private val launcherApps = mock(LauncherApps::class.java)
+
+    private val persistedBubbles = SparseArray<List<BubbleEntity>>()
+
+    private lateinit var dataRepository: BubbleDataRepository
+    private lateinit var persistentRepository: BubblePersistentRepository
+
+    @Before
+    fun setup() {
+        persistentRepository = spy(BubblePersistentRepository(mContext))
+        dataRepository = BubbleDataRepository(launcherApps, mainExecutor, persistentRepository)
+
+        // Add the bubbles to the persistent repository
+        persistedBubbles.put(0, user0BubbleEntities)
+        persistedBubbles.put(1, user1BubbleEntities)
+        persistentRepository.persistsToDisk(persistedBubbles)
+    }
+
+    @After
+    fun teardown() {
+        // Clean up any persisted bubbles for the next run
+        persistentRepository.persistsToDisk(SparseArray())
+    }
+
+    @Test
+    fun testLoadBubbles_invalidParent() {
+        val activeUserIds = listOf(10, 1, 12) // Missing user 0 in persistedBubbles
+        dataRepository.loadBubbles(1, activeUserIds) {
+            // Verify that user 0 has been removed from the persisted list
+            val entitiesByUser = persistentRepository.readFromDisk()
+            assertThat(entitiesByUser.get(0)).isNull()
+        }
+    }
+
+    @Test
+    fun testLoadBubbles_invalidChild() {
+        val activeUserIds = listOf(0, 10, 1) // Missing user 1's child user 12
+        dataRepository.loadBubbles(1, activeUserIds) {
+            // Build a list to compare against
+            val user1BubblesWithoutUser12 = mutableListOf<Bubble>()
+            val user1EntitiesWithoutUser12 = mutableListOf<BubbleEntity>()
+            for (entity in user1BubbleEntities) {
+                if (entity.userId != 12) {
+                    user1BubblesWithoutUser12.add(entity.toBubble())
+                    user1EntitiesWithoutUser12.add(entity)
+                }
+            }
+
+            // Verify that user 12 has been removed from the persisted list
+            val entitiesByUser = persistentRepository.readFromDisk()
+            assertThat(entitiesByUser.get(1)).isEqualTo(user1EntitiesWithoutUser12)
+        }
+    }
+
+    private fun BubbleEntity.toBubble(): Bubble {
+        return Bubble(
+            key,
+            mock(ShortcutInfo::class.java),
+            desiredHeight,
+            desiredHeightResId,
+            title,
+            taskId,
+            locus,
+            isDismissable,
+            mainExecutor,
+            mock(Bubbles.BubbleMetadataFlagListener::class.java)
+        )
+    }
+}

packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java

@@ -35,6 +35,7 @@ import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyList;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.atLeastOnce;
@@ -1131,7 +1132,7 @@ public class BubblesTest extends SysuiTestCase {
         assertThat(mBubbleData.getOverflowBubbleWithKey(mBubbleEntry2User11.getKey())).isNotNull();
 
         // Would have loaded bubbles twice because of user switch
-        verify(mDataRepository, times(2)).loadBubbles(anyInt(), any());
+        verify(mDataRepository, times(2)).loadBubbles(anyInt(), anyList(), any());
     }
 
     @Test
@@ -1187,7 +1188,7 @@ public class BubblesTest extends SysuiTestCase {
         mEntryListener.onEntryRemoved(mRow2, REASON_APP_CANCEL);
         assertThat(mBubbleData.getOverflowBubbles()).isEmpty();
 
-        verify(mDataRepository, times(1)).loadBubbles(anyInt(), any());
+        verify(mDataRepository, times(1)).loadBubbles(anyInt(), anyList(), any());
     }
 
     /**