VDM IME 2/n: setInputMethodComponent API

New VirtualDeviceParams.Builder#setInputMethodComponent API that
for now only updates a map in IMMS.

The custom VD IME cannot be changed during the device lifetime.

Bug: 287269288
Test: CTS
Change-Id: Iad2401d19482030fca072bdec48935c9b574d465

6 files changed, 115 insertions, 5 deletions

diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index 7df979fccfe6..c9f2942d676d 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -3247,6 +3247,7 @@ package android.companion.virtual {
     method @Deprecated public int getDefaultNavigationPolicy();
     method public int getDevicePolicy(int);
     method @FlaggedApi("android.companion.virtual.flags.vdm_custom_home") @Nullable public android.content.ComponentName getHomeComponent();
+    method @FlaggedApi("android.companion.virtual.flags.vdm_custom_ime") @Nullable public android.content.ComponentName getInputMethodComponent();
     method public int getLockState();
     method @Nullable public String getName();
     method @NonNull public java.util.Set<android.os.UserHandle> getUsersWithMatchingAccounts();
@@ -3280,6 +3281,7 @@ package android.companion.virtual {
     method @Deprecated @NonNull public android.companion.virtual.VirtualDeviceParams.Builder setBlockedCrossTaskNavigations(@NonNull java.util.Set<android.content.ComponentName>);
     method @NonNull public android.companion.virtual.VirtualDeviceParams.Builder setDevicePolicy(int, int);
     method @FlaggedApi("android.companion.virtual.flags.vdm_custom_home") @NonNull public android.companion.virtual.VirtualDeviceParams.Builder setHomeComponent(@Nullable android.content.ComponentName);
+    method @FlaggedApi("android.companion.virtual.flags.vdm_custom_ime") @NonNull public android.companion.virtual.VirtualDeviceParams.Builder setInputMethodComponent(@Nullable android.content.ComponentName);
     method @NonNull @RequiresPermission(value=android.Manifest.permission.ADD_ALWAYS_UNLOCKED_DISPLAY, conditional=true) public android.companion.virtual.VirtualDeviceParams.Builder setLockState(int);
     method @NonNull public android.companion.virtual.VirtualDeviceParams.Builder setName(@NonNull String);
     method @NonNull public android.companion.virtual.VirtualDeviceParams.Builder setUsersWithMatchingAccounts(@NonNull java.util.Set<android.os.UserHandle>);
diff --git a/core/java/android/companion/virtual/VirtualDeviceParams.java b/core/java/android/companion/virtual/VirtualDeviceParams.java
index 0d73e44f5197..0253ddd93a44 100644
--- a/core/java/android/companion/virtual/VirtualDeviceParams.java
+++ b/core/java/android/companion/virtual/VirtualDeviceParams.java
@@ -258,6 +258,7 @@ public final class VirtualDeviceParams implements Parcelable {
     // Mapping of @PolicyType to @DevicePolicy
     @NonNull private final SparseIntArray mDevicePolicies;
     @Nullable private final ComponentName mHomeComponent;
+    @Nullable private final ComponentName mInputMethodComponent;
     @NonNull private final List<VirtualSensorConfig> mVirtualSensorConfigs;
     @Nullable private final IVirtualSensorCallback mVirtualSensorCallback;
     private final int mAudioPlaybackSessionId;
@@ -273,6 +274,7 @@ public final class VirtualDeviceParams implements Parcelable {
             @Nullable String name,
             @NonNull SparseIntArray devicePolicies,
             @Nullable ComponentName homeComponent,
+            @Nullable ComponentName inputMethodComponent,
             @NonNull List<VirtualSensorConfig> virtualSensorConfigs,
             @Nullable IVirtualSensorCallback virtualSensorCallback,
             int audioPlaybackSessionId,
@@ -289,6 +291,7 @@ public final class VirtualDeviceParams implements Parcelable {
         mName = name;
         mDevicePolicies = Objects.requireNonNull(devicePolicies);
         mHomeComponent = homeComponent;
+        mInputMethodComponent = inputMethodComponent;
         mVirtualSensorConfigs = Objects.requireNonNull(virtualSensorConfigs);
         mVirtualSensorCallback = virtualSensorCallback;
         mAudioPlaybackSessionId = audioPlaybackSessionId;
@@ -312,6 +315,7 @@ public final class VirtualDeviceParams implements Parcelable {
         mAudioPlaybackSessionId = parcel.readInt();
         mAudioRecordingSessionId = parcel.readInt();
         mHomeComponent = parcel.readTypedObject(ComponentName.CREATOR);
+        mInputMethodComponent = parcel.readTypedObject(ComponentName.CREATOR);
     }
 
     /**
@@ -336,6 +340,18 @@ public final class VirtualDeviceParams implements Parcelable {
     }
 
     /**
+     * Returns the custom component used as input method on all displays owned by this virtual
+     * device.
+     *
+     * @see Builder#setInputMethodComponent
+     */
+    @FlaggedApi(Flags.FLAG_VDM_CUSTOM_IME)
+    @Nullable
+    public ComponentName getInputMethodComponent() {
+        return mInputMethodComponent;
+    }
+
+    /**
      * Returns the user handles with matching managed accounts on the remote device to which
      * this virtual device is streaming.
      *
@@ -532,6 +548,7 @@ public final class VirtualDeviceParams implements Parcelable {
         dest.writeInt(mAudioPlaybackSessionId);
         dest.writeInt(mAudioRecordingSessionId);
         dest.writeTypedObject(mHomeComponent, flags);
+        dest.writeTypedObject(mInputMethodComponent, flags);
     }
 
     @Override
@@ -563,6 +580,8 @@ public final class VirtualDeviceParams implements Parcelable {
                 && Objects.equals(mActivityPolicyExemptions, that.mActivityPolicyExemptions)
                 && mDefaultActivityPolicy == that.mDefaultActivityPolicy
                 && Objects.equals(mName, that.mName)
+                && Objects.equals(mHomeComponent, that.mHomeComponent)
+                && Objects.equals(mInputMethodComponent, that.mInputMethodComponent)
                 && mAudioPlaybackSessionId == that.mAudioPlaybackSessionId
                 && mAudioRecordingSessionId == that.mAudioRecordingSessionId;
     }
@@ -572,7 +591,8 @@ public final class VirtualDeviceParams implements Parcelable {
         int hashCode = Objects.hash(
                 mLockState, mUsersWithMatchingAccounts, mCrossTaskNavigationExemptions,
                 mDefaultNavigationPolicy, mActivityPolicyExemptions, mDefaultActivityPolicy, mName,
-                mDevicePolicies, mHomeComponent, mAudioPlaybackSessionId, mAudioRecordingSessionId);
+                mDevicePolicies, mHomeComponent, mInputMethodComponent, mAudioPlaybackSessionId,
+                mAudioRecordingSessionId);
         for (int i = 0; i < mDevicePolicies.size(); i++) {
             hashCode = 31 * hashCode + mDevicePolicies.keyAt(i);
             hashCode = 31 * hashCode + mDevicePolicies.valueAt(i);
@@ -593,6 +613,7 @@ public final class VirtualDeviceParams implements Parcelable {
                 + " mName=" + mName
                 + " mDevicePolicies=" + mDevicePolicies
                 + " mHomeComponent=" + mHomeComponent
+                + " mInputMethodComponent=" + mInputMethodComponent
                 + " mAudioPlaybackSessionId=" + mAudioPlaybackSessionId
                 + " mAudioRecordingSessionId=" + mAudioRecordingSessionId
                 + ")";
@@ -612,6 +633,8 @@ public final class VirtualDeviceParams implements Parcelable {
         pw.println(prefix + "mActivityPolicyExemptions=" + mActivityPolicyExemptions);
         pw.println(prefix + "mDevicePolicies=" + mDevicePolicies);
         pw.println(prefix + "mVirtualSensorConfigs=" + mVirtualSensorConfigs);
+        pw.println(prefix + "mHomeComponent=" + mHomeComponent);
+        pw.println(prefix + "mInputMethodComponent=" + mInputMethodComponent);
         pw.println(prefix + "mAudioPlaybackSessionId=" + mAudioPlaybackSessionId);
         pw.println(prefix + "mAudioRecordingSessionId=" + mAudioRecordingSessionId);
     }
@@ -644,16 +667,17 @@ public final class VirtualDeviceParams implements Parcelable {
         private int mDefaultActivityPolicy = ACTIVITY_POLICY_DEFAULT_ALLOWED;
         private boolean mDefaultActivityPolicyConfigured = false;
         @Nullable private String mName;
-        @NonNull private SparseIntArray mDevicePolicies = new SparseIntArray();
+        @NonNull private final SparseIntArray mDevicePolicies = new SparseIntArray();
         private int mAudioPlaybackSessionId = AUDIO_SESSION_ID_GENERATE;
         private int mAudioRecordingSessionId = AUDIO_SESSION_ID_GENERATE;
 
-        @NonNull private List<VirtualSensorConfig> mVirtualSensorConfigs = new ArrayList<>();
+        @NonNull private final List<VirtualSensorConfig> mVirtualSensorConfigs = new ArrayList<>();
         @Nullable private Executor mVirtualSensorCallbackExecutor;
         @Nullable private VirtualSensorCallback mVirtualSensorCallback;
         @Nullable private Executor mVirtualSensorDirectChannelCallbackExecutor;
         @Nullable private VirtualSensorDirectChannelCallback mVirtualSensorDirectChannelCallback;
         @Nullable private ComponentName mHomeComponent;
+        @Nullable private ComponentName mInputMethodComponent;
 
         private static class VirtualSensorCallbackDelegate extends IVirtualSensorCallback.Stub {
             @NonNull
@@ -749,6 +773,28 @@ public final class VirtualDeviceParams implements Parcelable {
         }
 
         /**
+         * Specifies a component to be used as input method on all displays owned by this virtual
+         * device.
+         *
+         * @param inputMethodComponent The component name to be used as input method. Must comply to
+         *   all general input method requirements described in the guide to
+         *   <a href="{@docRoot}guide/topics/text/creating-input-method.html">
+         *   Creating an Input Method</a>. If the given component is not available for any user that
+         *   may interact with the virtual device, then there will effectively be no IME on this
+         *   device's displays for that user.
+         *
+         * @see android.inputmethodservice.InputMethodService
+         * @attr ref android.R.styleable#InputMethod_isVirtualDeviceOnly
+         * @attr ref android.R.styleable#InputMethod_showInInputMethodPicker
+         */
+        @FlaggedApi(Flags.FLAG_VDM_CUSTOM_IME)
+        @NonNull
+        public Builder setInputMethodComponent(@Nullable ComponentName inputMethodComponent) {
+            mInputMethodComponent = inputMethodComponent;
+            return this;
+        }
+
+        /**
          * Sets the user handles with matching managed accounts on the remote device to which
          * this virtual device is streaming. The caller is responsible for verifying the presence
          * and legitimacy of a matching managed account on the remote device.
@@ -1136,6 +1182,7 @@ public final class VirtualDeviceParams implements Parcelable {
                     mName,
                     mDevicePolicies,
                     mHomeComponent,
+                    mInputMethodComponent,
                     mVirtualSensorConfigs,
                     virtualSensorCallbackDelegate,
                     mAudioPlaybackSessionId,
diff --git a/core/res/res/values/attrs.xml b/core/res/res/values/attrs.xml
index 034a55a4a65d..e1c1a42eed81 100644
--- a/core/res/res/values/attrs.xml
+++ b/core/res/res/values/attrs.xml
@@ -3855,8 +3855,9 @@
              device -->
         <attr name="isVrOnly" format="boolean"/>
         <!-- Specifies if an IME can only be used on a display created by a virtual device.
-             @hide @SystemApi
-             @FlaggedApi("android.companion.virtual.flags.vdm_custom_ime") -->
+             @see android.companion.virtual.VirtualDeviceParams.Builder#setInputMethodComponent
+             @hide @SystemApi -->
+        <!-- @FlaggedApi("android.companion.virtual.flags.vdm_custom_ime") -->
         <attr name="isVirtualDeviceOnly" format="boolean"/>
         <attr name="__removed2" format="boolean" />
         <!-- Specifies whether the IME supports showing inline suggestions. -->
diff --git a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
index 45d7314f3fab..13c79248eb38 100644
--- a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
+++ b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
@@ -104,6 +104,7 @@ import com.android.server.LocalServices;
 import com.android.server.companion.virtual.GenericWindowPolicyController.RunningAppsChangedListener;
 import com.android.server.companion.virtual.audio.VirtualAudioController;
 import com.android.server.companion.virtual.camera.VirtualCameraController;
+import com.android.server.inputmethod.InputMethodManagerInternal;
 
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
@@ -352,6 +353,14 @@ final class VirtualDeviceImpl extends IVirtualDevice.Stub
             flags |= DisplayManager.VIRTUAL_DISPLAY_FLAG_ALWAYS_UNLOCKED;
         }
         mBaseVirtualDisplayFlags = flags;
+
+        if (Flags.vdmCustomIme() && mParams.getInputMethodComponent() != null) {
+            final String imeId = mParams.getInputMethodComponent().flattenToShortString();
+            Slog.d(TAG, "Setting custom input method " + imeId + " as default for virtual device "
+                    + deviceId);
+            InputMethodManagerInternal.get().setVirtualDeviceInputMethodForAllUsers(
+                    mDeviceId, imeId);
+        }
     }
 
     @VisibleForTesting
@@ -556,6 +565,12 @@ final class VirtualDeviceImpl extends IVirtualDevice.Stub
                 mCameraAccessController.stopObservingIfNeeded();
             }
 
+            // Clear any previously set custom IME components.
+            if (Flags.vdmCustomIme() && mParams.getInputMethodComponent() != null) {
+                InputMethodManagerInternal.get().setVirtualDeviceInputMethodForAllUsers(
+                        mDeviceId, null);
+            }
+
             mInputController.close();
             mSensorController.close();
         } finally {
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java
index 14daf62a9ed2..82a5c5aa8047 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java
@@ -120,6 +120,25 @@ public abstract class InputMethodManagerInternal {
             @UserIdInt int userId);
 
     /**
+     * Makes the input method associated with {@code imeId} the default input method for all users
+     * on displays that are owned by the virtual device with the given {@code deviceId}. If the
+     * input method associated with {@code imeId} is not available, there will be no IME on the
+     * relevant displays.
+     *
+     * <p>The caller of this method is responsible for resetting it to {@code null} after the
+     * virtual device is closed.</p>
+     *
+     * @param deviceId the device ID on which to use the given input method as default.
+     * @param imeId  the input method ID to be used as default on the given device. If {@code null},
+     *               then any existing input method association with that device will be removed.
+     * @throws IllegalArgumentException if a non-{@code null} input method ID is passed for a
+     *                                  device ID that already has a custom input method set or if
+     *                                  the device ID is not a valid virtual device.
+     */
+    public abstract void setVirtualDeviceInputMethodForAllUsers(
+            int deviceId, @Nullable String imeId);
+
+    /**
      * Registers a new {@link InputMethodListListener}.
      *
      * @param listener the listener to add
@@ -247,6 +266,11 @@ public abstract class InputMethodManagerInternal {
                 }
 
                 @Override
+                public void setVirtualDeviceInputMethodForAllUsers(
+                        int deviceId, @Nullable String imeId) {
+                }
+
+                @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 b7007854dab8..2d0035d7e3e1 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -169,6 +169,7 @@ import com.android.internal.os.TransferPipe;
 import com.android.internal.util.ArrayUtils;
 import com.android.internal.util.ConcurrentUtils;
 import com.android.internal.util.DumpUtils;
+import com.android.internal.util.Preconditions;
 import com.android.internal.view.IInputMethodManager;
 import com.android.server.AccessibilityManagerInternal;
 import com.android.server.EventLogTags;
@@ -312,6 +313,9 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
     // All known input methods.
     final ArrayList<InputMethodInfo> mMethodList = new ArrayList<>();
     private final ArrayMap<String, InputMethodInfo> mMethodMap = new ArrayMap<>();
+    // Mapping from deviceId to the device-specific imeId for that device.
+    private final SparseArray<String> mVirtualDeviceMethodMap = new SparseArray<>();
+
     final InputMethodSubtypeSwitchingController mSwitchingController;
     final HardwareKeyboardShortcutController mHardwareKeyboardShortcutController =
             new HardwareKeyboardShortcutController();
@@ -5627,6 +5631,23 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
         }
 
         @Override
+        public void setVirtualDeviceInputMethodForAllUsers(int deviceId, @Nullable String imeId) {
+            // TODO(b/287269288): validate that id belongs to a valid virtual device instead.
+            Preconditions.checkArgument(deviceId == Context.DEVICE_ID_DEFAULT,
+                    "DeviceId " + deviceId + " does not belong to a virtual device.");
+            synchronized (ImfLock.class) {
+                if (imeId == null) {
+                    mVirtualDeviceMethodMap.remove(deviceId);
+                } else if (mVirtualDeviceMethodMap.contains(deviceId)) {
+                    throw new IllegalArgumentException("Virtual device " + deviceId
+                            + " already has a custom input method component");
+                } else {
+                    mVirtualDeviceMethodMap.put(deviceId, imeId);
+                }
+            }
+        }
+
+        @Override
         public void registerInputMethodListListener(InputMethodListListener listener) {
             mInputMethodListListeners.addIfAbsent(listener);
         }