Allow focused window to override stem primary key.

This change hooks StemPrimaryKeyRule to DeferredKeyActionExecutor so
that stem primary key events get deferred until the focused app doesn't
handle the DOWN key event.

Several key points of this change:

1. Only send stem primary key events to apps with permission
   OVERRIDE_SYSTEM_KEY_BEHAVIOR_IN_FOCUSED_WINDOW.
2. We also need to send KEYCODE_STEM_PRIMARY to status bar only if it's
   not handled by the app.
3. We want to make sure the triple-press accessibility gesture always
   get triggered regardless of if the gesture is consumed by app.

Bug: 308482931
Test: atest WmTests:StemKeyGestureTests
      manually tested using a test app
Change-Id: I84791ca71416ec6c5d2f1c603c647031c76e059b

6 files changed, 240 insertions, 13 deletions

diff --git a/core/java/com/android/internal/policy/KeyInterceptionInfo.java b/core/java/com/android/internal/policy/KeyInterceptionInfo.java
index 964be01952ea..b20f6d225b69 100644
--- a/core/java/com/android/internal/policy/KeyInterceptionInfo.java
+++ b/core/java/com/android/internal/policy/KeyInterceptionInfo.java
@@ -26,10 +26,12 @@ public class KeyInterceptionInfo {
     public final int layoutParamsPrivateFlags;
     // Debug friendly name to help identify the window
     public final String windowTitle;
+    public final int windowOwnerUid;
 
-    public KeyInterceptionInfo(int type, int flags, String title) {
+    public KeyInterceptionInfo(int type, int flags, String title, int uid) {
         layoutParamsType = type;
         layoutParamsPrivateFlags = flags;
         windowTitle = title;
+        windowOwnerUid = uid;
     }
 }
diff --git a/services/core/java/com/android/server/policy/PhoneWindowManager.java b/services/core/java/com/android/server/policy/PhoneWindowManager.java
index 30bce2f41cf7..4e5dc1dd76fa 100644
--- a/services/core/java/com/android/server/policy/PhoneWindowManager.java
+++ b/services/core/java/com/android/server/policy/PhoneWindowManager.java
@@ -17,11 +17,13 @@
 package com.android.server.policy;
 
 import static android.Manifest.permission.INTERNAL_SYSTEM_WINDOW;
+import static android.Manifest.permission.OVERRIDE_SYSTEM_KEY_BEHAVIOR_IN_FOCUSED_WINDOW;
 import static android.Manifest.permission.SYSTEM_ALERT_WINDOW;
 import static android.Manifest.permission.SYSTEM_APPLICATION_OVERLAY;
 import static android.app.AppOpsManager.OP_CREATE_ACCESSIBILITY_OVERLAY;
 import static android.app.AppOpsManager.OP_SYSTEM_ALERT_WINDOW;
 import static android.app.AppOpsManager.OP_TOAST_WINDOW;
+import static android.content.PermissionChecker.PID_UNKNOWN;
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_LANDSCAPE;
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_REVERSE_LANDSCAPE;
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_SENSOR_LANDSCAPE;
@@ -117,6 +119,7 @@ import android.content.ContentResolver;
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
+import android.content.PermissionChecker;
 import android.content.pm.ActivityInfo;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
@@ -690,6 +693,8 @@ public class PhoneWindowManager implements WindowManagerPolicy {
 
     private final com.android.internal.policy.LogDecelerateInterpolator mLogDecelerateInterpolator
             = new LogDecelerateInterpolator(100, 0);
+    private final DeferredKeyActionExecutor mDeferredKeyActionExecutor =
+            new DeferredKeyActionExecutor();
 
     private volatile int mTopFocusedDisplayId = INVALID_DISPLAY;
 
@@ -698,6 +703,7 @@ public class PhoneWindowManager implements WindowManagerPolicy {
     private KeyCombinationManager mKeyCombinationManager;
     private SingleKeyGestureDetector mSingleKeyGestureDetector;
     private GestureLauncherService mGestureLauncherService;
+    private ButtonOverridePermissionChecker mButtonOverridePermissionChecker;
 
     private boolean mLockNowPending = false;
 
@@ -725,6 +731,7 @@ public class PhoneWindowManager implements WindowManagerPolicy {
     private static final int MSG_RINGER_TOGGLE_CHORD = 24;
     private static final int MSG_SWITCH_KEYBOARD_LAYOUT = 25;
     private static final int MSG_LOG_KEYBOARD_SYSTEM_EVENT = 26;
+    private static final int MSG_SET_DEFERRED_KEY_ACTIONS_EXECUTABLE = 27;
 
     private class PolicyHandler extends Handler {
 
@@ -792,7 +799,9 @@ public class PhoneWindowManager implements WindowManagerPolicy {
                     mAutofillManagerInternal.onBackKeyPressed();
                     break;
                 case MSG_SYSTEM_KEY_PRESS:
-                    sendSystemKeyToStatusBar((KeyEvent) msg.obj);
+                    KeyEvent event = (KeyEvent) msg.obj;
+                    sendSystemKeyToStatusBar(event);
+                    event.recycle();
                     break;
                 case MSG_HANDLE_ALL_APPS:
                     launchAllAppsAction();
@@ -809,6 +818,11 @@ public class PhoneWindowManager implements WindowManagerPolicy {
                 case MSG_LOG_KEYBOARD_SYSTEM_EVENT:
                     handleKeyboardSystemEvent(KeyboardLogEvent.from(msg.arg1), (KeyEvent) msg.obj);
                     break;
+                case MSG_SET_DEFERRED_KEY_ACTIONS_EXECUTABLE:
+                    final int keyCode = msg.arg1;
+                    final long downTime = (Long) msg.obj;
+                    mDeferredKeyActionExecutor.setActionsExecutable(keyCode, downTime);
+                    break;
             }
         }
     }
@@ -2234,6 +2248,10 @@ public class PhoneWindowManager implements WindowManagerPolicy {
         IActivityManager getActivityManagerService() {
             return ActivityManager.getService();
         }
+
+        ButtonOverridePermissionChecker getButtonOverridePermissionChecker() {
+            return new ButtonOverridePermissionChecker();
+        }
     }
 
     /** {@inheritDoc} */
@@ -2499,6 +2517,7 @@ public class PhoneWindowManager implements WindowManagerPolicy {
         mKeyguardDelegate = injector.getKeyguardServiceDelegate();
         initKeyCombinationRules();
         initSingleKeyGestureRules(injector.getLooper());
+        mButtonOverridePermissionChecker = injector.getButtonOverridePermissionChecker();
         mSideFpsEventHandler = new SideFpsEventHandler(mContext, mHandler, mPowerManager);
     }
 
@@ -2768,17 +2787,33 @@ public class PhoneWindowManager implements WindowManagerPolicy {
             if (mShouldEarlyShortPressOnStemPrimary) {
                 return;
             }
-            stemPrimaryPress(1 /*count*/);
+            // Short-press should be triggered only if app doesn't handle it.
+            mDeferredKeyActionExecutor.queueKeyAction(
+                    KeyEvent.KEYCODE_STEM_PRIMARY, downTime, () -> stemPrimaryPress(1 /*count*/));
         }
 
         @Override
         void onLongPress(long eventTime) {
-            stemPrimaryLongPress(eventTime);
+            // Long-press should be triggered only if app doesn't handle it.
+            mDeferredKeyActionExecutor.queueKeyAction(
+                    KeyEvent.KEYCODE_STEM_PRIMARY,
+                    eventTime,
+                    () -> stemPrimaryLongPress(eventTime));
         }
 
         @Override
         void onMultiPress(long downTime, int count, int unusedDisplayId) {
-            stemPrimaryPress(count);
+            // Triple-press stem to toggle accessibility gesture should always be triggered
+            // regardless of if app handles it.
+            if (count == 3
+                    && mTriplePressOnStemPrimaryBehavior
+                    == TRIPLE_PRESS_PRIMARY_TOGGLE_ACCESSIBILITY) {
+                stemPrimaryPress(count);
+            } else {
+                // Other multi-press gestures should be triggered only if app doesn't handle it.
+                mDeferredKeyActionExecutor.queueKeyAction(
+                        KeyEvent.KEYCODE_STEM_PRIMARY, downTime, () -> stemPrimaryPress(count));
+            }
         }
 
         @Override
@@ -2792,7 +2827,9 @@ public class PhoneWindowManager implements WindowManagerPolicy {
                 mBackgroundRecentTaskInfoOnStemPrimarySingleKeyUp =
                         mActivityTaskManagerInternal.getMostRecentTaskFromBackground();
                 if (mShouldEarlyShortPressOnStemPrimary) {
-                    stemPrimaryPress(1 /*pressCount*/);
+                    // Key-up gesture should be triggered only if app doesn't handle it.
+                    mDeferredKeyActionExecutor.queueKeyAction(
+                            KeyEvent.KEYCODE_STEM_PRIMARY, eventTime, () -> stemPrimaryPress(1));
                 }
             }
         }
@@ -3750,6 +3787,15 @@ public class PhoneWindowManager implements WindowManagerPolicy {
                     return true;
                 }
                 break;
+            case KeyEvent.KEYCODE_STEM_PRIMARY:
+                if (prepareToSendSystemKeyToApplication(focusedToken, event)) {
+                    // Send to app.
+                    return false;
+                } else {
+                    // Intercepted.
+                    sendSystemKeyToStatusBarAsync(event);
+                    return true;
+                }
         }
         if (isValidGlobalKey(keyCode)
                 && mGlobalKeyManager.handleGlobalKey(mContext, keyCode, event)) {
@@ -3760,6 +3806,60 @@ public class PhoneWindowManager implements WindowManagerPolicy {
         return (metaState & KeyEvent.META_META_ON) != 0;
     }
 
+    /**
+     * In this function, we check whether a system key should be sent to the application. We also
+     * detect the key gesture on this key, even if the key will be sent to the app. The gesture
+     * action, if any, will not be executed immediately. It will be queued and execute only after
+     * the application tells us that it didn't handle this key.
+     *
+     * @return true if this key should be sent to the application. This also means that the target
+     *     application has the necessary permissions to receive this key. Return false otherwise.
+     */
+    private boolean prepareToSendSystemKeyToApplication(IBinder focusedToken, KeyEvent event) {
+        final int keyCode = event.getKeyCode();
+        if (!event.isSystem()) {
+            Log.wtf(
+                    TAG,
+                    "Illegal keycode provided to prepareToSendSystemKeyToApplication: "
+                            + KeyEvent.keyCodeToString(keyCode));
+            return false;
+        }
+        final boolean isDown = event.getAction() == KeyEvent.ACTION_DOWN;
+        if (isDown && event.getRepeatCount() == 0) {
+            // This happens at the initial DOWN event. Check focused window permission now.
+            final KeyInterceptionInfo info =
+                    mWindowManagerInternal.getKeyInterceptionInfoFromToken(focusedToken);
+            if (info != null
+                    && mButtonOverridePermissionChecker.canAppOverrideSystemKey(
+                            mContext, info.windowOwnerUid)) {
+                // Focused window has the permission. Pass the event to it.
+                return true;
+            } else {
+                // Focused window doesn't have the permission. Intercept the event.
+                // If the initial DOWN event is intercepted, follow-up events will be intercepted
+                // too. So we know the gesture won't be handled by app, and can handle the gesture
+                // in system.
+                setDeferredKeyActionsExecutableAsync(keyCode, event.getDownTime());
+                return false;
+            }
+        } else {
+            // This happens after the initial DOWN event. We will just reuse the initial decision.
+            // I.e., if the initial DOWN event was dispatched, follow-up events should be
+            // dispatched. Otherwise, follow-up events should be consumed.
+            final Set<Integer> consumedKeys = mConsumedKeysForDevice.get(event.getDeviceId());
+            final boolean wasConsumed = consumedKeys != null && consumedKeys.contains(keyCode);
+            return !wasConsumed;
+        }
+    }
+
+    private void setDeferredKeyActionsExecutableAsync(int keyCode, long downTime) {
+        Message msg = Message.obtain(mHandler, MSG_SET_DEFERRED_KEY_ACTIONS_EXECUTABLE);
+        msg.arg1 = keyCode;
+        msg.obj = downTime;
+        msg.setAsynchronous(true);
+        msg.sendToTarget();
+    }
+
     @SuppressLint("MissingPermission")
     private void injectBackGesture(long downtime) {
         // Create and inject down event
@@ -3977,11 +4077,34 @@ public class PhoneWindowManager implements WindowManagerPolicy {
                     mContext.closeSystemDialogs();
                 }
                 return true;
+            case KeyEvent.KEYCODE_STEM_PRIMARY:
+                handleUnhandledSystemKey(event);
+                sendSystemKeyToStatusBarAsync(event);
+                return true;
         }
 
         return false;
     }
 
+    /**
+     * Called when a system key was sent to application and was unhandled. We will execute any
+     * queued actions associated with this key code at this point.
+     */
+    private void handleUnhandledSystemKey(KeyEvent event) {
+        if (!event.isSystem()) {
+            Log.wtf(
+                    TAG,
+                    "Illegal keycode provided to handleUnhandledSystemKey: "
+                            + KeyEvent.keyCodeToString(event.getKeyCode()));
+            return;
+        }
+        if (event.getAction() == KeyEvent.ACTION_DOWN && event.getRepeatCount() == 0) {
+            // If the initial DOWN event is unhandled by app, follow-up events will also be
+            // unhandled by app. So we can handle the key event in system.
+            setDeferredKeyActionsExecutableAsync(event.getKeyCode(), event.getDownTime());
+        }
+    }
+
     private void sendSwitchKeyboardLayout(@NonNull KeyEvent event, int direction) {
         mHandler.obtainMessage(MSG_SWITCH_KEYBOARD_LAYOUT, event.getDeviceId(),
                 direction).sendToTarget();
@@ -4904,9 +5027,6 @@ public class PhoneWindowManager implements WindowManagerPolicy {
             case KeyEvent.KEYCODE_MACRO_4:
                 result &= ~ACTION_PASS_TO_USER;
                 break;
-            case KeyEvent.KEYCODE_STEM_PRIMARY:
-                sendSystemKeyToStatusBarAsync(event);
-                break;
         }
 
         if (useHapticFeedback) {
@@ -5016,7 +5136,8 @@ public class PhoneWindowManager implements WindowManagerPolicy {
      * Notify the StatusBar that a system key was pressed without blocking the current thread.
      */
     private void sendSystemKeyToStatusBarAsync(KeyEvent keyEvent) {
-        Message message = mHandler.obtainMessage(MSG_SYSTEM_KEY_PRESS, keyEvent);
+        // Make a copy because the event may be recycled.
+        Message message = mHandler.obtainMessage(MSG_SYSTEM_KEY_PRESS, KeyEvent.obtain(keyEvent));
         message.setAsynchronous(true);
         mHandler.sendMessage(message);
     }
@@ -6468,6 +6589,7 @@ public class PhoneWindowManager implements WindowManagerPolicy {
         mGlobalKeyManager.dump(prefix, pw);
         mKeyCombinationManager.dump(prefix, pw);
         mSingleKeyGestureDetector.dump(prefix, pw);
+        mDeferredKeyActionExecutor.dump(prefix, pw);
 
         if (mWakeGestureListener != null) {
             mWakeGestureListener.dump(pw, prefix);
@@ -6793,4 +6915,19 @@ public class PhoneWindowManager implements WindowManagerPolicy {
                     + " name.");
         }
     }
+
+    /** A helper class to check button override permission. */
+    static class ButtonOverridePermissionChecker {
+        boolean canAppOverrideSystemKey(Context context, int uid) {
+            return PermissionChecker.checkPermissionForDataDelivery(
+                            context,
+                            OVERRIDE_SYSTEM_KEY_BEHAVIOR_IN_FOCUSED_WINDOW,
+                            PID_UNKNOWN,
+                            uid,
+                            null,
+                            null,
+                            null)
+                    == PERMISSION_GRANTED;
+        }
+    }
 }
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index 7bc7e2cb780b..e87b7926ccd8 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -5630,9 +5630,10 @@ class WindowState extends WindowContainer<WindowState> implements WindowManagerP
         if (mKeyInterceptionInfo == null
                 || mKeyInterceptionInfo.layoutParamsPrivateFlags != getAttrs().privateFlags
                 || mKeyInterceptionInfo.layoutParamsType != getAttrs().type
-                || mKeyInterceptionInfo.windowTitle != getWindowTag()) {
+                || mKeyInterceptionInfo.windowTitle != getWindowTag()
+                || mKeyInterceptionInfo.windowOwnerUid != getOwningUid()) {
             mKeyInterceptionInfo = new KeyInterceptionInfo(getAttrs().type, getAttrs().privateFlags,
-                    getWindowTag().toString());
+                    getWindowTag().toString(), getOwningUid());
         }
         return mKeyInterceptionInfo;
     }
diff --git a/services/tests/wmtests/src/com/android/server/policy/ShortcutKeyTestBase.java b/services/tests/wmtests/src/com/android/server/policy/ShortcutKeyTestBase.java
index ab35da69da7c..9cdec2588501 100644
--- a/services/tests/wmtests/src/com/android/server/policy/ShortcutKeyTestBase.java
+++ b/services/tests/wmtests/src/com/android/server/policy/ShortcutKeyTestBase.java
@@ -64,6 +64,7 @@ class ShortcutKeyTestBase {
     @Rule public FakeSettingsProviderRule mSettingsProviderRule = FakeSettingsProvider.rule();
 
     TestPhoneWindowManager mPhoneWindowManager;
+    DispatchedKeyHandler mDispatchedKeyHandler = event -> false;
     final Context mContext = spy(getInstrumentation().getTargetContext());
 
     /** Modifier key to meta state */
@@ -102,6 +103,10 @@ class ShortcutKeyTestBase {
         mPhoneWindowManager = new TestPhoneWindowManager(mContext, supportSettingsUpdate);
     }
 
+    protected final void setDispatchedKeyHandler(DispatchedKeyHandler keyHandler) {
+        mDispatchedKeyHandler = keyHandler;
+    }
+
     @After
     public void tearDown() {
         if (mPhoneWindowManager != null) {
@@ -174,9 +179,20 @@ class ShortcutKeyTestBase {
         int actions = mPhoneWindowManager.interceptKeyBeforeQueueing(keyEvent);
         if ((actions & ACTION_PASS_TO_USER) != 0) {
             if (0 == mPhoneWindowManager.interceptKeyBeforeDispatching(keyEvent)) {
-                mPhoneWindowManager.dispatchUnhandledKey(keyEvent);
+                if (!mDispatchedKeyHandler.onKeyDispatched(keyEvent)) {
+                    mPhoneWindowManager.dispatchUnhandledKey(keyEvent);
+                }
             }
         }
         mPhoneWindowManager.dispatchAllPendingEvents();
     }
+
+    interface DispatchedKeyHandler {
+        /**
+         * Called when a key event is dispatched to app.
+         *
+         * @return true if the event is consumed by app.
+         */
+        boolean onKeyDispatched(KeyEvent event);
+    }
 }
diff --git a/services/tests/wmtests/src/com/android/server/policy/StemKeyGestureTests.java b/services/tests/wmtests/src/com/android/server/policy/StemKeyGestureTests.java
index 912e1d3df945..f7ad2a8f5243 100644
--- a/services/tests/wmtests/src/com/android/server/policy/StemKeyGestureTests.java
+++ b/services/tests/wmtests/src/com/android/server/policy/StemKeyGestureTests.java
@@ -97,6 +97,35 @@ public class StemKeyGestureTests extends ShortcutKeyTestBase {
     }
 
     @Test
+    public void stemSingleKey_appHasOverridePermission_consumedByApp_notOpenAllApp() {
+        overrideBehavior(STEM_PRIMARY_BUTTON_SHORT_PRESS, SHORT_PRESS_PRIMARY_LAUNCH_ALL_APPS);
+        setUpPhoneWindowManager(/* supportSettingsUpdate= */ true);
+        mPhoneWindowManager.overrideStartActivity();
+        mPhoneWindowManager.setKeyguardServiceDelegateIsShowing(false);
+        mPhoneWindowManager.overrideIsUserSetupComplete(true);
+        mPhoneWindowManager.overrideFocusedWindowButtonOverridePermission(true);
+        setDispatchedKeyHandler(keyEvent -> true);
+
+        sendKey(KEYCODE_STEM_PRIMARY);
+
+        mPhoneWindowManager.assertNotOpenAllAppView();
+    }
+
+    @Test
+    public void stemSingleKey_appHasOverridePermission_notConsumedByApp_openAllApp() {
+        overrideBehavior(STEM_PRIMARY_BUTTON_SHORT_PRESS, SHORT_PRESS_PRIMARY_LAUNCH_ALL_APPS);
+        setUpPhoneWindowManager(/* supportSettingsUpdate= */ true);
+        mPhoneWindowManager.overrideStartActivity();
+        mPhoneWindowManager.setKeyguardServiceDelegateIsShowing(false);
+        mPhoneWindowManager.overrideIsUserSetupComplete(true);
+        mPhoneWindowManager.overrideFocusedWindowButtonOverridePermission(true);
+
+        sendKey(KEYCODE_STEM_PRIMARY);
+
+        mPhoneWindowManager.assertOpenAllAppView();
+    }
+
+    @Test
     public void stemLongKey_triggerSearchServiceToLaunchAssist() {
         overrideBehavior(
                 STEM_PRIMARY_BUTTON_LONG_PRESS,
@@ -165,6 +194,30 @@ public class StemKeyGestureTests extends ShortcutKeyTestBase {
         mPhoneWindowManager.assertSwitchToRecent(referenceId);
     }
 
+    @Test
+    public void stemDoubleKey_earlyShortPress_firstPressConsumedByApp_switchToMostRecent()
+            throws RemoteException {
+        overrideBehavior(STEM_PRIMARY_BUTTON_DOUBLE_PRESS, SHORT_PRESS_PRIMARY_LAUNCH_ALL_APPS);
+        setUpPhoneWindowManager(/* supportSettingsUpdate= */ true);
+        mPhoneWindowManager.overrideShouldEarlyShortPressOnStemPrimary(true);
+        mPhoneWindowManager.setKeyguardServiceDelegateIsShowing(false);
+        mPhoneWindowManager.overrideIsUserSetupComplete(true);
+        mPhoneWindowManager.overrideFocusedWindowButtonOverridePermission(true);
+        RecentTaskInfo recentTaskInfo = new RecentTaskInfo();
+        int referenceId = 666;
+        recentTaskInfo.persistentId = referenceId;
+        doReturn(recentTaskInfo).when(
+                mPhoneWindowManager.mActivityTaskManagerInternal).getMostRecentTaskFromBackground();
+
+        setDispatchedKeyHandler(keyEvent -> true);
+        sendKey(KEYCODE_STEM_PRIMARY);
+        setDispatchedKeyHandler(keyEvent -> false);
+        sendKey(KEYCODE_STEM_PRIMARY);
+
+        mPhoneWindowManager.assertNotOpenAllAppView();
+        mPhoneWindowManager.assertSwitchToRecent(referenceId);
+    }
+
     private void overrideBehavior(String key, int expectedBehavior) {
         Settings.Global.putLong(mContext.getContentResolver(), key, expectedBehavior);
     }
diff --git a/services/tests/wmtests/src/com/android/server/policy/TestPhoneWindowManager.java b/services/tests/wmtests/src/com/android/server/policy/TestPhoneWindowManager.java
index 43c47458d19f..d057226836a3 100644
--- a/services/tests/wmtests/src/com/android/server/policy/TestPhoneWindowManager.java
+++ b/services/tests/wmtests/src/com/android/server/policy/TestPhoneWindowManager.java
@@ -90,6 +90,7 @@ import android.view.autofill.AutofillManagerInternal;
 
 import com.android.dx.mockito.inline.extended.StaticMockitoSession;
 import com.android.internal.accessibility.AccessibilityShortcutController;
+import com.android.internal.policy.KeyInterceptionInfo;
 import com.android.internal.util.FrameworkStatsLog;
 import com.android.server.GestureLauncherService;
 import com.android.server.LocalServices;
@@ -162,6 +163,9 @@ class TestPhoneWindowManager {
 
     @Mock private KeyguardServiceDelegate mKeyguardServiceDelegate;
 
+    @Mock
+    private PhoneWindowManager.ButtonOverridePermissionChecker mButtonOverridePermissionChecker;
+
     private StaticMockitoSession mMockitoSession;
     private OffsettableClock mClock = new OffsettableClock();
     private TestLooper mTestLooper = new TestLooper(() -> mClock.now());
@@ -189,6 +193,10 @@ class TestPhoneWindowManager {
         IActivityManager getActivityManagerService() {
             return mActivityManagerService;
         }
+
+        PhoneWindowManager.ButtonOverridePermissionChecker getButtonOverridePermissionChecker() {
+            return mButtonOverridePermissionChecker;
+        }
     }
 
     TestPhoneWindowManager(Context context, boolean supportSettingsUpdate) {
@@ -304,6 +312,11 @@ class TestPhoneWindowManager {
         doReturn(false).when(mPhoneWindowManager).keyguardOn();
         doNothing().when(mContext).startActivityAsUser(any(), any());
         doNothing().when(mContext).startActivityAsUser(any(), any(), any());
+
+        KeyInterceptionInfo interceptionInfo = new KeyInterceptionInfo(0, 0, null, 0);
+        doReturn(interceptionInfo)
+                .when(mWindowManagerInternal).getKeyInterceptionInfoFromToken(any());
+
         Mockito.reset(mContext);
     }
 
@@ -525,6 +538,11 @@ class TestPhoneWindowManager {
         mPhoneWindowManager.mPrimaryShortPressTargetActivity = component;
     }
 
+    void overrideFocusedWindowButtonOverridePermission(boolean granted) {
+        doReturn(granted)
+                .when(mButtonOverridePermissionChecker).canAppOverrideSystemKey(any(), anyInt());
+    }
+
     /**
      * Below functions will check the policy behavior could be invoked.
      */