Merge "Instantiate InputMethodManager for each display"

7 files changed, 292 insertions, 49 deletions

diff --git a/core/java/android/app/SystemServiceRegistry.java b/core/java/android/app/SystemServiceRegistry.java
index 0044005c51f2..6c87fe75740e 100644
--- a/core/java/android/app/SystemServiceRegistry.java
+++ b/core/java/android/app/SystemServiceRegistry.java
@@ -377,11 +377,15 @@ final class SystemServiceRegistry {
                 return new DisplayManager(ctx.getOuterContext());
             }});
 
+        // InputMethodManager has its own cache strategy based on display id to support apps that
+        // still assume InputMethodManager is a per-process singleton and it's safe to directly
+        // access internal fields via reflection.  Hence directly use ServiceFetcher instead of
+        // StaticServiceFetcher/CachedServiceFetcher.
         registerService(Context.INPUT_METHOD_SERVICE, InputMethodManager.class,
-                new StaticServiceFetcher<InputMethodManager>() {
+                new ServiceFetcher<InputMethodManager>() {
             @Override
-            public InputMethodManager createService() {
-                return InputMethodManager.getInstanceInternal();
+            public InputMethodManager getService(ContextImpl ctx) {
+                return InputMethodManager.forContext(ctx);
             }});
 
         registerService(Context.TEXT_SERVICES_MANAGER_SERVICE, TextServicesManager.class,
diff --git a/core/java/android/view/inputmethod/InputMethodManager.java b/core/java/android/view/inputmethod/InputMethodManager.java
index ca2ccaf224db..08ed9d17fb77 100644
--- a/core/java/android/view/inputmethod/InputMethodManager.java
+++ b/core/java/android/view/inputmethod/InputMethodManager.java
@@ -48,6 +48,7 @@ import android.util.Pools.SimplePool;
 import android.util.PrintWriterPrinter;
 import android.util.Printer;
 import android.util.SparseArray;
+import android.view.Display;
 import android.view.InputChannel;
 import android.view.InputEvent;
 import android.view.InputEventSender;
@@ -265,7 +266,7 @@ public final class InputMethodManager {
      * @hide
      */
     public static void ensureDefaultInstanceForDefaultDisplayIfNecessary() {
-        getInstanceInternal();
+        forContextInternal(Display.DEFAULT_DISPLAY, Looper.getMainLooper());
     }
 
     private static final Object sLock = new Object();
@@ -279,6 +280,17 @@ public final class InputMethodManager {
     static InputMethodManager sInstance;
 
     /**
+     * Global map between display to {@link InputMethodManager}.
+     *
+     * <p>Currently this map works like a so-called leaky singleton.  Once an instance is registered
+     * for the associated display ID, that instance will never be garbage collected.</p>
+     *
+     * <p>TODO(Bug 116699479): Implement instance clean up mechanism.</p>
+     */
+    @GuardedBy("sLock")
+    private static final SparseArray<InputMethodManager> sInstanceMap = new SparseArray<>();
+
+    /**
      * @hide Flag for IInputMethodManager.windowGainedFocus: a view in
      * the window has input focus.
      */
@@ -335,6 +347,8 @@ public final class InputMethodManager {
     // Our generic input connection if the current target does not have its own.
     final IInputContext mIInputContext;
 
+    private final int mDisplayId;
+
     /**
      * True if this input method client is active, initially false.
      */
@@ -452,6 +466,29 @@ public final class InputMethodManager {
         return afm != null && afm.isAutofillUiShowing();
     }
 
+    /**
+     * Checks the consistency between {@link InputMethodManager} state and {@link View} state.
+     *
+     * @param view {@link View} to be checked
+     * @return {@code true} if {@code view} is not {@code null} and there is a {@link Context}
+     *         mismatch between {@link InputMethodManager} and {@code view}
+     */
+    private boolean shouldDispatchToViewContext(@Nullable View view) {
+        if (view == null) {
+            return false;
+        }
+        final int viewDisplayId = getDisplayId(view.getContext());
+        if (viewDisplayId != mDisplayId) {
+            Log.w(TAG, "b/117267690: Context mismatch found. view=" + view + " belongs to"
+                    + " displayId=" + viewDisplayId
+                    + " but InputMethodManager belongs to displayId=" + mDisplayId
+                    + ". Use the right InputMethodManager instance to avoid performance overhead.",
+                    new Throwable());
+            return true;
+        }
+        return false;
+    }
+
     private static boolean canStartInput(View servedView) {
         // We can start input ether the servedView has window focus
         // or the activity is showing autofill ui.
@@ -733,33 +770,57 @@ public final class InputMethodManager {
                 });
     }
 
-    InputMethodManager(Looper looper) throws ServiceNotFoundException {
+    InputMethodManager(int displayId, Looper looper) throws ServiceNotFoundException {
         mService = getIInputMethodManager();
         mMainLooper = looper;
         mH = new H(looper);
+        mDisplayId = displayId;
         mIInputContext = new ControlledInputConnectionWrapper(looper,
                 mDummyInputConnection, this);
     }
 
+    private static int getDisplayId(Context context) {
+        final Display display = context.getDisplay();
+        return display != null ? display.getDisplayId() : Display.DEFAULT_DISPLAY;
+    }
+
     /**
-     * Retrieve the global {@link InputMethodManager} instance, creating it if it doesn't already
-     * exist.
+     * Retrieve an instance for the given {@link Context}, creating it if it doesn't already exist.
      *
-     * @return global {@link InputMethodManager} instance
+     * @param context {@link Context} for which IME APIs need to work
+     * @return {@link InputMethodManager} instance
      * @hide
      */
-    public static InputMethodManager getInstanceInternal() {
+    @Nullable
+    public static InputMethodManager forContext(Context context) {
+        final int displayId = getDisplayId(context);
+        // For better backward compatibility, we always use Looper.getMainLooper() for the default
+        // display case.
+        final Looper looper = displayId == Display.DEFAULT_DISPLAY
+                ? Looper.getMainLooper() : context.getMainLooper();
+        return forContextInternal(displayId, looper);
+    }
+
+    @Nullable
+    private static InputMethodManager forContextInternal(int displayId, Looper looper) {
+        final boolean isDefaultDisplay = displayId == Display.DEFAULT_DISPLAY;
         synchronized (sLock) {
-            if (sInstance == null) {
-                try {
-                    final InputMethodManager imm = new InputMethodManager(Looper.getMainLooper());
-                    imm.mService.addClient(imm.mClient, imm.mIInputContext);
-                    sInstance = imm;
-                } catch (ServiceNotFoundException | RemoteException e) {
-                    throw new IllegalStateException(e);
-                }
+            InputMethodManager instance = sInstanceMap.get(displayId);
+            if (instance != null) {
+                return instance;
             }
-            return sInstance;
+            try {
+                instance = new InputMethodManager(displayId, looper);
+                instance.mService.addClient(instance.mClient, instance.mIInputContext, displayId);
+            } catch (ServiceNotFoundException | RemoteException e) {
+                throw new IllegalStateException(e);
+            }
+            // For backward compatibility, store the instance also to sInstance for default display.
+            if (sInstance == null && isDefaultDisplay) {
+                sInstance = instance;
+            }
+            sInstanceMap.put(displayId, instance);
+            return instance;
         }
     }
 
@@ -916,6 +977,11 @@ public final class InputMethodManager {
      * input method.
      */
     public boolean isActive(View view) {
+        // Re-dispatch if there is a context mismatch.
+        if (shouldDispatchToViewContext(view)) {
+            return view.getContext().getSystemService(InputMethodManager.class).isActive(view);
+        }
+
         checkFocus();
         synchronized (mH) {
             return (mServedView == view
@@ -1006,6 +1072,13 @@ public final class InputMethodManager {
     }
 
     public void displayCompletions(View view, CompletionInfo[] completions) {
+        // Re-dispatch if there is a context mismatch.
+        if (shouldDispatchToViewContext(view)) {
+            view.getContext().getSystemService(InputMethodManager.class)
+                    .displayCompletions(view, completions);
+            return;
+        }
+
         checkFocus();
         synchronized (mH) {
             if (mServedView != view && (mServedView == null
@@ -1024,6 +1097,13 @@ public final class InputMethodManager {
     }
 
     public void updateExtractedText(View view, int token, ExtractedText text) {
+        // Re-dispatch if there is a context mismatch.
+        if (shouldDispatchToViewContext(view)) {
+            view.getContext().getSystemService(InputMethodManager.class)
+                    .updateExtractedText(view, token, text);
+            return;
+        }
+
         checkFocus();
         synchronized (mH) {
             if (mServedView != view && (mServedView == null
@@ -1065,6 +1145,12 @@ public final class InputMethodManager {
      * 0 or have the {@link #SHOW_IMPLICIT} bit set.
      */
     public boolean showSoftInput(View view, int flags) {
+        // Re-dispatch if there is a context mismatch.
+        if (shouldDispatchToViewContext(view)) {
+            return view.getContext().getSystemService(InputMethodManager.class)
+                    .showSoftInput(view, flags);
+        }
+
         return showSoftInput(view, flags, null);
     }
 
@@ -1127,6 +1213,12 @@ public final class InputMethodManager {
      * {@link #RESULT_HIDDEN}.
      */
     public boolean showSoftInput(View view, int flags, ResultReceiver resultReceiver) {
+        // Re-dispatch if there is a context mismatch.
+        if (shouldDispatchToViewContext(view)) {
+            return view.getContext().getSystemService(InputMethodManager.class)
+                    .showSoftInput(view, flags, resultReceiver);
+        }
+
         checkFocus();
         synchronized (mH) {
             if (mServedView != view && (mServedView == null
@@ -1290,6 +1382,12 @@ public final class InputMethodManager {
      * @param view The view whose text has changed.
      */
     public void restartInput(View view) {
+        // Re-dispatch if there is a context mismatch.
+        if (shouldDispatchToViewContext(view)) {
+            view.getContext().getSystemService(InputMethodManager.class).restartInput(view);
+            return;
+        }
+
         checkFocus();
         synchronized (mH) {
             if (mServedView != view && (mServedView == null
@@ -1714,6 +1812,13 @@ public final class InputMethodManager {
      */
     public void updateSelection(View view, int selStart, int selEnd,
             int candidatesStart, int candidatesEnd) {
+        // Re-dispatch if there is a context mismatch.
+        if (shouldDispatchToViewContext(view)) {
+            view.getContext().getSystemService(InputMethodManager.class)
+                    .updateSelection(view, selStart, selEnd, candidatesStart, candidatesEnd);
+            return;
+        }
+
         checkFocus();
         synchronized (mH) {
             if ((mServedView != view && (mServedView == null
@@ -1751,6 +1856,12 @@ public final class InputMethodManager {
      * Notify the event when the user tapped or clicked the text view.
      */
     public void viewClicked(View view) {
+        // Re-dispatch if there is a context mismatch.
+        if (shouldDispatchToViewContext(view)) {
+            view.getContext().getSystemService(InputMethodManager.class).viewClicked(view);
+            return;
+        }
+
         final boolean focusChanged = mServedView != mNextServedView;
         checkFocus();
         synchronized (mH) {
@@ -1815,6 +1926,13 @@ public final class InputMethodManager {
      */
     @Deprecated
     public void updateCursor(View view, int left, int top, int right, int bottom) {
+        // Re-dispatch if there is a context mismatch.
+        if (shouldDispatchToViewContext(view)) {
+            view.getContext().getSystemService(InputMethodManager.class)
+                    .updateCursor(view, left, top, right, bottom);
+            return;
+        }
+
         checkFocus();
         synchronized (mH) {
             if ((mServedView != view && (mServedView == null
@@ -1846,6 +1964,13 @@ public final class InputMethodManager {
         if (view == null || cursorAnchorInfo == null) {
             return;
         }
+        // Re-dispatch if there is a context mismatch.
+        if (shouldDispatchToViewContext(view)) {
+            view.getContext().getSystemService(InputMethodManager.class)
+                    .updateCursorAnchorInfo(view, cursorAnchorInfo);
+            return;
+        }
+
         checkFocus();
         synchronized (mH) {
             if ((mServedView != view &&
@@ -1891,6 +2016,13 @@ public final class InputMethodManager {
      * @param data Any data to include with the command.
      */
     public void sendAppPrivateCommand(View view, String action, Bundle data) {
+        // Re-dispatch if there is a context mismatch.
+        if (shouldDispatchToViewContext(view)) {
+            view.getContext().getSystemService(InputMethodManager.class)
+                    .sendAppPrivateCommand(view, action, data);
+            return;
+        }
+
         checkFocus();
         synchronized (mH) {
             if ((mServedView != view && (mServedView == null
@@ -2062,6 +2194,13 @@ public final class InputMethodManager {
      */
     public void dispatchKeyEventFromInputMethod(@Nullable View targetView,
             @NonNull KeyEvent event) {
+        // Re-dispatch if there is a context mismatch.
+        if (shouldDispatchToViewContext(targetView)) {
+            targetView.getContext().getSystemService(InputMethodManager.class)
+                    .dispatchKeyEventFromInputMethod(targetView, event);
+            return;
+        }
+
         synchronized (mH) {
             ViewRootImpl viewRootImpl = targetView != null ? targetView.getViewRootImpl() : null;
             if (viewRootImpl == null) {
@@ -2551,6 +2690,7 @@ public final class InputMethodManager {
         sb.append(",windowFocus=" + view.hasWindowFocus());
         sb.append(",autofillUiShowing=" + isAutofillUIShowing(view));
         sb.append(",window=" + view.getWindowToken());
+        sb.append(",displayId=" + getDisplayId(view.getContext()));
         sb.append(",temporaryDetach=" + view.isTemporarilyDetached());
         return sb.toString();
     }
diff --git a/core/java/com/android/internal/view/IInputMethodManager.aidl b/core/java/com/android/internal/view/IInputMethodManager.aidl
index 5f1243f37542..dceacda5d4a3 100644
--- a/core/java/com/android/internal/view/IInputMethodManager.aidl
+++ b/core/java/com/android/internal/view/IInputMethodManager.aidl
@@ -31,7 +31,8 @@ import com.android.internal.view.IInputMethodClient;
  * applications.
  */
 interface IInputMethodManager {
-    void addClient(in IInputMethodClient client, in IInputContext inputContext);
+    void addClient(in IInputMethodClient client, in IInputContext inputContext,
+            int untrustedDisplayId);
 
     // TODO: Use ParceledListSlice instead
     List<InputMethodInfo> getInputMethodList();
diff --git a/core/java/com/android/internal/view/InputBindResult.java b/core/java/com/android/internal/view/InputBindResult.java
index 101fd41f2925..ec8e8dacb9db 100644
--- a/core/java/com/android/internal/view/InputBindResult.java
+++ b/core/java/com/android/internal/view/InputBindResult.java
@@ -51,6 +51,9 @@ public final class InputBindResult implements Parcelable {
             ResultCode.ERROR_INVALID_USER,
             ResultCode.ERROR_NULL_EDITOR_INFO,
             ResultCode.ERROR_NOT_IME_TARGET_WINDOW,
+            ResultCode.ERROR_NO_EDITOR,
+            ResultCode.ERROR_DISPLAY_ID_MISMATCH,
+            ResultCode.ERROR_INVALID_DISPLAY_ID,
     })
     public @interface ResultCode {
         /**
@@ -139,13 +142,22 @@ public final class InputBindResult implements Parcelable {
          * The client should try to restart input when its {@link android.view.Window} is focused
          * again.</p>
          *
-         * @see com.android.server.wm.WindowManagerInternal#isInputMethodClientFocus(int, int)
+         * @see com.android.server.wm.WindowManagerInternal#isInputMethodClientFocus(int, int, int)
          */
         int ERROR_NOT_IME_TARGET_WINDOW = 11;
         /**
          * Indicates that focused view in the current window is not an editor.
          */
         int ERROR_NO_EDITOR = 12;
+        /**
+         * Indicates that there is a mismatch in display ID between IME client and focused Window.
+         */
+        int ERROR_DISPLAY_ID_MISMATCH = 13;
+        /**
+         * Indicates that current IME client is no longer allowed to access to the associated
+         * display.
+         */
+        int ERROR_INVALID_DISPLAY_ID = 14;
     }
 
     @ResultCode
@@ -271,6 +283,10 @@ public final class InputBindResult implements Parcelable {
                 return "ERROR_NULL_EDITOR_INFO";
             case ResultCode.ERROR_NOT_IME_TARGET_WINDOW:
                 return "ERROR_NOT_IME_TARGET_WINDOW";
+            case ResultCode.ERROR_DISPLAY_ID_MISMATCH:
+                return "ERROR_DISPLAY_ID_MISMATCH";
+            case ResultCode.ERROR_INVALID_DISPLAY_ID:
+                return "ERROR_INVALID_DISPLAY_ID";
             default:
                 return "Unknown(" + result + ")";
         }
@@ -316,4 +332,15 @@ public final class InputBindResult implements Parcelable {
      */
     public static final InputBindResult INVALID_USER = error(ResultCode.ERROR_INVALID_USER);
 
+    /**
+     * Predefined error object for {@link ResultCode#ERROR_DISPLAY_ID_MISMATCH}.
+     */
+    public static final InputBindResult DISPLAY_ID_MISMATCH =
+            error(ResultCode.ERROR_DISPLAY_ID_MISMATCH);
+
+    /**
+     * Predefined error object for {@link ResultCode#ERROR_INVALID_DISPLAY_ID}.
+     */
+    public static final InputBindResult INVALID_DISPLAY_ID =
+            error(ResultCode.ERROR_INVALID_DISPLAY_ID);
 }
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index a9b0d5c42f73..71c419f50790 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -421,6 +421,7 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
         final IInputContext inputContext;
         final int uid;
         final int pid;
+        final int selfReportedDisplayId;
         final InputBinding binding;
         final ClientDeathRecipient clientDeathRecipient;
 
@@ -430,16 +431,18 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
         @Override
         public String toString() {
             return "ClientState{" + Integer.toHexString(
-                    System.identityHashCode(this)) + " uid " + uid
-                    + " pid " + pid + "}";
+                    System.identityHashCode(this)) + " uid=" + uid
+                    + " pid=" + pid + " displayId=" + selfReportedDisplayId + "}";
         }
 
         ClientState(IInputMethodClient _client, IInputContext _inputContext,
-                int _uid, int _pid, ClientDeathRecipient _clientDeathRecipient) {
+                int _uid, int _pid, int _selfReportedDisplayId,
+                ClientDeathRecipient _clientDeathRecipient) {
             client = _client;
             inputContext = _inputContext;
             uid = _uid;
             pid = _pid;
+            selfReportedDisplayId = _selfReportedDisplayId;
             binding = new InputBinding(null, inputContext.asBinder(), uid, pid);
             clientDeathRecipient = _clientDeathRecipient;
         }
@@ -1745,15 +1748,21 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
      *               process
      * @param inputContext communication channel for the dummy
      *                     {@link android.view.inputmethod.InputConnection}
+     * @param selfReportedDisplayId self-reported display ID to which the client is associated.
+     *                              Whether the client is still allowed to access to this display
+     *                              or not needs to be evaluated every time the client interacts
+     *                              with the display
      */
     @Override
-    public void addClient(IInputMethodClient client, IInputContext inputContext) {
+    public void addClient(IInputMethodClient client, IInputContext inputContext,
+            int selfReportedDisplayId) {
         final int callerUid = Binder.getCallingUid();
         final int callerPid = Binder.getCallingPid();
         synchronized (mMethodMap) {
             // TODO: Optimize this linear search.
             for (ClientState state : mClients.values()) {
-                if (state.uid == callerUid && state.pid == callerPid) {
+                if (state.uid == callerUid && state.pid == callerPid
+                        && state.selfReportedDisplayId == selfReportedDisplayId) {
                     throw new SecurityException("uid=" + callerUid + "/pid=" + callerPid
                             + " is already registered");
                 }
@@ -1764,11 +1773,25 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
             } catch (RemoteException e) {
                 throw new IllegalStateException(e);
             }
-            mClients.put(client.asBinder(),
-                    new ClientState(client, inputContext, callerUid, callerPid, deathRecipient));
+            // We cannot fully avoid race conditions where the client UID already lost the access to
+            // the given self-reported display ID, even if the client is not maliciously reporting
+            // a fake display ID. Unconditionally returning SecurityException just because the
+            // client doesn't pass display ID verification can cause many test failures hence not an
+            // option right now.  At the same time
+            //    context.getSystemService(InputMethodManager.class)
+            // is expected to return a valid non-null instance at any time if we do not choose to
+            // have the client crash.  Thus we do not verify the display ID at all here.  Instead we
+            // later check the display ID every time the client needs to interact with the specified
+            // display.
+            mClients.put(client.asBinder(), new ClientState(client, inputContext, callerUid,
+                    callerPid, selfReportedDisplayId, deathRecipient));
         }
     }
 
+    private boolean verifyDisplayId(ClientState cs) {
+        return mWindowManagerInternal.isUidAllowedOnDisplay(cs.selfReportedDisplayId, cs.uid);
+    }
+
     void removeClient(IInputMethodClient client) {
         synchronized (mMethodMap) {
             ClientState cs = mClients.remove(client.asBinder());
@@ -1918,8 +1941,9 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
         mCurAttribute = attribute;
 
         // Check if the input method is changing.
-        final int displayId = mWindowManagerInternal.getDisplayIdForWindow(
-                mCurFocusedWindow);
+        // We expect the caller has already verified that the client is allowed to access this
+        // display ID.
+        final int displayId = mCurFocusedWindowClient.selfReportedDisplayId;
         if (mCurId != null && mCurId.equals(mCurMethodId) && displayId == mCurTokenDisplayId) {
             if (cs.curSession != null) {
                 // Fast case: if we are already connected to the input method,
@@ -1984,7 +2008,11 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
                 com.android.internal.R.string.input_method_binding_label);
         mCurIntent.putExtra(Intent.EXTRA_CLIENT_INTENT, PendingIntent.getActivity(
                 mContext, 0, new Intent(Settings.ACTION_INPUT_METHOD_SETTINGS), 0));
-        final int displayId = mWindowManagerInternal.getDisplayIdForWindow(mCurFocusedWindow);
+        if (!verifyDisplayId(mCurFocusedWindowClient)) {
+            // Wait, the client no longer has access to the display.
+            return InputBindResult.INVALID_DISPLAY_ID;
+        }
+        final int displayId = mCurFocusedWindowClient.selfReportedDisplayId;
         mCurTokenDisplayId = (displayId != INVALID_DISPLAY) ? displayId : DEFAULT_DISPLAY;
 
         if (bindCurrentInputMethodServiceLocked(mCurIntent, this, IME_CONNECTION_BIND_FLAGS)) {
@@ -2584,7 +2612,8 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
                     if (cs == null) {
                         throw new IllegalArgumentException("unknown client " + client.asBinder());
                     }
-                    if (!mWindowManagerInternal.isInputMethodClientFocus(cs.uid, cs.pid)) {
+                    if (!mWindowManagerInternal.isInputMethodClientFocus(cs.uid, cs.pid,
+                            cs.selfReportedDisplayId)) {
                         Slog.w(TAG, "Ignoring showSoftInput of uid " + uid + ": " + client);
                         return false;
                     }
@@ -2668,7 +2697,8 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
                     if (cs == null) {
                         throw new IllegalArgumentException("unknown client " + client.asBinder());
                     }
-                    if (!mWindowManagerInternal.isInputMethodClientFocus(cs.uid, cs.pid)) {
+                    if (!mWindowManagerInternal.isInputMethodClientFocus(cs.uid, cs.pid,
+                            cs.selfReportedDisplayId)) {
                         if (DEBUG) {
                             Slog.w(TAG, "Ignoring hideSoftInput of uid " + uid + ": " + client);
                         }
@@ -2767,6 +2797,8 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
         InputBindResult res = null;
         long ident = Binder.clearCallingIdentity();
         try {
+            final int windowDisplayId =
+                    mWindowManagerInternal.getDisplayIdForWindow(windowToken);
             synchronized (mMethodMap) {
                 if (DEBUG) Slog.v(TAG, "startInputOrWindowGainedFocusInternal: reason="
                         + InputMethodClient.getStartInputReason(startInputReason)
@@ -2785,8 +2817,15 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
                     throw new IllegalArgumentException("unknown client "
                             + client.asBinder());
                 }
+                if (cs.selfReportedDisplayId != windowDisplayId) {
+                    Slog.e(TAG, "startInputOrWindowGainedFocusInternal: display ID mismatch."
+                            + " from client:" + cs.selfReportedDisplayId
+                            + " from window:" + windowDisplayId);
+                    return InputBindResult.DISPLAY_ID_MISMATCH;
+                }
 
-                if (!mWindowManagerInternal.isInputMethodClientFocus(cs.uid, cs.pid)) {
+                if (!mWindowManagerInternal.isInputMethodClientFocus(cs.uid, cs.pid,
+                        cs.selfReportedDisplayId)) {
                     // Check with the window manager to make sure this client actually
                     // has a window with focus.  If not, reject.  This is thread safe
                     // because if the focus changes some time before or after, the
@@ -2858,9 +2897,9 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
                                 // If focused display changed, we should unbind current method
                                 // to make app window in previous display relayout after Ime
                                 // window token removed.
-                                final int newFocusDisplayId =
-                                        mWindowManagerInternal.getDisplayIdForWindow(windowToken);
-                                if (newFocusDisplayId != mCurTokenDisplayId) {
+                                // Note that we can trust client's display ID as long as it matches
+                                // to the display ID obtained from the window.
+                                if (cs.selfReportedDisplayId != mCurTokenDisplayId) {
                                     unbindCurrentMethodLocked();
                                 }
                             }
@@ -2958,6 +2997,7 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
     }
 
     private boolean canShowInputMethodPickerLocked(IInputMethodClient client) {
+        // TODO(yukawa): multi-display support.
         final int uid = Binder.getCallingUid();
         if (UserHandle.getAppId(uid) == Process.SYSTEM_UID) {
             return true;
@@ -3034,6 +3074,7 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
     @Override
     public void showInputMethodAndSubtypeEnablerFromClient(
             IInputMethodClient client, String inputMethodId) {
+        // TODO(yukawa): Should we verify the display ID?
         if (!calledFromValidUser()) {
             return;
         }
@@ -3229,6 +3270,7 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
      */
     @Override
     public int getInputMethodWindowVisibleHeight() {
+        // TODO(yukawa): Should we verify the display ID?
         return mWindowManagerInternal.getInputMethodWindowVisibleHeight(mCurTokenDisplayId);
     }
 
diff --git a/services/core/java/com/android/server/wm/WindowManagerInternal.java b/services/core/java/com/android/server/wm/WindowManagerInternal.java
index 5410676a49fc..b096bf2d0738 100644
--- a/services/core/java/com/android/server/wm/WindowManagerInternal.java
+++ b/services/core/java/com/android/server/wm/WindowManagerInternal.java
@@ -430,14 +430,25 @@ public abstract class WindowManagerInternal {
     public abstract boolean isUidFocused(int uid);
 
     /**
-     * Checks whether the specified process has IME focus or not.
+     * Checks whether the specified IME client has IME focus or not.
      *
      * @param uid UID of the process to be queried
      * @param pid PID of the process to be queried
-     * @return {@code true} if a process that is identified by {@code uid} and {@code pid} has IME
-     *         focus
+     * @param displayId Display ID reported from the client. Note that this method also verifies
+     *                  whether the specified process is allowed to access to this display or not
+     * @return {@code true} if the IME client specified with {@code uid}, {@code pid}, and
+     *         {@code displayId} has IME focus
      */
-    public abstract boolean isInputMethodClientFocus(int uid, int pid);
+    public abstract boolean isInputMethodClientFocus(int uid, int pid, int displayId);
+
+    /**
+     * Checks whether the given {@code uid} is allowed to use the given {@code displayId} or not.
+     *
+     * @param displayId Display ID to be checked
+     * @param uid UID to be checked.
+     * @return {@code true} if the given {@code uid} is allowed to use the given {@code displayId}
+     */
+    public abstract boolean isUidAllowedOnDisplay(int displayId, int uid);
 
     /**
      * Return the display Id for given window.
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 10ba63e0a9f7..e153a1d6594e 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -7201,16 +7201,20 @@ public class WindowManagerService extends IWindowManager.Stub
         }
 
         @Override
-        public boolean isInputMethodClientFocus(int uid, int pid) {
+        public boolean isInputMethodClientFocus(int uid, int pid, int displayId) {
+            if (displayId == Display.INVALID_DISPLAY) {
+                return false;
+            }
             synchronized (mWindowMap) {
-                // Check all displays if any input method window has focus.
-                for (int i = mRoot.mChildren.size() - 1; i >= 0; --i) {
-                    final DisplayContent displayContent = mRoot.mChildren.get(i);
-                    if (displayContent.isInputMethodClientFocus(uid, pid)) {
-                        return true;
-                    }
+                final DisplayContent displayContent = mRoot.getTopFocusedDisplayContent();
+                if (displayContent == null
+                        || displayContent.getDisplayId() != displayId
+                        || !displayContent.hasAccess(uid)) {
+                    return false;
+                }
+                if (displayContent.isInputMethodClientFocus(uid, pid)) {
+                    return true;
                 }
-
                 // Okay, how about this...  what is the current focus?
                 // It seems in some cases we may not have moved the IM
                 // target window, such as when it was in a pop-up window,
@@ -7219,7 +7223,7 @@ public class WindowManagerService extends IWindowManager.Stub
                 // press home.  Sometimes the IME won't go down.)
                 // Would be nice to fix this more correctly, but it's
                 // way at the end of a release, and this should be good enough.
-                final WindowState currentFocus = mRoot.getTopFocusedDisplayContent().mCurrentFocus;
+                final WindowState currentFocus = displayContent.mCurrentFocus;
                 if (currentFocus != null && currentFocus.mSession.mUid == uid
                         && currentFocus.mSession.mPid == pid) {
                     return true;
@@ -7229,6 +7233,20 @@ public class WindowManagerService extends IWindowManager.Stub
         }
 
         @Override
+        public boolean isUidAllowedOnDisplay(int displayId, int uid) {
+            if (displayId == Display.DEFAULT_DISPLAY) {
+                return true;
+            }
+            if (displayId == Display.INVALID_DISPLAY) {
+                return false;
+            }
+            synchronized (mWindowMap) {
+                final DisplayContent displayContent = mRoot.getDisplayContent(displayId);
+                return displayContent != null && displayContent.hasAccess(uid);
+            }
+        }
+
+        @Override
         public int getDisplayIdForWindow(IBinder windowToken) {
             synchronized (mWindowMap) {
                 final WindowState window = mWindowMap.get(windowToken);