Add new setInputMethodEnabled API

Make accessibility services able to enable/disable a specified
IME in the same package for the current user.

Bug: 195476910
Fix: 191386474

Test: atest AccessibilitySoftKeyboardTest. Also tested with
modified talkback (use the swith input method API for enable
ime), and self created empty a11y service with an empty ime.
Mananged device and profile which set limit on user 0 are also
tested.

Change-Id: I4187468076705ac597d680f2f5dc32d7b166da1f
diff --git a/core/api/current.txt b/core/api/current.txt
index 56b6800..02f7b82 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -3144,8 +3144,12 @@
     method public void addOnShowModeChangedListener(@NonNull android.accessibilityservice.AccessibilityService.SoftKeyboardController.OnShowModeChangedListener, @Nullable android.os.Handler);
     method public int getShowMode();
     method public boolean removeOnShowModeChangedListener(@NonNull android.accessibilityservice.AccessibilityService.SoftKeyboardController.OnShowModeChangedListener);
+    method @CheckResult public int setInputMethodEnabled(@NonNull String, boolean) throws java.lang.SecurityException;
     method public boolean setShowMode(int);
     method public boolean switchToInputMethod(@NonNull String);
+    field public static final int ENABLE_IME_FAIL_BY_ADMIN = 1; // 0x1
+    field public static final int ENABLE_IME_FAIL_UNKNOWN = 2; // 0x2
+    field public static final int ENABLE_IME_SUCCESS = 0; // 0x0
   }
 
   public static interface AccessibilityService.SoftKeyboardController.OnShowModeChangedListener {
diff --git a/core/java/android/accessibilityservice/AccessibilityService.java b/core/java/android/accessibilityservice/AccessibilityService.java
index 09af72d..abc18f6 100644
--- a/core/java/android/accessibilityservice/AccessibilityService.java
+++ b/core/java/android/accessibilityservice/AccessibilityService.java
@@ -20,6 +20,7 @@
 
 import android.accessibilityservice.GestureDescription.MotionEventGenerator;
 import android.annotation.CallbackExecutor;
+import android.annotation.CheckResult;
 import android.annotation.ColorInt;
 import android.annotation.IntDef;
 import android.annotation.NonNull;
@@ -1640,6 +1641,29 @@
         private ArrayMap<OnShowModeChangedListener, Handler> mListeners;
         private final Object mLock;
 
+        /** @hide */
+        @Retention(RetentionPolicy.SOURCE)
+        @IntDef({
+                ENABLE_IME_SUCCESS,
+                ENABLE_IME_FAIL_BY_ADMIN,
+                ENABLE_IME_FAIL_UNKNOWN
+        })
+        public @interface EnableImeResult {}
+        /**
+         * Return value for {@link #setInputMethodEnabled(String, boolean)}. The action succeeded.
+         */
+        public static final int ENABLE_IME_SUCCESS = 0;
+        /**
+         * Return value for {@link #setInputMethodEnabled(String, boolean)}. The action failed
+         * because the InputMethod is not permitted by device policy manager.
+         */
+        public static final int ENABLE_IME_FAIL_BY_ADMIN = 1;
+        /**
+         * Return value for {@link #setInputMethodEnabled(String, boolean)}. The action failed
+         * and the reason is unknown.
+         */
+        public static final int ENABLE_IME_FAIL_UNKNOWN = 2;
+
         SoftKeyboardController(@NonNull AccessibilityService service, @NonNull Object lock) {
             mService = service;
             mLock = lock;
@@ -1868,6 +1892,39 @@
             }
             return false;
         }
+
+        /**
+         * Enable or disable the specified IME for the user for whom the service is activated. The
+         * IME needs to be in the same package as the service and needs to be allowed by device
+         * policy, if there is one. The change will persist until the specified IME is next
+         * explicitly enabled or disabled by whatever means, such as user choice, and may persist
+         * beyond the life cycle of the requesting service.
+         *
+         * @param imeId The ID of the input method to enable or disable. This IME must be installed.
+         * @param enabled {@code true} if the input method associated with {@code imeId} should be
+         *                enabled.
+         * @return status code for the result of enabling/disabling the input method associated
+         *         with {@code imeId}.
+         * @throws SecurityException if the input method is not in the same package as the service.
+         *
+         * @see android.view.inputmethod.InputMethodInfo#getId()
+         */
+        @CheckResult
+        @EnableImeResult
+        public int setInputMethodEnabled(@NonNull String imeId, boolean enabled)
+                throws SecurityException {
+            final IAccessibilityServiceConnection connection =
+                    AccessibilityInteractionClient.getInstance(mService).getConnection(
+                            mService.mConnectionId);
+            if (connection != null) {
+                try {
+                    return connection.setInputMethodEnabled(imeId, enabled);
+                } catch (RemoteException re) {
+                    throw new RuntimeException(re);
+                }
+            }
+            return ENABLE_IME_FAIL_UNKNOWN;
+        }
     }
 
     /**
diff --git a/core/java/android/accessibilityservice/IAccessibilityServiceConnection.aidl b/core/java/android/accessibilityservice/IAccessibilityServiceConnection.aidl
index 81457eb..2a3a680 100644
--- a/core/java/android/accessibilityservice/IAccessibilityServiceConnection.aidl
+++ b/core/java/android/accessibilityservice/IAccessibilityServiceConnection.aidl
@@ -100,6 +100,8 @@
 
     boolean switchToInputMethod(String imeId);
 
+    int setInputMethodEnabled(String imeId, boolean enabled);
+
     boolean isAccessibilityButtonAvailable();
 
     void sendGesture(int sequence, in ParceledListSlice gestureSteps);
diff --git a/core/tests/coretests/src/android/view/accessibility/AccessibilityServiceConnectionImpl.java b/core/tests/coretests/src/android/view/accessibility/AccessibilityServiceConnectionImpl.java
index c241e36..c4a1c4a 100644
--- a/core/tests/coretests/src/android/view/accessibility/AccessibilityServiceConnectionImpl.java
+++ b/core/tests/coretests/src/android/view/accessibility/AccessibilityServiceConnectionImpl.java
@@ -16,6 +16,7 @@
 
 package android.view.accessibility;
 
+import android.accessibilityservice.AccessibilityService;
 import android.accessibilityservice.AccessibilityServiceInfo;
 import android.accessibilityservice.IAccessibilityServiceConnection;
 import android.content.pm.ParceledListSlice;
@@ -138,6 +139,10 @@
         return false;
     }
 
+    public int setInputMethodEnabled(String imeId, boolean enabled) {
+        return AccessibilityService.SoftKeyboardController.ENABLE_IME_FAIL_UNKNOWN;
+    }
+
     public boolean isAccessibilityButtonAvailable() {
         return false;
     }
diff --git a/packages/SettingsLib/RestrictedLockUtils/Android.bp b/packages/SettingsLib/RestrictedLockUtils/Android.bp
index c0623ed..ef548b5 100644
--- a/packages/SettingsLib/RestrictedLockUtils/Android.bp
+++ b/packages/SettingsLib/RestrictedLockUtils/Android.bp
@@ -7,6 +7,12 @@
     default_applicable_licenses: ["frameworks_base_license"],
 }
 
+filegroup {
+    name: "SettingsLibRestrictedLockUtilsSrc",
+    srcs: ["src/**/*.java"],
+    visibility: ["//frameworks/base/services/accessibility"],
+}
+
 android_library {
     name: "SettingsLibRestrictedLockUtils",
 
diff --git a/services/accessibility/Android.bp b/services/accessibility/Android.bp
index 1698e9a..bf8a9af 100644
--- a/services/accessibility/Android.bp
+++ b/services/accessibility/Android.bp
@@ -16,7 +16,15 @@
 
 java_library_static {
     name: "services.accessibility",
-    defaults: ["platform_service_defaults"],
-    srcs: [":services.accessibility-sources"],
-    libs: ["services.core"],
+    defaults: [
+        "platform_service_defaults",
+    ],
+    srcs: [
+        ":services.accessibility-sources",
+        "//frameworks/base/packages/SettingsLib/RestrictedLockUtils:SettingsLibRestrictedLockUtilsSrc",
+    ],
+    libs: [
+        "services.core",
+        "androidx.annotation_annotation",
+    ],
 }
diff --git a/services/accessibility/java/com/android/server/accessibility/AccessibilitySecurityPolicy.java b/services/accessibility/java/com/android/server/accessibility/AccessibilitySecurityPolicy.java
index 0ab0c89..e251bcc 100644
--- a/services/accessibility/java/com/android/server/accessibility/AccessibilitySecurityPolicy.java
+++ b/services/accessibility/java/com/android/server/accessibility/AccessibilitySecurityPolicy.java
@@ -16,13 +16,18 @@
 
 package com.android.server.accessibility;
 
+import static android.accessibilityservice.AccessibilityService.SoftKeyboardController.ENABLE_IME_FAIL_BY_ADMIN;
+import static android.accessibilityservice.AccessibilityService.SoftKeyboardController.ENABLE_IME_SUCCESS;
 import static android.content.pm.PackageManagerInternal.PACKAGE_INSTALLER;
 
 import android.Manifest;
+import android.accessibilityservice.AccessibilityService;
 import android.accessibilityservice.AccessibilityServiceInfo;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.annotation.UserIdInt;
 import android.app.AppOpsManager;
+import android.app.admin.DevicePolicyManager;
 import android.appwidget.AppWidgetManagerInternal;
 import android.content.ComponentName;
 import android.content.Context;
@@ -41,13 +46,17 @@
 import android.util.ArraySet;
 import android.util.Slog;
 import android.view.accessibility.AccessibilityEvent;
+import android.view.inputmethod.InputMethodInfo;
 
 import com.android.internal.util.ArrayUtils;
 import com.android.server.LocalServices;
+import com.android.server.inputmethod.InputMethodManagerInternal;
+import com.android.settingslib.RestrictedLockUtils;
 
 import libcore.util.EmptyArray;
 
 import java.util.ArrayList;
+import java.util.List;
 import java.util.Set;
 
 /**
@@ -364,6 +373,115 @@
     }
 
     /**
+     * Check whether the input method can be enabled or disabled by the accessibility service.
+     *
+     * @param imeId The id of the input method.
+     * @param service The accessibility service connection.
+     * @return Whether the input method can be enabled/disabled or the reason why it can't be
+     *         enabled/disabled.
+     * @throws SecurityException if the input method is not in the same package as the service.
+     */
+    @AccessibilityService.SoftKeyboardController.EnableImeResult
+    int canEnableDisableInputMethod(String imeId, AbstractAccessibilityServiceConnection service)
+            throws SecurityException {
+        final String servicePackageName = service.getComponentName().getPackageName();
+        final int callingUserId = UserHandle.getCallingUserId();
+
+        InputMethodInfo inputMethodInfo = null;
+        List<InputMethodInfo> inputMethodInfoList =
+                InputMethodManagerInternal.get().getInputMethodListAsUser(callingUserId);
+        if (inputMethodInfoList != null) {
+            for (InputMethodInfo info : inputMethodInfoList) {
+                if (info.getId().equals(imeId)) {
+                    inputMethodInfo = info;
+                    break;
+                }
+            }
+        }
+
+        if (inputMethodInfo == null
+                || !inputMethodInfo.getPackageName().equals(servicePackageName)) {
+            throw new SecurityException("The input method is in a different package with the "
+                    + "accessibility service");
+        }
+
+        // TODO(b/207697949, b/208872785): Add cts test for managed device.
+        //  Use RestrictedLockUtilsInternal in AccessibilitySecurityPolicy
+        if (checkIfInputMethodDisallowed(
+                mContext, inputMethodInfo.getPackageName(), callingUserId) != null) {
+            return ENABLE_IME_FAIL_BY_ADMIN;
+        }
+
+        return ENABLE_IME_SUCCESS;
+    }
+
+    /**
+     * @return the UserHandle for a userId. Return null for USER_NULL
+     */
+    private static UserHandle getUserHandleOf(@UserIdInt int userId) {
+        if (userId == UserHandle.USER_NULL) {
+            return null;
+        } else {
+            return UserHandle.of(userId);
+        }
+    }
+
+    private static int getManagedProfileId(Context context, int userId) {
+        UserManager um = context.getSystemService(UserManager.class);
+        List<UserInfo> userProfiles = um.getProfiles(userId);
+        for (UserInfo uInfo : userProfiles) {
+            if (uInfo.id == userId) {
+                continue;
+            }
+            if (uInfo.isManagedProfile()) {
+                return uInfo.id;
+            }
+        }
+        return UserHandle.USER_NULL;
+    }
+
+    private static RestrictedLockUtils.EnforcedAdmin checkIfInputMethodDisallowed(Context context,
+            String packageName, int userId) {
+        DevicePolicyManager dpm = context.getSystemService(DevicePolicyManager.class);
+        if (dpm == null) {
+            return null;
+        }
+        RestrictedLockUtils.EnforcedAdmin admin =
+                RestrictedLockUtils.getProfileOrDeviceOwner(context, getUserHandleOf(userId));
+        boolean permitted = true;
+        if (admin != null) {
+            permitted = dpm.isInputMethodPermittedByAdmin(admin.component,
+                    packageName, userId);
+        }
+
+        boolean permittedByParentAdmin = true;
+        RestrictedLockUtils.EnforcedAdmin profileAdmin = null;
+        int managedProfileId = getManagedProfileId(context, userId);
+        if (managedProfileId != UserHandle.USER_NULL) {
+            profileAdmin = RestrictedLockUtils.getProfileOrDeviceOwner(
+                    context, getUserHandleOf(managedProfileId));
+            // If the device is an organization-owned device with a managed profile, the
+            // managedProfileId will be used instead of the affected userId. This is because
+            // isInputMethodPermittedByAdmin is called on the parent DPM instance, which will
+            // return results affecting the personal profile.
+            if (profileAdmin != null && dpm.isOrganizationOwnedDeviceWithManagedProfile()) {
+                DevicePolicyManager parentDpm = dpm.getParentProfileInstance(
+                        UserManager.get(context).getUserInfo(managedProfileId));
+                permittedByParentAdmin = parentDpm.isInputMethodPermittedByAdmin(
+                        profileAdmin.component, packageName, managedProfileId);
+            }
+        }
+        if (!permitted && !permittedByParentAdmin) {
+            return RestrictedLockUtils.EnforcedAdmin.MULTIPLE_ENFORCED_ADMIN;
+        } else if (!permitted) {
+            return admin;
+        } else if (!permittedByParentAdmin) {
+            return profileAdmin;
+        }
+        return null;
+    }
+
+    /**
      * Returns the parent userId of the profile according to the specified userId.
      *
      * @param userId The userId to check
diff --git a/services/accessibility/java/com/android/server/accessibility/AccessibilityServiceConnection.java b/services/accessibility/java/com/android/server/accessibility/AccessibilityServiceConnection.java
index bcb3413..8f7260f 100644
--- a/services/accessibility/java/com/android/server/accessibility/AccessibilityServiceConnection.java
+++ b/services/accessibility/java/com/android/server/accessibility/AccessibilityServiceConnection.java
@@ -16,9 +16,13 @@
 
 package com.android.server.accessibility;
 
+import static android.accessibilityservice.AccessibilityService.SoftKeyboardController.ENABLE_IME_FAIL_UNKNOWN;
+import static android.accessibilityservice.AccessibilityService.SoftKeyboardController.ENABLE_IME_SUCCESS;
+
 import static com.android.internal.util.function.pooled.PooledLambda.obtainMessage;
 
 import android.Manifest;
+import android.accessibilityservice.AccessibilityService;
 import android.accessibilityservice.AccessibilityServiceInfo;
 import android.accessibilityservice.AccessibilityTrace;
 import android.accessibilityservice.IAccessibilityServiceClient;
@@ -310,6 +314,41 @@
     }
 
     @Override
+    @AccessibilityService.SoftKeyboardController.EnableImeResult
+    public int setInputMethodEnabled(String imeId, boolean enabled) throws SecurityException {
+        if (svcConnTracingEnabled()) {
+            logTraceSvcConn("switchToInputMethod", "imeId=" + imeId);
+        }
+        synchronized (mLock) {
+            if (!hasRightsToCurrentUserLocked()) {
+                return ENABLE_IME_FAIL_UNKNOWN;
+            }
+        }
+
+        final int callingUserId = UserHandle.getCallingUserId();
+        final InputMethodManagerInternal inputMethodManagerInternal =
+                InputMethodManagerInternal.get();
+
+        final @AccessibilityService.SoftKeyboardController.EnableImeResult int checkResult;
+        final long identity = Binder.clearCallingIdentity();
+        try {
+            synchronized (mLock) {
+                checkResult = mSecurityPolicy.canEnableDisableInputMethod(imeId, this);
+            }
+            if (checkResult != ENABLE_IME_SUCCESS) {
+                return checkResult;
+            }
+            if (inputMethodManagerInternal.setInputMethodEnabled(imeId,
+                        enabled, callingUserId)) {
+                return ENABLE_IME_SUCCESS;
+            }
+        } finally {
+            Binder.restoreCallingIdentity(identity);
+        }
+        return ENABLE_IME_FAIL_UNKNOWN;
+    }
+
+    @Override
     public boolean isAccessibilityButtonAvailable() {
         if (svcConnTracingEnabled()) {
             logTraceSvcConn("isAccessibilityButtonAvailable", "");
diff --git a/services/accessibility/java/com/android/server/accessibility/UiAutomationManager.java b/services/accessibility/java/com/android/server/accessibility/UiAutomationManager.java
index 7ee0690..6cc2214 100644
--- a/services/accessibility/java/com/android/server/accessibility/UiAutomationManager.java
+++ b/services/accessibility/java/com/android/server/accessibility/UiAutomationManager.java
@@ -16,6 +16,7 @@
 
 package com.android.server.accessibility;
 
+import android.accessibilityservice.AccessibilityService;
 import android.accessibilityservice.AccessibilityServiceInfo;
 import android.accessibilityservice.AccessibilityTrace;
 import android.accessibilityservice.IAccessibilityServiceClient;
@@ -333,6 +334,11 @@
         }
 
         @Override
+        public int setInputMethodEnabled(String imeId, boolean enabled) {
+            return AccessibilityService.SoftKeyboardController.ENABLE_IME_FAIL_UNKNOWN;
+        }
+
+        @Override
         public boolean isAccessibilityButtonAvailable() {
             return false;
         }
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java
index 2328dfc..9a53d83 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java
@@ -95,6 +95,21 @@
     public abstract boolean switchToInputMethod(String imeId, @UserIdInt int userId);
 
     /**
+     * Force enable or disable the input method associated with {@code imeId} for given user. If
+     * the input method associated with {@code imeId} is not installed, do nothing.
+     *
+     * @param imeId  The input method ID to be enabled or disabled.
+     * @param enabled {@code true} if the input method associated with {@code imeId} should be
+     *                enabled.
+     * @param userId The user ID to be queried.
+     * @return {@code true} if the input method associated with {@code imeId} was successfully
+     *         enabled or disabled, {@code false} if the input method specified is not installed
+     *         or was unable to be enabled/disabled for some other reason.
+     */
+    public abstract boolean setInputMethodEnabled(String imeId, boolean enabled,
+            @UserIdInt int userId);
+
+    /**
      * Registers a new {@link InputMethodListListener}.
      */
     public abstract void registerInputMethodListListener(InputMethodListListener listener);
@@ -168,6 +183,11 @@
                 }
 
                 @Override
+                public boolean setInputMethodEnabled(String imeId, boolean enabled, int userId) {
+                    return false;
+                }
+
+                @Override
                 public void registerInputMethodListListener(InputMethodListListener listener) {
                 }
 
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index 9499e51..dc13328 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -1937,13 +1937,7 @@
         if (userId == mSettings.getCurrentUserId()) {
             return mSettings.getEnabledInputMethodListLocked();
         }
-        final ArrayMap<String, InputMethodInfo> methodMap = new ArrayMap<>();
-        final ArrayList<InputMethodInfo> methodList = new ArrayList<>();
-        final ArrayMap<String, List<InputMethodSubtype>> additionalSubtypeMap =
-                new ArrayMap<>();
-        AdditionalSubtypeUtils.load(additionalSubtypeMap, userId);
-        queryInputMethodServicesInternal(mContext, userId, additionalSubtypeMap, methodMap,
-                methodList);
+        final ArrayMap<String, InputMethodInfo> methodMap = queryMethodMapForUser(userId);
         final InputMethodSettings settings = new InputMethodSettings(mContext.getResources(),
                 mContext.getContentResolver(), methodMap, userId, true);
         return settings.getEnabledInputMethodListLocked();
@@ -2108,13 +2102,7 @@
             return mSettings.getEnabledInputMethodSubtypeListLocked(
                     mContext, imi, allowsImplicitlySelectedSubtypes);
         }
-        final ArrayMap<String, InputMethodInfo> methodMap = new ArrayMap<>();
-        final ArrayList<InputMethodInfo> methodList = new ArrayList<>();
-        final ArrayMap<String, List<InputMethodSubtype>> additionalSubtypeMap =
-                new ArrayMap<>();
-        AdditionalSubtypeUtils.load(additionalSubtypeMap, userId);
-        queryInputMethodServicesInternal(mContext, userId, additionalSubtypeMap, methodMap,
-                methodList);
+        final ArrayMap<String, InputMethodInfo> methodMap = queryMethodMapForUser(userId);
         final InputMethodInfo imi = methodMap.get(imiId);
         if (imi == null) {
             return Collections.emptyList();
@@ -4814,31 +4802,36 @@
         }
     }
 
+    private ArrayMap<String, InputMethodInfo> queryMethodMapForUser(@UserIdInt int userId) {
+        final ArrayMap<String, InputMethodInfo> methodMap = new ArrayMap<>();
+        final ArrayList<InputMethodInfo> methodList = new ArrayList<>();
+        final ArrayMap<String, List<InputMethodSubtype>> additionalSubtypeMap =
+                new ArrayMap<>();
+        AdditionalSubtypeUtils.load(additionalSubtypeMap, userId);
+        queryInputMethodServicesInternal(mContext, userId, additionalSubtypeMap,
+                methodMap, methodList);
+        return methodMap;
+    }
+
     private boolean switchToInputMethod(String imeId, @UserIdInt int userId) {
         synchronized (mMethodMap) {
             if (userId == mSettings.getCurrentUserId()) {
                 if (!mMethodMap.containsKey(imeId)
                         || !mSettings.getEnabledInputMethodListLocked()
                                 .contains(mMethodMap.get(imeId))) {
-                    return false; // IME is not is found or not enabled.
+                    return false; // IME is not found or not enabled.
                 }
                 setInputMethodLocked(imeId, NOT_A_SUBTYPE_ID);
                 return true;
             }
-            final ArrayMap<String, InputMethodInfo> methodMap = new ArrayMap<>();
-            final ArrayList<InputMethodInfo> methodList = new ArrayList<>();
-            final ArrayMap<String, List<InputMethodSubtype>> additionalSubtypeMap =
-                    new ArrayMap<>();
-            AdditionalSubtypeUtils.load(additionalSubtypeMap, userId);
-            queryInputMethodServicesInternal(mContext, userId, additionalSubtypeMap,
-                    methodMap, methodList);
+            final ArrayMap<String, InputMethodInfo> methodMap = queryMethodMapForUser(userId);
             final InputMethodSettings settings = new InputMethodSettings(
                     mContext.getResources(), mContext.getContentResolver(), methodMap,
                     userId, false);
             if (!methodMap.containsKey(imeId)
                     || !settings.getEnabledInputMethodListLocked()
                             .contains(methodMap.get(imeId))) {
-                return false; // IME is not is found or not enabled.
+                return false; // IME is not found or not enabled.
             }
             settings.putSelectedInputMethod(imeId);
             settings.putSelectedSubtype(NOT_A_SUBTYPE_ID);
@@ -4846,6 +4839,35 @@
         }
     }
 
+    private boolean setInputMethodEnabled(String imeId, boolean enabled, @UserIdInt int userId) {
+        synchronized (mMethodMap) {
+            if (userId == mSettings.getCurrentUserId()) {
+                if (!mMethodMap.containsKey(imeId)) {
+                    return false; // IME is not found.
+                }
+                setInputMethodEnabledLocked(imeId, enabled);
+                return true;
+            }
+            final ArrayMap<String, InputMethodInfo> methodMap = queryMethodMapForUser(userId);
+            final InputMethodSettings settings = new InputMethodSettings(
+                    mContext.getResources(), mContext.getContentResolver(), methodMap,
+                    userId, false);
+            if (!methodMap.containsKey(imeId)) {
+                return false; // IME is not found.
+            }
+            if (enabled) {
+                if (!settings.getEnabledInputMethodListLocked().contains(methodMap.get(imeId))) {
+                    settings.appendAndPutEnabledInputMethodLocked(imeId, false);
+                }
+            } else {
+                settings.buildAndPutEnabledInputMethodsStrRemovingIdLocked(
+                        new StringBuilder(),
+                        settings.getEnabledInputMethodsAndSubtypeListLocked(), imeId);
+            }
+            return true;
+        }
+    }
+
     private boolean transferTouchFocusToImeWindow(@NonNull IBinder sourceInputToken,
             int displayId) {
         //TODO(b/150843766): Check if Input Token is valid.
@@ -4918,6 +4940,11 @@
         }
 
         @Override
+        public boolean setInputMethodEnabled(String imeId, boolean enabled, int userId) {
+            return mService.setInputMethodEnabled(imeId, enabled, userId);
+        }
+
+        @Override
         public void registerInputMethodListListener(InputMethodListListener listener) {
             mService.mInputMethodListListeners.addIfAbsent(listener);
         }
@@ -5513,13 +5540,7 @@
                 previouslyEnabled = setInputMethodEnabledLocked(imeId, enabled);
             }
         } else {
-            final ArrayMap<String, InputMethodInfo> methodMap = new ArrayMap<>();
-            final ArrayList<InputMethodInfo> methodList = new ArrayList<>();
-            final ArrayMap<String, List<InputMethodSubtype>> additionalSubtypeMap =
-                    new ArrayMap<>();
-            AdditionalSubtypeUtils.load(additionalSubtypeMap, userId);
-            queryInputMethodServicesInternal(mContext, userId, additionalSubtypeMap,
-                    methodMap, methodList);
+            final ArrayMap<String, InputMethodInfo> methodMap = queryMethodMapForUser(userId);
             final InputMethodSettings settings = new InputMethodSettings(mContext.getResources(),
                     mContext.getContentResolver(), methodMap, userId, false);
             if (enabled) {
diff --git a/services/tests/servicestests/src/com/android/server/accessibility/AbstractAccessibilityServiceConnectionTest.java b/services/tests/servicestests/src/com/android/server/accessibility/AbstractAccessibilityServiceConnectionTest.java
index a83d51b..f0c90f4 100644
--- a/services/tests/servicestests/src/com/android/server/accessibility/AbstractAccessibilityServiceConnectionTest.java
+++ b/services/tests/servicestests/src/com/android/server/accessibility/AbstractAccessibilityServiceConnectionTest.java
@@ -65,6 +65,7 @@
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 
+import android.accessibilityservice.AccessibilityService;
 import android.accessibilityservice.AccessibilityServiceInfo;
 import android.accessibilityservice.AccessibilityTrace;
 import android.accessibilityservice.IAccessibilityServiceClient;
@@ -840,6 +841,11 @@
         }
 
         @Override
+        public int setInputMethodEnabled(String imeId, boolean enabled) throws RemoteException {
+            return AccessibilityService.SoftKeyboardController.ENABLE_IME_FAIL_UNKNOWN;
+        }
+
+        @Override
         public boolean isAccessibilityButtonAvailable() throws RemoteException {
             return false;
         }