Add VDMInternal#isInputDeviceOwnedByVirtualDevice API

And remove the unused #isDisplayOwnedByAnyVirtualDevice API.

Test: atest VirtualDeviceManagerServiceTest
Fix: 319631178
Change-Id: If172fad34787c037c25ba56c2e085c578bf95f88
diff --git a/services/companion/java/com/android/server/companion/virtual/InputController.java b/services/companion/java/com/android/server/companion/virtual/InputController.java
index 1f89e57b..98ecca0 100644
--- a/services/companion/java/com/android/server/companion/virtual/InputController.java
+++ b/services/companion/java/com/android/server/companion/virtual/InputController.java
@@ -437,7 +437,7 @@
         }
     }
 
-    @VisibleForTesting
+    @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
     Map<IBinder, InputDeviceDescriptor> getInputDeviceDescriptors() {
         final Map<IBinder, InputDeviceDescriptor> inputDeviceDescriptors = new ArrayMap<>();
         synchronized (mLock) {
@@ -529,7 +529,8 @@
         }
     }
 
-    @VisibleForTesting static final class InputDeviceDescriptor {
+    @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
+    static final class InputDeviceDescriptor {
 
         static final int TYPE_KEYBOARD = 1;
         static final int TYPE_MOUSE = 2;
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 58aa2c3..b6a22c2 100644
--- a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
+++ b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
@@ -1337,6 +1337,11 @@
         }
     }
 
+    boolean isInputDeviceOwnedByVirtualDevice(int inputDeviceId) {
+        return mInputController.getInputDeviceDescriptors().values().stream().anyMatch(
+                inputDeviceDescriptor -> inputDeviceDescriptor.getInputDeviceId() == inputDeviceId);
+    }
+
     void onEnteringPipBlocked(int uid) {
         // Do nothing. ActivityRecord#checkEnterPictureInPictureState logs that the display does not
         // support PiP.
diff --git a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceManagerService.java b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceManagerService.java
index 0d5cdcb..ef61498 100644
--- a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceManagerService.java
+++ b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceManagerService.java
@@ -838,10 +838,11 @@
         }
 
         @Override
-        public boolean isDisplayOwnedByAnyVirtualDevice(int displayId) {
+        public boolean isInputDeviceOwnedByVirtualDevice(int inputDeviceId) {
             ArrayList<VirtualDeviceImpl> virtualDevicesSnapshot = getVirtualDevicesSnapshot();
             for (int i = 0; i < virtualDevicesSnapshot.size(); i++) {
-                if (virtualDevicesSnapshot.get(i).isDisplayOwnedByVirtualDevice(displayId)) {
+                if (virtualDevicesSnapshot.get(i)
+                        .isInputDeviceOwnedByVirtualDevice(inputDeviceId)) {
                     return true;
                 }
             }
diff --git a/services/core/java/com/android/server/companion/virtual/VirtualDeviceManagerInternal.java b/services/core/java/com/android/server/companion/virtual/VirtualDeviceManagerInternal.java
index 823788f..b179783 100644
--- a/services/core/java/com/android/server/companion/virtual/VirtualDeviceManagerInternal.java
+++ b/services/core/java/com/android/server/companion/virtual/VirtualDeviceManagerInternal.java
@@ -137,9 +137,9 @@
     public abstract boolean isAppRunningOnAnyVirtualDevice(int uid);
 
     /**
-     * Returns true if the {@code displayId} is owned by any virtual device
+     * @return whether the input device with the given id was created by a virtual device.
      */
-    public abstract boolean isDisplayOwnedByAnyVirtualDevice(int displayId);
+    public abstract boolean isInputDeviceOwnedByVirtualDevice(int inputDeviceId);
 
     /**
      * Gets the ids of VirtualDisplays owned by a VirtualDevice.
diff --git a/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java b/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java
index 995d1f4..a360288 100644
--- a/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java
@@ -607,6 +607,20 @@
     }
 
     @Test
+    public void testIsInputDeviceOwnedByVirtualDevice() {
+        assertThat(mLocalService.isInputDeviceOwnedByVirtualDevice(INPUT_DEVICE_ID)).isFalse();
+
+        final int fd = 1;
+        mInputController.addDeviceForTesting(BINDER, fd,
+                InputController.InputDeviceDescriptor.TYPE_KEYBOARD, DISPLAY_ID_1, PHYS,
+                DEVICE_NAME_1, INPUT_DEVICE_ID);
+        assertThat(mLocalService.isInputDeviceOwnedByVirtualDevice(INPUT_DEVICE_ID)).isTrue();
+
+        mInputController.unregisterInputDevice(BINDER);
+        assertThat(mLocalService.isInputDeviceOwnedByVirtualDevice(INPUT_DEVICE_ID)).isFalse();
+    }
+
+    @Test
     public void getDeviceIdsForUid_noRunningApps_returnsNull() {
         assertThat(mLocalService.getDeviceIdsForUid(UID_1)).isEmpty();
         assertThat(mVdmNative.getDeviceIdsForUid(UID_1)).isEmpty();