Merge "Fix crash when a shortcut update for notes happens before user unlocked" into udc-dev

packages/SystemUI/src/com/android/systemui/notetask/NoteTaskController.kt

@@ -250,6 +250,11 @@ constructor(
      * Widget Picker to all users.
      */
     fun setNoteTaskShortcutEnabled(value: Boolean, user: UserHandle) {
+        if (!userManager.isUserUnlocked(user)) {
+            debugLog { "setNoteTaskShortcutEnabled call but user locked: user=$user" }
+            return
+        }
+
         val componentName = ComponentName(context, CreateNoteTaskShortcutActivity::class.java)
 
         val enabledState =
@@ -305,6 +310,10 @@ constructor(
     /** @see OnRoleHoldersChangedListener */
     fun onRoleHoldersChanged(roleName: String, user: UserHandle) {
         if (roleName != ROLE_NOTES) return
+        if (!userManager.isUserUnlocked(user)) {
+            debugLog { "onRoleHoldersChanged call but user locked: role=$roleName, user=$user" }
+            return
+        }
 
         if (user == userTracker.userHandle) {
             updateNoteTaskAsUser(user)

packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInitializer.kt

@@ -18,8 +18,13 @@ package com.android.systemui.notetask
 import android.app.role.RoleManager
 import android.os.UserHandle
 import android.view.KeyEvent
-import androidx.annotation.VisibleForTesting
+import android.view.KeyEvent.KEYCODE_N
+import android.view.KeyEvent.KEYCODE_STYLUS_BUTTON_TAIL
+import com.android.keyguard.KeyguardUpdateMonitor
+import com.android.keyguard.KeyguardUpdateMonitorCallback
 import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.notetask.NoteTaskEntryPoint.KEYBOARD_SHORTCUT
+import com.android.systemui.notetask.NoteTaskEntryPoint.TAIL_BUTTON
 import com.android.systemui.settings.UserTracker
 import com.android.systemui.statusbar.CommandQueue
 import com.android.wm.shell.bubbles.Bubbles
@@ -35,35 +40,82 @@ constructor(
     private val roleManager: RoleManager,
     private val commandQueue: CommandQueue,
     private val optionalBubbles: Optional<Bubbles>,
+    private val userTracker: UserTracker,
+    private val keyguardUpdateMonitor: KeyguardUpdateMonitor,
     @Background private val backgroundExecutor: Executor,
     @NoteTaskEnabledKey private val isEnabled: Boolean,
-    private val userTracker: UserTracker,
 ) {
 
-    @VisibleForTesting
+    /** Initializes note task related features and glue it with other parts of the SystemUI. */
-    val callbacks =
-        object : CommandQueue.Callbacks {
-            override fun handleSystemKey(key: KeyEvent) {
-                if (key.keyCode == KeyEvent.KEYCODE_STYLUS_BUTTON_TAIL) {
-                    controller.showNoteTask(NoteTaskEntryPoint.TAIL_BUTTON)
-                } else if (
-                    key.keyCode == KeyEvent.KEYCODE_N && key.isMetaPressed && key.isCtrlPressed
-                ) {
-                    controller.showNoteTask(NoteTaskEntryPoint.KEYBOARD_SHORTCUT)
-                }
-            }
-        }
-
     fun initialize() {
         // Guard against feature not being enabled or mandatory dependencies aren't available.
         if (!isEnabled || optionalBubbles.isEmpty) return
 
-        controller.setNoteTaskShortcutEnabled(true, userTracker.userHandle)
+        initializeHandleSystemKey()
+        initializeOnRoleHoldersChanged()
+        initializeOnUserUnlocked()
+    }
+
+    /**
+     * Initializes a callback for [CommandQueue] which will redirect [KeyEvent] from a Stylus to
+     * [NoteTaskController], ensure custom actions can be triggered (i.e., keyboard shortcut).
+     */
+    private fun initializeHandleSystemKey() {
+        val callbacks =
+            object : CommandQueue.Callbacks {
+                override fun handleSystemKey(key: KeyEvent) {
+                    key.toNoteTaskEntryPointOrNull()?.let(controller::showNoteTask)
+                }
+            }
         commandQueue.addCallback(callbacks)
+    }
+
+    /**
+     * Initializes the [RoleManager] role holder changed listener to ensure [NoteTaskController]
+     * will always update whenever the role holder app changes. Keep in mind that a role may change
+     * by direct user interaction (i.e., user goes to settings and change it) or by indirect
+     * interaction (i.e., the current role holder app is uninstalled).
+     */
+    private fun initializeOnRoleHoldersChanged() {
         roleManager.addOnRoleHoldersChangedListenerAsUser(
             backgroundExecutor,
             controller::onRoleHoldersChanged,
             UserHandle.ALL,
         )
     }
+
+    /**
+     * Initializes a [KeyguardUpdateMonitor] listener that will ensure [NoteTaskController] is in
+     * correct state during system initialization (after a direct boot user unlocked event).
+     *
+     * Once the system is unlocked, we will force trigger [NoteTaskController.onRoleHoldersChanged]
+     * with a hardcoded [RoleManager.ROLE_NOTES] for the current user.
+     */
+    private fun initializeOnUserUnlocked() {
+        if (keyguardUpdateMonitor.isUserUnlocked(userTracker.userId)) {
+            controller.setNoteTaskShortcutEnabled(true, userTracker.userHandle)
+        } else {
+            keyguardUpdateMonitor.registerCallback(onUserUnlockedCallback)
+        }
+    }
+
+    // KeyguardUpdateMonitor.registerCallback uses a weak reference, so we need a hard reference.
+    private val onUserUnlockedCallback =
+        object : KeyguardUpdateMonitorCallback() {
+            override fun onUserUnlocked() {
+                controller.setNoteTaskShortcutEnabled(true, userTracker.userHandle)
+                keyguardUpdateMonitor.removeCallback(this)
+            }
+        }
+}
+
+/**
+ * Maps a [KeyEvent] to a [NoteTaskEntryPoint]. If the [KeyEvent] does not represent a
+ * [NoteTaskEntryPoint], returns null.
+ */
+private fun KeyEvent.toNoteTaskEntryPointOrNull(): NoteTaskEntryPoint? =
+    when {
+        keyCode == KEYCODE_STYLUS_BUTTON_TAIL -> TAIL_BUTTON
+        keyCode == KEYCODE_N && isMetaPressed && isCtrlPressed -> KEYBOARD_SHORTCUT
+        else -> null
     }

packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskControllerTest.kt

@@ -108,6 +108,8 @@ internal class NoteTaskControllerTest : SysuiTestCase() {
         whenever(context.packageManager).thenReturn(packageManager)
         whenever(resolver.resolveInfo(any(), any(), any())).thenReturn(NOTE_TASK_INFO)
         whenever(userManager.isUserUnlocked).thenReturn(true)
+        whenever(userManager.isUserUnlocked(any<Int>())).thenReturn(true)
+        whenever(userManager.isUserUnlocked(any<UserHandle>())).thenReturn(true)
         whenever(
                 devicePolicyManager.getKeyguardDisabledFeatures(
                     /* admin= */ eq(null),

packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskInitializerTest.kt

@@ -16,124 +16,169 @@
 package com.android.systemui.notetask
 
 import android.app.role.RoleManager
-import android.test.suitebuilder.annotation.SmallTest
+import android.app.role.RoleManager.ROLE_NOTES
+import android.os.UserHandle
+import android.os.UserManager
+import android.testing.AndroidTestingRunner
 import android.view.KeyEvent
-import androidx.test.runner.AndroidJUnit4
+import android.view.KeyEvent.ACTION_DOWN
+import android.view.KeyEvent.KEYCODE_STYLUS_BUTTON_TAIL
+import androidx.test.filters.SmallTest
+import com.android.keyguard.KeyguardUpdateMonitor
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.settings.FakeUserTracker
 import com.android.systemui.statusbar.CommandQueue
 import com.android.systemui.util.concurrency.FakeExecutor
 import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.argumentCaptor
 import com.android.systemui.util.mockito.eq
+import com.android.systemui.util.mockito.whenever
 import com.android.systemui.util.time.FakeSystemClock
 import com.android.wm.shell.bubbles.Bubbles
 import java.util.Optional
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.ArgumentCaptor
 import org.mockito.Mock
 import org.mockito.Mockito.never
+import org.mockito.Mockito.times
 import org.mockito.Mockito.verify
 import org.mockito.Mockito.verifyZeroInteractions
-import org.mockito.MockitoAnnotations
+import org.mockito.MockitoAnnotations.initMocks
 
 /** atest SystemUITests:NoteTaskInitializerTest */
+@OptIn(ExperimentalCoroutinesApi::class, InternalNoteTaskApi::class)
 @SmallTest
-@RunWith(AndroidJUnit4::class)
+@RunWith(AndroidTestingRunner::class)
 internal class NoteTaskInitializerTest : SysuiTestCase() {
 
     @Mock lateinit var commandQueue: CommandQueue
     @Mock lateinit var bubbles: Bubbles
     @Mock lateinit var controller: NoteTaskController
     @Mock lateinit var roleManager: RoleManager
-    private val clock = FakeSystemClock()
+    @Mock lateinit var userManager: UserManager
-    private val executor = FakeExecutor(clock)
+    @Mock lateinit var keyguardMonitor: KeyguardUpdateMonitor
+
+    private val executor = FakeExecutor(FakeSystemClock())
     private val userTracker = FakeUserTracker()
 
     @Before
     fun setUp() {
-        MockitoAnnotations.initMocks(this)
+        initMocks(this)
+        whenever(keyguardMonitor.isUserUnlocked(userTracker.userId)).thenReturn(true)
     }
 
-    private fun createNoteTaskInitializer(
+    private fun createUnderTest(
-        isEnabled: Boolean = true,
+        isEnabled: Boolean,
-        bubbles: Bubbles? = this.bubbles,
+        bubbles: Bubbles?,
-    ): NoteTaskInitializer {
+    ): NoteTaskInitializer =
-        return NoteTaskInitializer(
+        NoteTaskInitializer(
             controller = controller,
             commandQueue = commandQueue,
             optionalBubbles = Optional.ofNullable(bubbles),
             isEnabled = isEnabled,
             roleManager = roleManager,
-            backgroundExecutor = executor,
             userTracker = userTracker,
+            keyguardUpdateMonitor = keyguardMonitor,
+            backgroundExecutor = executor,
         )
-    }
 
-    // region initializer
     @Test
-    fun initialize() {
+    fun initialize_withUserUnlocked() {
-        createNoteTaskInitializer().initialize()
+        whenever(keyguardMonitor.isUserUnlocked(userTracker.userId)).thenReturn(true)
+
+        createUnderTest(isEnabled = true, bubbles = bubbles).initialize()
 
-        verify(controller).setNoteTaskShortcutEnabled(eq(true), eq(userTracker.userHandle))
         verify(commandQueue).addCallback(any())
         verify(roleManager).addOnRoleHoldersChangedListenerAsUser(any(), any(), any())
+        verify(controller).setNoteTaskShortcutEnabled(any(), any())
+        verify(keyguardMonitor, never()).registerCallback(any())
+    }
+
+    @Test
+    fun initialize_withUserLocked() {
+        whenever(keyguardMonitor.isUserUnlocked(userTracker.userId)).thenReturn(false)
+
+        createUnderTest(isEnabled = true, bubbles = bubbles).initialize()
+
+        verify(commandQueue).addCallback(any())
+        verify(roleManager).addOnRoleHoldersChangedListenerAsUser(any(), any(), any())
+        verify(controller, never()).setNoteTaskShortcutEnabled(any(), any())
+        verify(keyguardMonitor).registerCallback(any())
     }
 
     @Test
     fun initialize_flagDisabled() {
-        createNoteTaskInitializer(isEnabled = false).initialize()
+        val underTest = createUnderTest(isEnabled = false, bubbles = bubbles)
 
-        verify(controller, never()).setNoteTaskShortcutEnabled(any(), any())
+        underTest.initialize()
-        verify(commandQueue, never()).addCallback(any())
+
-        verify(roleManager, never()).addOnRoleHoldersChangedListenerAsUser(any(), any(), any())
+        verifyZeroInteractions(
+            commandQueue,
+            bubbles,
+            controller,
+            roleManager,
+            userManager,
+            keyguardMonitor,
+        )
     }
 
     @Test
     fun initialize_bubblesNotPresent() {
-        createNoteTaskInitializer(bubbles = null).initialize()
+        val underTest = createUnderTest(isEnabled = true, bubbles = null)
 
-        verify(controller, never()).setNoteTaskShortcutEnabled(any(), any())
+        underTest.initialize()
-        verify(commandQueue, never()).addCallback(any())
-        verify(roleManager, never()).addOnRoleHoldersChangedListenerAsUser(any(), any(), any())
-    }
-    // endregion
 
-    // region handleSystemKey
+        verifyZeroInteractions(
-    @Test
+            commandQueue,
-    fun handleSystemKey_receiveValidSystemKey_shouldShowNoteTask() {
+            bubbles,
-        createNoteTaskInitializer()
+            controller,
-            .callbacks
+            roleManager,
-            .handleSystemKey(KeyEvent(KeyEvent.ACTION_DOWN, KeyEvent.KEYCODE_STYLUS_BUTTON_TAIL))
+            userManager,
-
+            keyguardMonitor,
-        verify(controller).showNoteTask(entryPoint = NoteTaskEntryPoint.TAIL_BUTTON)
-    }
-
-    @Test
-    fun handleSystemKey_receiveKeyboardShortcut_shouldShowNoteTask() {
-        createNoteTaskInitializer()
-            .callbacks
-            .handleSystemKey(
-                KeyEvent(
-                    0,
-                    0,
-                    KeyEvent.ACTION_DOWN,
-                    KeyEvent.KEYCODE_N,
-                    0,
-                    KeyEvent.META_META_ON or KeyEvent.META_CTRL_ON
         )
-            )
-
-        verify(controller).showNoteTask(entryPoint = NoteTaskEntryPoint.KEYBOARD_SHORTCUT)
     }
 
     @Test
-    fun handleSystemKey_receiveInvalidSystemKey_shouldDoNothing() {
+    fun initialize_handleSystemKey() {
-        createNoteTaskInitializer()
+        val expectedKeyEvent = KeyEvent(ACTION_DOWN, KEYCODE_STYLUS_BUTTON_TAIL)
-            .callbacks
+        val underTest = createUnderTest(isEnabled = true, bubbles = bubbles)
-            .handleSystemKey(KeyEvent(KeyEvent.ACTION_DOWN, KeyEvent.KEYCODE_UNKNOWN))
+        underTest.initialize()
+        val callback = captureArgument { verify(commandQueue).addCallback(capture()) }
 
-        verifyZeroInteractions(controller)
+        callback.handleSystemKey(expectedKeyEvent)
+
+        verify(controller).showNoteTask(any())
     }
-    // endregion
+
+    @Test
+    fun initialize_userUnlocked() {
+        whenever(keyguardMonitor.isUserUnlocked(userTracker.userId)).thenReturn(false)
+        val underTest = createUnderTest(isEnabled = true, bubbles = bubbles)
+        underTest.initialize()
+        val callback = captureArgument { verify(keyguardMonitor).registerCallback(capture()) }
+        whenever(keyguardMonitor.isUserUnlocked(userTracker.userId)).thenReturn(true)
+
+        callback.onUserUnlocked()
+        verify(controller).setNoteTaskShortcutEnabled(any(), any())
     }
+
+    @Test
+    fun initialize_onRoleHoldersChanged() {
+        val underTest = createUnderTest(isEnabled = true, bubbles = bubbles)
+        underTest.initialize()
+        val callback = captureArgument {
+            verify(roleManager)
+                .addOnRoleHoldersChangedListenerAsUser(any(), capture(), eq(UserHandle.ALL))
+        }
+
+        callback.onRoleHoldersChanged(ROLE_NOTES, userTracker.userHandle)
+
+        verify(controller).onRoleHoldersChanged(ROLE_NOTES, userTracker.userHandle)
+    }
+}
+
+private inline fun <reified T : Any> captureArgument(block: ArgumentCaptor<T>.() -> Unit) =
+    argumentCaptor<T>().apply(block).value