Check keyguard policy before launching note shortcut at lock screen

Test: atest SystemUITests:NoteTaskControllerTest
atest SystemUITests:DevicePolicyManagersTest
atest SystemUITests:KeyguardQuickAffordanceInteractorTest
Bug: 268210531

Change-Id: Id69047e37aa845ef95e26a8fb34ddf7ad0c094cf
diff --git a/packages/SystemUI/src/com/android/systemui/devicepolicy/DevicePolicyManagerExt.kt b/packages/SystemUI/src/com/android/systemui/devicepolicy/DevicePolicyManagerExt.kt
new file mode 100644
index 0000000..1390b4d
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/devicepolicy/DevicePolicyManagerExt.kt
@@ -0,0 +1,32 @@
+/*
+ * 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.systemui.devicepolicy
+
+import android.app.admin.DevicePolicyManager
+import android.app.admin.DevicePolicyManager.KEYGUARD_DISABLE_FEATURES_ALL
+import android.app.admin.DevicePolicyManager.KEYGUARD_DISABLE_SHORTCUTS_ALL
+import android.content.ComponentName
+
+/** Returns true if the admin of [userId] disallows keyguard shortcuts. */
+fun DevicePolicyManager.areKeyguardShortcutsDisabled(
+    admin: ComponentName? = null,
+    userId: Int
+): Boolean {
+    val flags = getKeyguardDisabledFeatures(admin, userId)
+    return flags and KEYGUARD_DISABLE_SHORTCUTS_ALL == KEYGUARD_DISABLE_SHORTCUTS_ALL ||
+        flags and KEYGUARD_DISABLE_FEATURES_ALL == KEYGUARD_DISABLE_FEATURES_ALL
+}
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt
index dfbe1c2..9b5f7f6 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt
@@ -26,6 +26,7 @@
 import com.android.systemui.animation.Expandable
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.devicepolicy.areKeyguardShortcutsDisabled
 import com.android.systemui.flags.FeatureFlags
 import com.android.systemui.flags.Flags
 import com.android.systemui.keyguard.data.quickaffordance.KeyguardQuickAffordanceConfig
@@ -410,16 +411,10 @@
         )
     }
 
-    private suspend fun isFeatureDisabledByDevicePolicy(): Boolean {
-        val flags =
-            withContext(backgroundDispatcher) {
-                devicePolicyManager.getKeyguardDisabledFeatures(null, userTracker.userId)
-            }
-        val flagsToCheck =
-            DevicePolicyManager.KEYGUARD_DISABLE_FEATURES_ALL or
-                DevicePolicyManager.KEYGUARD_DISABLE_SHORTCUTS_ALL
-        return flagsToCheck and flags != 0
-    }
+    private suspend fun isFeatureDisabledByDevicePolicy(): Boolean =
+        withContext(backgroundDispatcher) {
+            devicePolicyManager.areKeyguardShortcutsDisabled(userId = userTracker.userId)
+        }
 
     companion object {
         private const val TAG = "KeyguardQuickAffordanceInteractor"
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskController.kt b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskController.kt
index be615d6..5702681 100644
--- a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskController.kt
+++ b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskController.kt
@@ -17,17 +17,21 @@
 package com.android.systemui.notetask
 
 import android.app.KeyguardManager
+import android.app.admin.DevicePolicyManager
 import android.content.ActivityNotFoundException
 import android.content.ComponentName
 import android.content.Context
 import android.content.Intent
 import android.content.pm.PackageManager
+import android.os.Build
 import android.os.UserManager
 import android.util.Log
 import com.android.internal.logging.UiEvent
 import com.android.internal.logging.UiEventLogger
 import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.devicepolicy.areKeyguardShortcutsDisabled
 import com.android.systemui.notetask.shortcut.CreateNoteTaskShortcutActivity
+import com.android.systemui.settings.UserTracker
 import com.android.systemui.util.kotlin.getOrNull
 import com.android.wm.shell.bubbles.Bubbles
 import java.util.Optional
@@ -49,8 +53,10 @@
     private val optionalBubbles: Optional<Bubbles>,
     private val optionalKeyguardManager: Optional<KeyguardManager>,
     private val optionalUserManager: Optional<UserManager>,
+    private val devicePolicyManager: DevicePolicyManager,
     @NoteTaskEnabledKey private val isEnabled: Boolean,
     private val uiEventLogger: UiEventLogger,
+    private val userTracker: UserTracker,
 ) {
 
     /**
@@ -80,6 +86,18 @@
         // TODO(b/249954038): We should handle direct boot (isUserUnlocked). For now, we do nothing.
         if (!userManager.isUserUnlocked) return
 
+        val isKeyguardLocked = keyguardManager.isKeyguardLocked
+        // KeyguardQuickAffordanceInteractor blocks the quick affordance from showing in the
+        // keyguard if it is not allowed by the admin policy. Here we block any other way to show
+        // note task when the screen is locked.
+        if (
+            isKeyguardLocked &&
+                devicePolicyManager.areKeyguardShortcutsDisabled(userId = userTracker.userId)
+        ) {
+            logDebug { "Enterprise policy disallows launching note app when the screen is locked." }
+            return
+        }
+
         val noteTaskInfo = resolver.resolveInfo() ?: return
 
         uiEvent?.let { uiEventLogger.log(it, noteTaskInfo.uid, noteTaskInfo.packageName) }
@@ -87,7 +105,7 @@
         // TODO(b/266686199): We should handle when app not available. For now, we log.
         val intent = noteTaskInfo.toCreateNoteIntent()
         try {
-            if (isInMultiWindowMode || keyguardManager.isKeyguardLocked) {
+            if (isInMultiWindowMode || isKeyguardLocked) {
                 context.startActivity(intent)
             } else {
                 bubbles.showOrHideAppBubble(intent)
@@ -144,7 +162,7 @@
     }
 
     companion object {
-        private val TAG = NoteTaskController::class.simpleName.orEmpty()
+        val TAG = NoteTaskController::class.simpleName.orEmpty()
 
         private fun NoteTaskInfoResolver.NoteTaskInfo.toCreateNoteIntent(): Intent {
             return Intent(ACTION_CREATE_NOTE)
@@ -165,3 +183,9 @@
         const val INTENT_EXTRA_USE_STYLUS_MODE = "android.intent.extra.USE_STYLUS_MODE"
     }
 }
+
+private inline fun logDebug(message: () -> String) {
+    if (Build.IS_DEBUGGABLE) {
+        Log.d(NoteTaskController.TAG, message())
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/devicepolicy/DevicePolicyManagerExtTest.kt b/packages/SystemUI/tests/src/com/android/systemui/devicepolicy/DevicePolicyManagerExtTest.kt
new file mode 100644
index 0000000..34d5661
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/devicepolicy/DevicePolicyManagerExtTest.kt
@@ -0,0 +1,93 @@
+/*
+ * 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.systemui.devicepolicy
+
+import android.app.admin.DevicePolicyManager
+import android.app.admin.DevicePolicyManager.KEYGUARD_DISABLE_FACE
+import android.app.admin.DevicePolicyManager.KEYGUARD_DISABLE_FEATURES_ALL
+import android.app.admin.DevicePolicyManager.KEYGUARD_DISABLE_FEATURES_NONE
+import android.app.admin.DevicePolicyManager.KEYGUARD_DISABLE_SECURE_CAMERA
+import android.app.admin.DevicePolicyManager.KEYGUARD_DISABLE_SHORTCUTS_ALL
+import androidx.test.filters.SmallTest
+import com.android.systemui.settings.UserTracker
+import com.android.systemui.util.mockito.eq
+import com.android.systemui.util.mockito.whenever
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+import org.mockito.ArgumentMatchers.anyInt
+import org.mockito.Mock
+import org.mockito.MockitoAnnotations
+
+@SmallTest
+@RunWith(JUnit4::class)
+class DevicePolicyManagerExtTest {
+
+    @Mock lateinit var devicePolicyManager: DevicePolicyManager
+    @Mock private lateinit var userTracker: UserTracker
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+
+        whenever(userTracker.userId).thenReturn(CURRENT_USER_ID)
+    }
+
+    // region areKeyguardShortcutsDisabled
+    @Test
+    fun areKeyguardShortcutsDisabled_noDisabledKeyguardFeature_shouldReturnFalse() {
+        whenever(devicePolicyManager.getKeyguardDisabledFeatures(eq(null), anyInt()))
+            .thenReturn(KEYGUARD_DISABLE_FEATURES_NONE)
+
+        assertThat(devicePolicyManager.areKeyguardShortcutsDisabled(userId = CURRENT_USER_ID))
+            .isFalse()
+    }
+
+    @Test
+    fun areKeyguardShortcutsDisabled_otherDisabledKeyguardFeatures_shouldReturnFalse() {
+        whenever(devicePolicyManager.getKeyguardDisabledFeatures(eq(null), anyInt()))
+            .thenReturn(KEYGUARD_DISABLE_SECURE_CAMERA or KEYGUARD_DISABLE_FACE)
+
+        assertThat(devicePolicyManager.areKeyguardShortcutsDisabled(userId = CURRENT_USER_ID))
+            .isFalse()
+    }
+
+    @Test
+    fun areKeyguardShortcutsDisabled_disabledShortcutsKeyguardFeature_shouldReturnTrue() {
+        whenever(devicePolicyManager.getKeyguardDisabledFeatures(eq(null), anyInt()))
+            .thenReturn(KEYGUARD_DISABLE_SHORTCUTS_ALL)
+
+        assertThat(devicePolicyManager.areKeyguardShortcutsDisabled(userId = CURRENT_USER_ID))
+            .isTrue()
+    }
+
+    @Test
+    fun areKeyguardShortcutsDisabled_disabledAllKeyguardFeatures_shouldReturnTrue() {
+        whenever(devicePolicyManager.getKeyguardDisabledFeatures(eq(null), anyInt()))
+            .thenReturn(KEYGUARD_DISABLE_FEATURES_ALL)
+
+        assertThat(devicePolicyManager.areKeyguardShortcutsDisabled(userId = CURRENT_USER_ID))
+            .isTrue()
+    }
+    // endregion
+
+    private companion object {
+        const val CURRENT_USER_ID = 123
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskControllerTest.kt
index 39c4e06..52b29ac 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskControllerTest.kt
@@ -16,6 +16,7 @@
 package com.android.systemui.notetask
 
 import android.app.KeyguardManager
+import android.app.admin.DevicePolicyManager
 import android.content.ComponentName
 import android.content.Context
 import android.content.Intent
@@ -29,6 +30,7 @@
 import com.android.systemui.notetask.NoteTaskController.ShowNoteTaskUiEvent
 import com.android.systemui.notetask.NoteTaskInfoResolver.NoteTaskInfo
 import com.android.systemui.notetask.shortcut.CreateNoteTaskShortcutActivity
+import com.android.systemui.settings.UserTracker
 import com.android.systemui.util.mockito.argumentCaptor
 import com.android.systemui.util.mockito.capture
 import com.android.systemui.util.mockito.eq
@@ -39,6 +41,7 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.Mock
 import org.mockito.Mockito.verify
 import org.mockito.Mockito.verifyZeroInteractions
@@ -64,6 +67,8 @@
     @Mock lateinit var optionalUserManager: Optional<UserManager>
     @Mock lateinit var userManager: UserManager
     @Mock lateinit var uiEventLogger: UiEventLogger
+    @Mock private lateinit var userTracker: UserTracker
+    @Mock private lateinit var devicePolicyManager: DevicePolicyManager
 
     @Before
     fun setUp() {
@@ -75,6 +80,13 @@
         whenever(optionalKeyguardManager.orElse(null)).thenReturn(keyguardManager)
         whenever(optionalUserManager.orElse(null)).thenReturn(userManager)
         whenever(userManager.isUserUnlocked).thenReturn(true)
+        whenever(
+                devicePolicyManager.getKeyguardDisabledFeatures(
+                    /* admin= */ eq(null),
+                    /* userHandle= */ anyInt()
+                )
+            )
+            .thenReturn(DevicePolicyManager.KEYGUARD_DISABLE_FEATURES_NONE)
     }
 
     private fun createNoteTaskController(isEnabled: Boolean = true): NoteTaskController {
@@ -84,8 +96,10 @@
             optionalBubbles = optionalBubbles,
             optionalKeyguardManager = optionalKeyguardManager,
             optionalUserManager = optionalUserManager,
+            devicePolicyManager = devicePolicyManager,
             isEnabled = isEnabled,
             uiEventLogger = uiEventLogger,
+            userTracker = userTracker,
         )
     }
 
@@ -291,6 +305,86 @@
     }
     // endregion
 
+    // region keyguard policy
+    @Test
+    fun showNoteTask_keyguardLocked_keyguardDisableShortcutsAll_shouldDoNothing() {
+        whenever(keyguardManager.isKeyguardLocked).thenReturn(true)
+        whenever(
+                devicePolicyManager.getKeyguardDisabledFeatures(
+                    /* admin= */ eq(null),
+                    /* userHandle= */ anyInt()
+                )
+            )
+            .thenReturn(DevicePolicyManager.KEYGUARD_DISABLE_SHORTCUTS_ALL)
+
+        createNoteTaskController().showNoteTask(isInMultiWindowMode = false, uiEvent = null)
+
+        verifyZeroInteractions(context, bubbles, uiEventLogger)
+    }
+
+    @Test
+    fun showNoteTask_keyguardLocked_keyguardDisableFeaturesAll_shouldDoNothing() {
+        whenever(keyguardManager.isKeyguardLocked).thenReturn(true)
+        whenever(
+                devicePolicyManager.getKeyguardDisabledFeatures(
+                    /* admin= */ eq(null),
+                    /* userHandle= */ anyInt()
+                )
+            )
+            .thenReturn(DevicePolicyManager.KEYGUARD_DISABLE_FEATURES_ALL)
+
+        createNoteTaskController().showNoteTask(isInMultiWindowMode = false, uiEvent = null)
+
+        verifyZeroInteractions(context, bubbles, uiEventLogger)
+    }
+
+    @Test
+    fun showNoteTask_keyguardUnlocked_keyguardDisableShortcutsAll_shouldStartBubble() {
+        whenever(keyguardManager.isKeyguardLocked).thenReturn(false)
+        whenever(
+                devicePolicyManager.getKeyguardDisabledFeatures(
+                    /* admin= */ eq(null),
+                    /* userHandle= */ anyInt()
+                )
+            )
+            .thenReturn(DevicePolicyManager.KEYGUARD_DISABLE_SHORTCUTS_ALL)
+
+        createNoteTaskController().showNoteTask(isInMultiWindowMode = false, uiEvent = null)
+
+        val intentCaptor = argumentCaptor<Intent>()
+        verify(bubbles).showOrHideAppBubble(capture(intentCaptor))
+        intentCaptor.value.let { intent ->
+            assertThat(intent.action).isEqualTo(NoteTaskController.ACTION_CREATE_NOTE)
+            assertThat(intent.`package`).isEqualTo(NOTES_PACKAGE_NAME)
+            assertThat(intent.flags).isEqualTo(Intent.FLAG_ACTIVITY_NEW_TASK)
+            assertThat(intent.getBooleanExtra(INTENT_EXTRA_USE_STYLUS_MODE, false)).isTrue()
+        }
+    }
+
+    @Test
+    fun showNoteTask_keyguardUnlocked_keyguardDisableFeaturesAll_shouldStartBubble() {
+        whenever(keyguardManager.isKeyguardLocked).thenReturn(false)
+        whenever(
+                devicePolicyManager.getKeyguardDisabledFeatures(
+                    /* admin= */ eq(null),
+                    /* userHandle= */ anyInt()
+                )
+            )
+            .thenReturn(DevicePolicyManager.KEYGUARD_DISABLE_FEATURES_ALL)
+
+        createNoteTaskController().showNoteTask(isInMultiWindowMode = false, uiEvent = null)
+
+        val intentCaptor = argumentCaptor<Intent>()
+        verify(bubbles).showOrHideAppBubble(capture(intentCaptor))
+        intentCaptor.value.let { intent ->
+            assertThat(intent.action).isEqualTo(NoteTaskController.ACTION_CREATE_NOTE)
+            assertThat(intent.`package`).isEqualTo(NOTES_PACKAGE_NAME)
+            assertThat(intent.flags).isEqualTo(Intent.FLAG_ACTIVITY_NEW_TASK)
+            assertThat(intent.getBooleanExtra(INTENT_EXTRA_USE_STYLUS_MODE, false)).isTrue()
+        }
+    }
+    // endregion
+
     private companion object {
         const val NOTES_PACKAGE_NAME = "com.android.note.app"
         const val NOTES_UID = 123456