Fix an issue which causes back key doesn't work

This issue is caused ag/5777021, which drops keys unexpectedly.

Change-Id: I23122dce9032555385fa1091527317d4bec92326
Fix: 121111244
Test: Manual
diff --git a/services/core/java/com/android/server/policy/PhoneWindowManager.java b/services/core/java/com/android/server/policy/PhoneWindowManager.java
index ebe202e..f370edf 100644
--- a/services/core/java/com/android/server/policy/PhoneWindowManager.java
+++ b/services/core/java/com/android/server/policy/PhoneWindowManager.java
@@ -2549,12 +2549,15 @@
     @Override
     public long interceptKeyBeforeDispatching(WindowState win, KeyEvent event, int policyFlags) {
         final long result = interceptKeyBeforeDispatchingInner(win, event, policyFlags);
+        final int eventDisplayId = event.getDisplayId();
         if (result == 0 && !mPerDisplayFocusEnabled
-                && event.getDisplayId() != mTopFocusedDisplayId) {
+                && eventDisplayId != INVALID_DISPLAY && eventDisplayId != mTopFocusedDisplayId) {
             // Someone tries to send a key event to a display which doesn't have a focused window.
             // We drop the event here, or it will cause ANR.
             // TODO (b/121057974): The user may be confused about why the key doesn't work, so we
             // may need to deal with this problem.
+            Slog.i(TAG, "Dropping this event targeting display #" + eventDisplayId
+                    + " because the focus is on display #" + mTopFocusedDisplayId);
             return -1;
         }
         return result;
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index 6e4f69e..801e5f2 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -184,6 +184,7 @@
         if (mTopFocusedDisplayId != topFocusedDisplayId) {
             mTopFocusedDisplayId = topFocusedDisplayId;
             mWmService.mInputManager.setFocusedDisplay(topFocusedDisplayId);
+            mWmService.mPolicy.setTopFocusedDisplay(topFocusedDisplayId);
             if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "New topFocusedDisplayId="
                     + topFocusedDisplayId);
         }