Introduce IMMS ClientController component

ClientController is the component responsible for storing and managing
IMMS clients.

This is a pure refactoring, no feature flag is required.

Bug: 314150112
Bug: 315227580
Test: atest CtsInputMethodTestCases CtsInputMethodInstallTestCases
Test: atest FrameworksInputMethodSystemServerTests
Test: atest --host FrameworksInputMethodSystemServerTests_host
Change-Id: I117396b145f098f7b852328a2f1dd3df7a9f3b28
Signed-off-by: Antonio Kantek <kanant@google.com>

5 files changed, 340 insertions, 125 deletions

diff --git a/services/core/java/com/android/server/inputmethod/ClientController.java b/services/core/java/com/android/server/inputmethod/ClientController.java
new file mode 100644
index 000000000000..293464054fdc
--- /dev/null
+++ b/services/core/java/com/android/server/inputmethod/ClientController.java
@@ -0,0 +1,162 @@
+/*
+ * 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.server.inputmethod;
+
+import android.annotation.NonNull;
+import android.content.pm.PackageManagerInternal;
+import android.os.IBinder;
+import android.os.RemoteException;
+import android.util.ArrayMap;
+import android.util.SparseArray;
+import android.view.inputmethod.InputBinding;
+
+import com.android.internal.annotations.GuardedBy;
+import com.android.internal.inputmethod.IInputMethodClient;
+import com.android.internal.inputmethod.IRemoteInputConnection;
+
+/**
+ * Store and manage {@link InputMethodManagerService} clients. This class was designed to be a
+ * singleton in {@link InputMethodManagerService} since it stores information about all clients,
+ * still the current client will be defined per display.
+ *
+ * <p>
+ * As part of the re-architecture plan (described in go/imms-rearchitecture-plan), the following
+ * fields and methods will be moved out from IMMS and placed here:
+ * <ul>
+ * <li>mCurClient (ClientState)</li>
+ * <li>mClients (ArrayMap of ClientState indexed by IBinder)</li>
+ * <li>mLastSwitchUserId</li>
+ * </ul>
+ * <p>
+ * Nested Classes (to move from IMMS):
+ * <ul>
+ * <li>ClientDeathRecipient</li>
+ * <li>ClientState<</li>
+ * </ul>
+ * <p>
+ * Methods to rewrite and/or extract from IMMS and move here:
+ * <ul>
+ * <li>addClient</li>
+ * <li>removeClient</li>
+ * <li>verifyClientAndPackageMatch</li>
+ * <li>setImeTraceEnabledForAllClients (make it reactive)</li>
+ * <li>unbindCurrentClient</li>
+ * </ul>
+ */
+// TODO(b/314150112): Update the Javadoc above, by removing the re-architecture steps, once this
+//  class is finalized
+final class ClientController {
+
+    // TODO(b/314150112): Make this field private when breaking the cycle with IMMS.
+    @GuardedBy("ImfLock.class")
+    final ArrayMap<IBinder, ClientState> mClients = new ArrayMap<>();
+
+    private final PackageManagerInternal mPackageManagerInternal;
+
+    ClientController(PackageManagerInternal packageManagerInternal) {
+        mPackageManagerInternal = packageManagerInternal;
+    }
+
+    @GuardedBy("ImfLock.class")
+    void addClient(IInputMethodClientInvoker clientInvoker,
+            IRemoteInputConnection inputConnection,
+            int selfReportedDisplayId, IBinder.DeathRecipient deathRecipient, int callerUid,
+            int callerPid) {
+        // TODO: Optimize this linear search.
+        final int numClients = mClients.size();
+        for (int i = 0; i < numClients; ++i) {
+            final ClientState state = mClients.valueAt(i);
+            if (state.mUid == callerUid && state.mPid == callerPid
+                    && state.mSelfReportedDisplayId == selfReportedDisplayId) {
+                throw new SecurityException("uid=" + callerUid + "/pid=" + callerPid
+                        + "/displayId=" + selfReportedDisplayId + " is already registered");
+            }
+        }
+        try {
+            clientInvoker.asBinder().linkToDeath(deathRecipient, 0 /* flags */);
+        } catch (RemoteException e) {
+            throw new IllegalStateException(e);
+        }
+        // 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(clientInvoker.asBinder(), new ClientState(clientInvoker, inputConnection,
+                callerUid, callerPid, selfReportedDisplayId, deathRecipient));
+    }
+
+    @GuardedBy("ImfLock.class")
+    boolean verifyClientAndPackageMatch(
+            @NonNull IInputMethodClient client, @NonNull String packageName) {
+        ClientState cs = mClients.get(client.asBinder());
+        if (cs == null) {
+            throw new IllegalArgumentException("unknown client " + client.asBinder());
+        }
+        return InputMethodUtils.checkIfPackageBelongsToUid(
+                mPackageManagerInternal, cs.mUid, packageName);
+    }
+
+    static final class ClientState {
+        final IInputMethodClientInvoker mClient;
+        final IRemoteInputConnection mFallbackInputConnection;
+        final int mUid;
+        final int mPid;
+        final int mSelfReportedDisplayId;
+        final InputBinding mBinding;
+        final IBinder.DeathRecipient mClientDeathRecipient;
+
+        @GuardedBy("ImfLock.class")
+        boolean mSessionRequested;
+
+        @GuardedBy("ImfLock.class")
+        boolean mSessionRequestedForAccessibility;
+
+        @GuardedBy("ImfLock.class")
+        InputMethodManagerService.SessionState mCurSession;
+
+        @GuardedBy("ImfLock.class")
+        SparseArray<InputMethodManagerService.AccessibilitySessionState> mAccessibilitySessions =
+                new SparseArray<>();
+
+        @Override
+        public String toString() {
+            return "ClientState{" + Integer.toHexString(
+                    System.identityHashCode(this)) + " mUid=" + mUid
+                    + " mPid=" + mPid + " mSelfReportedDisplayId=" + mSelfReportedDisplayId + "}";
+        }
+
+        ClientState(IInputMethodClientInvoker client,
+                IRemoteInputConnection fallbackInputConnection,
+                int uid, int pid, int selfReportedDisplayId,
+                IBinder.DeathRecipient clientDeathRecipient) {
+            mClient = client;
+            mFallbackInputConnection = fallbackInputConnection;
+            mUid = uid;
+            mPid = pid;
+            mSelfReportedDisplayId = selfReportedDisplayId;
+            mBinding = new InputBinding(null /*conn*/, mFallbackInputConnection.asBinder(), mUid,
+                    mPid);
+            mClientDeathRecipient = clientDeathRecipient;
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index d722242b20b0..6fe1885ce57b 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -48,6 +48,7 @@ import static android.view.Display.INVALID_DISPLAY;
 import static android.view.WindowManager.DISPLAY_IME_POLICY_HIDE;
 import static android.view.WindowManager.DISPLAY_IME_POLICY_LOCAL;
 
+import static com.android.server.inputmethod.ClientController.ClientState;
 import static com.android.server.inputmethod.ImeVisibilityStateComputer.ImeTargetWindowState;
 import static com.android.server.inputmethod.ImeVisibilityStateComputer.ImeVisibilityResult;
 import static com.android.server.inputmethod.ImeVisibilityStateComputer.STATE_HIDE_IME;
@@ -127,7 +128,6 @@ import android.view.WindowManager.LayoutParams.SoftInputModeFlags;
 import android.view.inputmethod.EditorInfo;
 import android.view.inputmethod.Flags;
 import android.view.inputmethod.ImeTracker;
-import android.view.inputmethod.InputBinding;
 import android.view.inputmethod.InputConnection;
 import android.view.inputmethod.InputMethod;
 import android.view.inputmethod.InputMethodEditorTraceProto.InputMethodClientsTraceFileProto;
@@ -273,6 +273,7 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
     @NonNull
     private final String[] mNonPreemptibleInputMethods;
 
+    // TODO(b/314150112): Move this to ClientController.
     @UserIdInt
     private int mLastSwitchUserId;
 
@@ -391,7 +392,7 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
     /**
      * Record session state for an accessibility service.
      */
-    private static class AccessibilitySessionState {
+    static class AccessibilitySessionState {
         final ClientState mClient;
         // Id of the accessibility service.
         final int mId;
@@ -415,58 +416,10 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
         }
     }
 
-    private static final class ClientDeathRecipient implements IBinder.DeathRecipient {
-        private final InputMethodManagerService mImms;
-        private final IInputMethodClient mClient;
-
-        ClientDeathRecipient(InputMethodManagerService imms, IInputMethodClient client) {
-            mImms = imms;
-            mClient = client;
-        }
-
-        @Override
-        public void binderDied() {
-            mImms.removeClient(mClient);
-        }
-    }
-
-    static final class ClientState {
-        final IInputMethodClientInvoker mClient;
-        final IRemoteInputConnection mFallbackInputConnection;
-        final int mUid;
-        final int mPid;
-        final int mSelfReportedDisplayId;
-        final InputBinding mBinding;
-        final ClientDeathRecipient mClientDeathRecipient;
-
-        boolean mSessionRequested;
-        boolean mSessionRequestedForAccessibility;
-        SessionState mCurSession;
-        SparseArray<AccessibilitySessionState> mAccessibilitySessions = new SparseArray<>();
-
-        @Override
-        public String toString() {
-            return "ClientState{" + Integer.toHexString(
-                    System.identityHashCode(this)) + " mUid=" + mUid
-                    + " mPid=" + mPid + " mSelfReportedDisplayId=" + mSelfReportedDisplayId + "}";
-        }
-
-        ClientState(IInputMethodClientInvoker client,
-                IRemoteInputConnection fallbackInputConnection,
-                int uid, int pid, int selfReportedDisplayId,
-                ClientDeathRecipient clientDeathRecipient) {
-            mClient = client;
-            mFallbackInputConnection = fallbackInputConnection;
-            mUid = uid;
-            mPid = pid;
-            mSelfReportedDisplayId = selfReportedDisplayId;
-            mBinding = new InputBinding(null, mFallbackInputConnection.asBinder(), mUid, mPid);
-            mClientDeathRecipient = clientDeathRecipient;
-        }
-    }
-
-    @GuardedBy("ImfLock.class")
-    final ArrayMap<IBinder, ClientState> mClients = new ArrayMap<>();
+    /**
+     * Manages the IME clients.
+     */
+    private final ClientController mClientController;
 
     /**
      * Set once the system is ready to run third party code.
@@ -524,6 +477,7 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
     /**
      * The client that is currently bound to an input method.
      */
+    // TODO(b/314150112): Move this to ClientController.
     @Nullable
     private ClientState mCurClient;
 
@@ -864,8 +818,9 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
             @Nullable
             final String mImeSurfaceParentName;
 
-            Entry(ClientState client, EditorInfo editorInfo, String focusedWindowName,
-                    @SoftInputModeFlags int softInputMode, @SoftInputShowHideReason int reason,
+            Entry(ClientState client, EditorInfo editorInfo,
+                    String focusedWindowName, @SoftInputModeFlags int softInputMode,
+                    @SoftInputShowHideReason int reason,
                     boolean inFullscreenMode, String requestWindowName,
                     @Nullable String imeControlTargetName, @Nullable String imeTargetName,
                     @Nullable String imeSurfaceParentName) {
@@ -1719,6 +1674,7 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
 
         mVisibilityStateComputer = new ImeVisibilityStateComputer(this);
         mVisibilityApplier = new DefaultImeVisibilityApplier(this);
+        mClientController = new ClientController(mPackageManagerInternal);
 
         mPreventImeStartupUnlessTextEditor = mRes.getBoolean(
                 com.android.internal.R.bool.config_preventImeStartupUnlessTextEditor);
@@ -1876,7 +1832,8 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
         mLastSwitchUserId = newUserId;
 
         if (mIsInteractive && clientToBeReset != null) {
-            final ClientState cs = mClients.get(clientToBeReset.asBinder());
+            final ClientState cs =
+                    mClientController.mClients.get(clientToBeReset.asBinder());
             if (cs == null) {
                 // The client is already gone.
                 return;
@@ -2214,43 +2171,22 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
         // actually running.
         final int callerUid = Binder.getCallingUid();
         final int callerPid = Binder.getCallingPid();
+
+        // TODO(b/314150112): Move the death recipient logic to ClientController when moving
+        //     removeClient method.
+        final IBinder.DeathRecipient deathRecipient = () -> removeClient(client);
+        final IInputMethodClientInvoker clientInvoker =
+                IInputMethodClientInvoker.create(client, mHandler);
         synchronized (ImfLock.class) {
-            // TODO: Optimize this linear search.
-            final int numClients = mClients.size();
-            for (int i = 0; i < numClients; ++i) {
-                final ClientState state = mClients.valueAt(i);
-                if (state.mUid == callerUid && state.mPid == callerPid
-                        && state.mSelfReportedDisplayId == selfReportedDisplayId) {
-                    throw new SecurityException("uid=" + callerUid + "/pid=" + callerPid
-                            + "/displayId=" + selfReportedDisplayId + " is already registered.");
-                }
-            }
-            final ClientDeathRecipient deathRecipient = new ClientDeathRecipient(this, client);
-            try {
-                client.asBinder().linkToDeath(deathRecipient, 0 /* flags */);
-            } catch (RemoteException e) {
-                throw new IllegalStateException(e);
-            }
-            // 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.
-            final IInputMethodClientInvoker clientInvoker =
-                    IInputMethodClientInvoker.create(client, mHandler);
-            mClients.put(client.asBinder(), new ClientState(clientInvoker, inputConnection,
-                    callerUid, callerPid, selfReportedDisplayId, deathRecipient));
+            mClientController.addClient(clientInvoker, inputConnection, selfReportedDisplayId,
+                    deathRecipient, callerUid, callerPid);
         }
     }
 
+    // TODO(b/314150112): Move this to ClientController.
     void removeClient(IInputMethodClient client) {
         synchronized (ImfLock.class) {
-            ClientState cs = mClients.remove(client.asBinder());
+            ClientState cs = mClientController.mClients.remove(client.asBinder());
             if (cs != null) {
                 client.asBinder().unlinkToDeath(cs.mClientDeathRecipient, 0 /* flags */);
                 clearClientSessionLocked(cs);
@@ -2280,6 +2216,7 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
         }
     }
 
+    // TODO(b/314150112): Move this to ClientController.
     @GuardedBy("ImfLock.class")
     void unbindCurrentClientLocked(@UnbindReason int unbindClientReason) {
         if (mCurClient != null) {
@@ -2332,7 +2269,10 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
         }
     }
 
-    /** {@code true} when a {@link ClientState} has attached from starting the input connection. */
+    /**
+     * {@code true} when a {@link ClientState} has attached from starting the
+     * input connection.
+     */
     @GuardedBy("ImfLock.class")
     boolean hasAttachedClient() {
         return mCurClient != null;
@@ -2976,10 +2916,10 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
     @GuardedBy("ImfLock.class")
     void clearClientSessionsLocked() {
         if (getCurMethodLocked() != null) {
-            final int numClients = mClients.size();
+            final int numClients = mClientController.mClients.size();
             for (int i = 0; i < numClients; ++i) {
-                clearClientSessionLocked(mClients.valueAt(i));
-                clearClientSessionForAccessibilityLocked(mClients.valueAt(i));
+                clearClientSessionLocked(mClientController.mClients.valueAt(i));
+                clearClientSessionForAccessibilityLocked(mClientController.mClients.valueAt(i));
             }
 
             finishSessionLocked(mEnabledSession);
@@ -3509,9 +3449,12 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
                     + " pref is disabled for user: " + userId);
             return;
         }
-        if (!verifyClientAndPackageMatch(client, delegatorPackageName)) {
-            Slog.w(TAG, "prepareStylusHandwritingDelegation() fail");
-            throw new IllegalArgumentException("Delegator doesn't match Uid");
+        synchronized (ImfLock.class) {
+            if (!mClientController.verifyClientAndPackageMatch(client,
+                    delegatorPackageName)) {
+                Slog.w(TAG, "prepareStylusHandwritingDelegation() fail");
+                throw new IllegalArgumentException("Delegator doesn't match Uid");
+            }
         }
         schedulePrepareStylusHandwritingDelegation(
                 userId, delegatePackageName, delegatorPackageName);
@@ -3537,30 +3480,17 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
         return true;
     }
 
-    private boolean verifyClientAndPackageMatch(
-            @NonNull IInputMethodClient client, @NonNull String packageName) {
-        ClientState cs;
-        synchronized (ImfLock.class) {
-            cs = mClients.get(client.asBinder());
-        }
-        if (cs == null) {
-            throw new IllegalArgumentException("unknown client " + client.asBinder());
-        }
-        return InputMethodUtils.checkIfPackageBelongsToUid(
-                mPackageManagerInternal, cs.mUid, packageName);
-    }
-
     private boolean verifyDelegator(
             @NonNull IInputMethodClient client,
             @NonNull String delegatePackageName,
             @NonNull String delegatorPackageName,
             @InputMethodManager.HandwritingDelegateFlags int flags) {
-        if (!verifyClientAndPackageMatch(client, delegatePackageName)) {
-            Slog.w(TAG, "Delegate package does not belong to the same user. Ignoring"
-                    + " startStylusHandwriting");
-            return false;
-        }
         synchronized (ImfLock.class) {
+            if (!mClientController.verifyClientAndPackageMatch(client, delegatePackageName)) {
+                Slog.w(TAG, "Delegate package does not belong to the same user. Ignoring"
+                        + " startStylusHandwriting");
+                return false;
+            }
             boolean homeDelegatorAllowed =
                     (flags & InputMethodManager.HANDWRITING_DELEGATE_FLAG_HOME_DELEGATOR_ALLOWED)
                             != 0;
@@ -3823,7 +3753,7 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
             return InputBindResult.INVALID_USER;
         }
 
-        final ClientState cs = mClients.get(client.asBinder());
+        final ClientState cs = mClientController.mClients.get(client.asBinder());
         if (cs == null) {
             throw new IllegalArgumentException("unknown client " + client.asBinder());
         }
@@ -3997,7 +3927,8 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
             // We need to check if this is the current client with
             // focus in the window manager, to allow this call to
             // be made before input is started in it.
-            final ClientState cs = mClients.get(client.asBinder());
+            final ClientState cs =
+                    mClientController.mClients.get(client.asBinder());
             if (cs == null) {
                 ImeTracker.forLogging().onFailed(statsToken, ImeTracker.PHASE_SERVER_CLIENT_KNOWN);
                 throw new IllegalArgumentException("unknown client " + client.asBinder());
@@ -4621,7 +4552,7 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
         ImeTracing.getInstance().startTrace(null /* printwriter */);
         ArrayMap<IBinder, ClientState> clients;
         synchronized (ImfLock.class) {
-            clients = new ArrayMap<>(mClients);
+            clients = new ArrayMap<>(mClientController.mClients);
         }
         for (ClientState state : clients.values()) {
             if (state != null) {
@@ -4639,7 +4570,7 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
         ImeTracing.getInstance().stopTrace(null /* printwriter */);
         ArrayMap<IBinder, ClientState> clients;
         synchronized (ImfLock.class) {
-            clients = new ArrayMap<>(mClients);
+            clients = new ArrayMap<>(mClientController.mClients);
         }
         for (ClientState state : clients.values()) {
             if (state != null) {
@@ -5878,10 +5809,10 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
                 // We only have sessions when we bound to an input method. Remove this session
                 // from all clients.
                 if (getCurMethodLocked() != null) {
-                    final int numClients = mClients.size();
+                    final int numClients = mClientController.mClients.size();
                     for (int i = 0; i < numClients; ++i) {
-                        clearClientSessionForAccessibilityLocked(mClients.valueAt(i),
-                                accessibilityConnectionId);
+                        clearClientSessionForAccessibilityLocked(
+                                mClientController.mClients.valueAt(i), accessibilityConnectionId);
                     }
                     AccessibilitySessionState session = mEnabledAccessibilitySessions.get(
                             accessibilityConnectionId);
@@ -6066,9 +5997,10 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
                 info.dump(p, "    ");
             }
             p.println("  ClientStates:");
-            final int numClients = mClients.size();
+            // TODO(b/314150112): move client related dump info to ClientController#dump
+            final int numClients = mClientController.mClients.size();
             for (int i = 0; i < numClients; ++i) {
-                final ClientState ci = mClients.valueAt(i);
+                final ClientState ci = mClientController.mClients.valueAt(i);
                 p.println("  " + ci + ":");
                 p.println("    client=" + ci.mClient);
                 p.println("    fallbackInputConnection=" + ci.mFallbackInputConnection);
@@ -6687,7 +6619,7 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
         boolean isImeTraceEnabled = ImeTracing.getInstance().isEnabled();
         ArrayMap<IBinder, ClientState> clients;
         synchronized (ImfLock.class) {
-            clients = new ArrayMap<>(mClients);
+            clients = new ArrayMap<>(mClientController.mClients);
         }
         for (ClientState state : clients.values()) {
             if (state != null) {
diff --git a/services/tests/InputMethodSystemServerTests/Android.bp b/services/tests/InputMethodSystemServerTests/Android.bp
index ffe6dc5d1c63..56423b961813 100644
--- a/services/tests/InputMethodSystemServerTests/Android.bp
+++ b/services/tests/InputMethodSystemServerTests/Android.bp
@@ -40,6 +40,7 @@ android_test {
         "frameworks-base-testutils",
         "mockito-target-extended-minus-junit4",
         "platform-test-annotations",
+        "ravenwood-junit",
         "services.core",
         "service-permission.stubs.system_server",
         "servicestests-core-utils",
@@ -66,6 +67,28 @@ android_test {
     },
 }
 
+android_ravenwood_test {
+    name: "FrameworksInputMethodSystemServerTests_host",
+    static_libs: [
+        "androidx.annotation_annotation",
+        "androidx.test.rules",
+        "framework",
+        "mockito_ravenwood",
+        "ravenwood-runtime",
+        "ravenwood-utils",
+        "services",
+    ],
+    libs: [
+        "android.test.base",
+        "android.test.runner",
+    ],
+    srcs: [
+        "src/com/android/server/inputmethod/**/ClientControllerTest.java",
+    ],
+    sdk_version: "test_current",
+    auto_gen_config: true,
+}
+
 android_test {
     name: "FrameworksImeTests",
     defaults: [
@@ -88,6 +111,7 @@ android_test {
         "frameworks-base-testutils",
         "mockito-target-extended-minus-junit4",
         "platform-test-annotations",
+        "ravenwood-junit",
         "services.core",
         "service-permission.stubs.system_server",
         "servicestests-core-utils",
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/ClientControllerTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/ClientControllerTest.java
new file mode 100644
index 000000000000..3c8f5c9578d3
--- /dev/null
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/ClientControllerTest.java
@@ -0,0 +1,97 @@
+/*
+ * 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.server.inputmethod;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.junit.Assert.assertThrows;
+import static org.mockito.Mockito.when;
+
+import android.content.pm.PackageManagerInternal;
+import android.os.Handler;
+import android.os.IBinder;
+import android.os.Looper;
+import android.platform.test.annotations.IgnoreUnderRavenwood;
+import android.platform.test.ravenwood.RavenwoodRule;
+import android.view.Display;
+import android.view.inputmethod.InputBinding;
+
+import com.android.internal.inputmethod.IInputMethodClient;
+import com.android.internal.inputmethod.IRemoteInputConnection;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+
+// This test is designed to run on both device and host (Ravenwood) side.
+public final class ClientControllerTest {
+    private static final int ANY_DISPLAY_ID = Display.DEFAULT_DISPLAY;
+    private static final int ANY_CALLER_UID = 1;
+    private static final int ANY_CALLER_PID = 1;
+
+    @Rule
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true).build();
+
+    @Mock
+    private PackageManagerInternal mMockPackageManagerInternal;
+
+    @Mock(extraInterfaces = IBinder.class)
+    private IInputMethodClient mClient;
+
+    @Mock
+    private IRemoteInputConnection mConnection;
+
+    @Mock
+    private IBinder.DeathRecipient mDeathRecipient;
+
+    private Handler mHandler;
+
+    private ClientController mController;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mHandler = new Handler(Looper.getMainLooper());
+        mController = new ClientController(mMockPackageManagerInternal);
+        when(mClient.asBinder()).thenReturn((IBinder) mClient);
+    }
+
+    @Test
+    // TODO(b/314150112): Enable host side mode for this test once b/315544364 is fixed.
+    @IgnoreUnderRavenwood(blockedBy = {InputBinding.class, IInputMethodClientInvoker.class})
+    public void testAddClient_cannotAddTheSameClientTwice() {
+        var invoker = IInputMethodClientInvoker.create(mClient, mHandler);
+
+        synchronized (ImfLock.class) {
+            mController.addClient(invoker, mConnection, ANY_DISPLAY_ID, mDeathRecipient,
+                    ANY_CALLER_UID, ANY_CALLER_PID);
+
+            SecurityException thrown = assertThrows(SecurityException.class,
+                    () -> {
+                        synchronized (ImfLock.class) {
+                            mController.addClient(invoker, mConnection, ANY_DISPLAY_ID,
+                                    mDeathRecipient, ANY_CALLER_UID, ANY_CALLER_PID);
+                        }
+                    });
+            assertThat(thrown.getMessage()).isEqualTo(
+                    "uid=1/pid=1/displayId=0 is already registered");
+        }
+    }
+}
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/DefaultImeVisibilityApplierTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/DefaultImeVisibilityApplierTest.java
index 3199e062418f..438bea458c47 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/DefaultImeVisibilityApplierTest.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/DefaultImeVisibilityApplierTest.java
@@ -22,6 +22,7 @@ import static android.view.WindowManager.LayoutParams.SOFT_INPUT_STATE_ALWAYS_VI
 import static com.android.internal.inputmethod.SoftInputShowHideReason.HIDE_SOFT_INPUT;
 import static com.android.internal.inputmethod.SoftInputShowHideReason.HIDE_SWITCH_USER;
 import static com.android.internal.inputmethod.SoftInputShowHideReason.SHOW_SOFT_INPUT;
+import static com.android.server.inputmethod.ClientController.ClientState;
 import static com.android.server.inputmethod.ImeVisibilityStateComputer.STATE_HIDE_IME;
 import static com.android.server.inputmethod.ImeVisibilityStateComputer.STATE_HIDE_IME_EXPLICIT;
 import static com.android.server.inputmethod.ImeVisibilityStateComputer.STATE_HIDE_IME_NOT_ALWAYS;
@@ -68,8 +69,7 @@ public class DefaultImeVisibilityApplierTest extends InputMethodManagerServiceTe
         super.setUp();
         mVisibilityApplier =
                 (DefaultImeVisibilityApplier) mInputMethodManagerService.getVisibilityApplier();
-        mInputMethodManagerService.setAttachedClientForTesting(
-                mock(InputMethodManagerService.ClientState.class));
+        mInputMethodManagerService.setAttachedClientForTesting(mock(ClientState.class));
     }
 
     @Test