From 185782ddbcd3c6e20010db05359e8ba2b2457950 Mon Sep 17 00:00:00 2001
From: Mady Mellor <madym@google.com>
Date: Tue, 17 May 2022 00:21:44 +0000
Subject: [PATCH] Fix potential NPE & clean up some logic

- I had accidentally swapped the values in removeBubblesForUser
  in one place which could lead to a crash, it's actually simpler
  to just do the removal there since it's the parentUser that was
  removed so we can remove all of those entries.
- In other places I've added a null check before performing the
  removeIf to be safe
- Added a test checking bad input

Bug: 232847471
Test: atest BubbleVolatileRepositoryTest
Change-Id: I10f5b24e591db62edbb63ad40e588ba19ab621c1
---
 .../shell/bubbles/storage/BubbleVolatileRepository.kt | 11 ++++++++---
 .../bubbles/storage/BubbleVolatileRepositoryTest.kt   |  6 ++++++
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/storage/BubbleVolatileRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/storage/BubbleVolatileRepository.kt
index 0d3ba9a78e353..1eee0291cb261 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/storage/BubbleVolatileRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/storage/BubbleVolatileRepository.kt
@@ -126,7 +126,11 @@ class BubbleVolatileRepository(private val launcherApps: LauncherApps) {
         @UserIdInt userId: Int,
         @UserIdInt parentUserId: Int
     ): Boolean {
-        return entitiesByUser.get(parentUserId).removeIf { b: BubbleEntity -> b.userId == userId }
+        if (entitiesByUser.get(parentUserId) != null) {
+            return entitiesByUser.get(parentUserId).removeIf {
+                b: BubbleEntity -> b.userId == userId }
+        }
+        return false
     }
 
     /**
@@ -141,8 +145,9 @@ class BubbleVolatileRepository(private val launcherApps: LauncherApps) {
             // First check if the user is a parent / top-level user
             val parentUserId = entitiesByUser.keyAt(i)
             if (!activeUsers.contains(parentUserId)) {
-                return removeBubblesForUser(parentUserId, -1)
-            } else {
+                entitiesByUser.remove(parentUserId)
+                return true
+            } else if (entitiesByUser.get(parentUserId) != null) {
                 // Then 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.
                 return entitiesByUser.get(parentUserId).removeIf { b: BubbleEntity ->
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/storage/BubbleVolatileRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/storage/BubbleVolatileRepositoryTest.kt
index 77c70557fb707..9f0d89bc31288 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/storage/BubbleVolatileRepositoryTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/storage/BubbleVolatileRepositoryTest.kt
@@ -269,6 +269,12 @@ class BubbleVolatileRepositoryTest : ShellTestCase() {
         assertThat(repository.getEntities(user11.identifier).toList())
                 .isEqualTo(listOf(bubble11, bubble12))
     }
+
+    @Test
+    fun testRemoveBubbleForUser_invalidInputDoesntCrash() {
+        repository.removeBubblesForUser(-1, 0)
+        repository.removeBubblesForUser(-1, -1)
+    }
 }
 
 private const val PKG_MESSENGER = "com.example.messenger"