Add surface_control_input_receiver native API

Bug: 324271765
Test: ASurfaceControlInputReceiverTest
Change-Id: Ia85a07af09878846c681e552c1f8471652f932fc
diff --git a/core/api/current.txt b/core/api/current.txt
index 925cf4c..cc7dcb1 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -54419,9 +54419,9 @@
     method @Deprecated public android.view.Display getDefaultDisplay();
     method @NonNull public default android.view.WindowMetrics getMaximumWindowMetrics();
     method public default boolean isCrossWindowBlurEnabled();
-    method @FlaggedApi("com.android.window.flags.surface_control_input_receiver") @NonNull public default android.window.InputTransferToken registerBatchedSurfaceControlInputReceiver(int, @NonNull android.window.InputTransferToken, @NonNull android.view.SurfaceControl, @NonNull android.view.Choreographer, @NonNull android.view.SurfaceControlInputReceiver);
+    method @FlaggedApi("com.android.window.flags.surface_control_input_receiver") @NonNull public default android.window.InputTransferToken registerBatchedSurfaceControlInputReceiver(@NonNull android.window.InputTransferToken, @NonNull android.view.SurfaceControl, @NonNull android.view.Choreographer, @NonNull android.view.SurfaceControlInputReceiver);
     method @FlaggedApi("com.android.window.flags.trusted_presentation_listener_for_window") public default void registerTrustedPresentationListener(@NonNull android.os.IBinder, @NonNull android.window.TrustedPresentationThresholds, @NonNull java.util.concurrent.Executor, @NonNull java.util.function.Consumer<java.lang.Boolean>);
-    method @FlaggedApi("com.android.window.flags.surface_control_input_receiver") @NonNull public default android.window.InputTransferToken registerUnbatchedSurfaceControlInputReceiver(int, @NonNull android.window.InputTransferToken, @NonNull android.view.SurfaceControl, @NonNull android.os.Looper, @NonNull android.view.SurfaceControlInputReceiver);
+    method @FlaggedApi("com.android.window.flags.surface_control_input_receiver") @NonNull public default android.window.InputTransferToken registerUnbatchedSurfaceControlInputReceiver(@NonNull android.window.InputTransferToken, @NonNull android.view.SurfaceControl, @NonNull android.os.Looper, @NonNull android.view.SurfaceControlInputReceiver);
     method public default void removeCrossWindowBlurEnabledListener(@NonNull java.util.function.Consumer<java.lang.Boolean>);
     method public default void removeProposedRotationListener(@NonNull java.util.function.IntConsumer);
     method @FlaggedApi("com.android.window.flags.screen_recording_callbacks") @RequiresPermission(android.Manifest.permission.DETECT_SCREEN_RECORDING) public default void removeScreenRecordingCallback(@NonNull java.util.function.Consumer<java.lang.Integer>);
diff --git a/core/java/android/view/WindowManager.java b/core/java/android/view/WindowManager.java
index 107c5f2..71f1c3438 100644
--- a/core/java/android/view/WindowManager.java
+++ b/core/java/android/view/WindowManager.java
@@ -6193,9 +6193,6 @@
      * caller must invoke {@link #unregisterSurfaceControlInputReceiver(SurfaceControl)} to clean up
      * the resources when no longer needing to use the {@link SurfaceControlInputReceiver}
      *
-     * @param displayId              The display that the SurfaceControl will be placed on. Input
-     *                               will only work if SurfaceControl is on that display and that
-     *                               display  was touched.
      * @param surfaceControl         The SurfaceControl to register the InputChannel for
      * @param hostInputTransferToken The host token to link the embedded. This is used to handle
      *                               transferring touch gesture from host to embedded and for ANRs
@@ -6210,7 +6207,7 @@
      */
     @FlaggedApi(Flags.FLAG_SURFACE_CONTROL_INPUT_RECEIVER)
     @NonNull
-    default InputTransferToken registerBatchedSurfaceControlInputReceiver(int displayId,
+    default InputTransferToken registerBatchedSurfaceControlInputReceiver(
             @NonNull InputTransferToken hostInputTransferToken,
             @NonNull SurfaceControl surfaceControl, @NonNull Choreographer choreographer,
             @NonNull SurfaceControlInputReceiver receiver) {
@@ -6221,15 +6218,12 @@
     /**
      * Registers a {@link SurfaceControlInputReceiver} for a {@link SurfaceControl} that will
      * receive every input event. This is different than calling
-     * {@link #registerBatchedSurfaceControlInputReceiver(int, InputTransferToken, SurfaceControl,
+     * {@link #registerBatchedSurfaceControlInputReceiver(InputTransferToken, SurfaceControl,
      * Choreographer, SurfaceControlInputReceiver)} in that the input events are received
      * unbatched.
      * The caller must invoke {@link #unregisterSurfaceControlInputReceiver(SurfaceControl)} to
      * clean up the resources when no longer needing to use the {@link SurfaceControlInputReceiver}
      *
-     * @param displayId              The display that the SurfaceControl will be placed on. Input
-     *                               will only work if SurfaceControl is on that display and that
-     *                               display  was touched.
      * @param surfaceControl         The SurfaceControl to register the InputChannel for
      * @param hostInputTransferToken The host token to link the embedded. This is used to handle
      *                               transferring touch gesture from host to embedded and for ANRs
@@ -6243,7 +6237,7 @@
      */
     @FlaggedApi(Flags.FLAG_SURFACE_CONTROL_INPUT_RECEIVER)
     @NonNull
-    default InputTransferToken registerUnbatchedSurfaceControlInputReceiver(int displayId,
+    default InputTransferToken registerUnbatchedSurfaceControlInputReceiver(
             @NonNull InputTransferToken hostInputTransferToken,
             @NonNull SurfaceControl surfaceControl, @NonNull Looper looper,
             @NonNull SurfaceControlInputReceiver receiver) {
@@ -6256,9 +6250,9 @@
      * specified token.
      * <p>
      * Must be called on the same {@link Looper} thread to which was passed to the
-     * {@link #registerBatchedSurfaceControlInputReceiver(int, InputTransferToken, SurfaceControl,
+     * {@link #registerBatchedSurfaceControlInputReceiver(InputTransferToken, SurfaceControl,
      * Choreographer, SurfaceControlInputReceiver)} or
-     * {@link #registerUnbatchedSurfaceControlInputReceiver(int, InputTransferToken, SurfaceControl,
+     * {@link #registerUnbatchedSurfaceControlInputReceiver(InputTransferToken, SurfaceControl,
      * Looper, SurfaceControlInputReceiver)}
      *
      * @param surfaceControl The SurfaceControl to remove and unregister the input channel for.
@@ -6272,9 +6266,9 @@
     /**
      * Returns the input client token for the {@link SurfaceControl}. This will only return non
      * null if the SurfaceControl was registered for input via
-     * {@link #registerBatchedSurfaceControlInputReceiver(int, InputTransferToken, SurfaceControl,
+     * {@link #registerBatchedSurfaceControlInputReceiver(InputTransferToken, SurfaceControl,
      * Choreographer, SurfaceControlInputReceiver)} or
-     * {@link #registerUnbatchedSurfaceControlInputReceiver(int, InputTransferToken,
+     * {@link #registerUnbatchedSurfaceControlInputReceiver(InputTransferToken,
      * SurfaceControl, Looper, SurfaceControlInputReceiver)}.
      * <p>
      * This is helpful for testing to ensure the test waits for the layer to be registered with
@@ -6304,9 +6298,9 @@
      * </li>
      * <li>
      * Registering a SurfaceControl for input and passing the host's token to either
-     * {@link #registerBatchedSurfaceControlInputReceiver(int, InputTransferToken, SurfaceControl,
+     * {@link #registerBatchedSurfaceControlInputReceiver(InputTransferToken, SurfaceControl,
      * Choreographer, SurfaceControlInputReceiver)} or
-     * {@link #registerUnbatchedSurfaceControlInputReceiver(int, InputTransferToken,
+     * {@link #registerUnbatchedSurfaceControlInputReceiver(InputTransferToken,
      * SurfaceControl, Looper, SurfaceControlInputReceiver)}.
      * </li>
      * </ul>
@@ -6321,9 +6315,9 @@
      * When the host wants to transfer touch gesture to the embedded, it can retrieve the embedded
      * token via {@link SurfaceControlViewHost.SurfacePackage#getInputTransferToken()} or use the
      * value returned from either
-     * {@link #registerBatchedSurfaceControlInputReceiver(int, InputTransferToken, SurfaceControl,
+     * {@link #registerBatchedSurfaceControlInputReceiver(InputTransferToken, SurfaceControl,
      * Choreographer, SurfaceControlInputReceiver)} or
-     * {@link #registerUnbatchedSurfaceControlInputReceiver(int, InputTransferToken, SurfaceControl,
+     * {@link #registerUnbatchedSurfaceControlInputReceiver(InputTransferToken, SurfaceControl,
      * Looper, SurfaceControlInputReceiver)} and pass its own token as the transferFromToken.
      * <p>
      * When the embedded wants to transfer touch gesture to the host, it can pass in its own
diff --git a/core/java/android/view/WindowManagerGlobal.java b/core/java/android/view/WindowManagerGlobal.java
index c6d0454..961a9c4 100644
--- a/core/java/android/view/WindowManagerGlobal.java
+++ b/core/java/android/view/WindowManagerGlobal.java
@@ -16,6 +16,7 @@
 
 package android.view;
 
+import static android.view.Display.INVALID_DISPLAY;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION;
 
 import android.animation.ValueAnimator;
@@ -839,20 +840,40 @@
         mTrustedPresentationListener.removeListener(listener);
     }
 
-    InputTransferToken registerBatchedSurfaceControlInputReceiver(int displayId,
+    private static InputChannel createInputChannel(@NonNull IBinder clientToken,
             @NonNull InputTransferToken hostToken, @NonNull SurfaceControl surfaceControl,
-            @NonNull Choreographer choreographer, @NonNull SurfaceControlInputReceiver receiver) {
-        IBinder clientToken = new Binder();
-        InputTransferToken inputTransferToken = new InputTransferToken();
+            @Nullable InputTransferToken inputTransferToken) {
         InputChannel inputChannel = new InputChannel();
         try {
-            WindowManagerGlobal.getWindowSession().grantInputChannel(displayId, surfaceControl,
-                    clientToken, hostToken, 0, 0, TYPE_APPLICATION, 0, null, inputTransferToken,
-                    surfaceControl.getName(), inputChannel);
+            // TODO (b/329860681): Use INVALID_DISPLAY for now because the displayId will be
+            // selected in  SurfaceFlinger. This should be cleaned up so grantInputChannel doesn't
+            // take in a displayId at all
+            WindowManagerGlobal.getWindowSession().grantInputChannel(INVALID_DISPLAY,
+                    surfaceControl, clientToken, hostToken, 0, 0, TYPE_APPLICATION, 0, null,
+                    inputTransferToken, surfaceControl.getName(), inputChannel);
         } catch (RemoteException e) {
             Log.e(TAG, "Failed to create input channel", e);
             e.rethrowAsRuntimeException();
         }
+        return inputChannel;
+    }
+
+    private static void removeInputChannel(IBinder clientToken) {
+        try {
+            WindowManagerGlobal.getWindowSession().remove(clientToken);
+        } catch (RemoteException e) {
+            Log.e(TAG, "Failed to remove input channel", e);
+            e.rethrowAsRuntimeException();
+        }
+    }
+
+    InputTransferToken registerBatchedSurfaceControlInputReceiver(
+            @NonNull InputTransferToken hostToken, @NonNull SurfaceControl surfaceControl,
+            @NonNull Choreographer choreographer, @NonNull SurfaceControlInputReceiver receiver) {
+        IBinder clientToken = new Binder();
+        InputTransferToken inputTransferToken = new InputTransferToken();
+        InputChannel inputChannel = createInputChannel(clientToken, hostToken,
+                surfaceControl, inputTransferToken);
 
         synchronized (mSurfaceControlInputReceivers) {
             mSurfaceControlInputReceivers.put(surfaceControl.getLayerId(),
@@ -869,20 +890,13 @@
         return inputTransferToken;
     }
 
-    InputTransferToken registerUnbatchedSurfaceControlInputReceiver(int displayId,
+    InputTransferToken registerUnbatchedSurfaceControlInputReceiver(
             @NonNull InputTransferToken hostToken, @NonNull SurfaceControl surfaceControl,
             @NonNull Looper looper, @NonNull SurfaceControlInputReceiver receiver) {
         IBinder clientToken = new Binder();
         InputTransferToken inputTransferToken = new InputTransferToken();
-        InputChannel inputChannel = new InputChannel();
-        try {
-            WindowManagerGlobal.getWindowSession().grantInputChannel(displayId, surfaceControl,
-                    clientToken, hostToken, 0, 0, TYPE_APPLICATION, 0, null, inputTransferToken,
-                    surfaceControl.getName(), inputChannel);
-        } catch (RemoteException e) {
-            Log.e(TAG, "Failed to create input channel", e);
-            e.rethrowAsRuntimeException();
-        }
+        InputChannel inputChannel = createInputChannel(clientToken, hostToken,
+                surfaceControl, inputTransferToken);
 
         synchronized (mSurfaceControlInputReceivers) {
             mSurfaceControlInputReceivers.put(surfaceControl.getLayerId(),
@@ -909,13 +923,7 @@
             Log.w(TAG, "No registered input event receiver with sc: " + surfaceControl);
             return;
         }
-        try {
-            WindowManagerGlobal.getWindowSession().remove(
-                    surfaceControlInputReceiverInfo.mClientToken);
-        } catch (RemoteException e) {
-            Log.e(TAG, "Failed to remove input channel", e);
-            e.rethrowAsRuntimeException();
-        }
+        removeInputChannel(surfaceControlInputReceiverInfo.mClientToken);
 
         surfaceControlInputReceiverInfo.mInputEventReceiver.dispose();
     }
diff --git a/core/java/android/view/WindowManagerImpl.java b/core/java/android/view/WindowManagerImpl.java
index df4fed6..b667427 100644
--- a/core/java/android/view/WindowManagerImpl.java
+++ b/core/java/android/view/WindowManagerImpl.java
@@ -535,7 +535,7 @@
 
     @NonNull
     @Override
-    public InputTransferToken registerBatchedSurfaceControlInputReceiver(int displayId,
+    public InputTransferToken registerBatchedSurfaceControlInputReceiver(
             @NonNull InputTransferToken hostInputTransferToken,
             @NonNull SurfaceControl surfaceControl, @NonNull Choreographer choreographer,
             @NonNull SurfaceControlInputReceiver receiver) {
@@ -543,13 +543,13 @@
         Objects.requireNonNull(surfaceControl);
         Objects.requireNonNull(choreographer);
         Objects.requireNonNull(receiver);
-        return mGlobal.registerBatchedSurfaceControlInputReceiver(displayId, hostInputTransferToken,
+        return mGlobal.registerBatchedSurfaceControlInputReceiver(hostInputTransferToken,
                 surfaceControl, choreographer, receiver);
     }
 
     @NonNull
     @Override
-    public InputTransferToken registerUnbatchedSurfaceControlInputReceiver(int displayId,
+    public InputTransferToken registerUnbatchedSurfaceControlInputReceiver(
             @NonNull InputTransferToken hostInputTransferToken,
             @NonNull SurfaceControl surfaceControl, @NonNull Looper looper,
             @NonNull SurfaceControlInputReceiver receiver) {
@@ -557,7 +557,7 @@
         Objects.requireNonNull(surfaceControl);
         Objects.requireNonNull(looper);
         Objects.requireNonNull(receiver);
-        return mGlobal.registerUnbatchedSurfaceControlInputReceiver(displayId,
+        return mGlobal.registerUnbatchedSurfaceControlInputReceiver(
                 hostInputTransferToken, surfaceControl, looper, receiver);
     }
 
diff --git a/core/java/android/window/InputTransferToken.java b/core/java/android/window/InputTransferToken.java
index c62eee4..5fab48f 100644
--- a/core/java/android/window/InputTransferToken.java
+++ b/core/java/android/window/InputTransferToken.java
@@ -38,9 +38,9 @@
  * {@link SurfaceControlViewHost} or {@link android.view.SurfaceControl} that has an input channel.
  * <p>
  * The {@link android.view.SurfaceControl} needs to have been registered for input via
- * {@link android.view.WindowManager#registerUnbatchedSurfaceControlInputReceiver(int,
+ * {@link android.view.WindowManager#registerUnbatchedSurfaceControlInputReceiver(
  * InputTransferToken, SurfaceControl, Looper, SurfaceControlInputReceiver)} or
- * {@link android.view.WindowManager#registerBatchedSurfaceControlInputReceiver(int,
+ * {@link android.view.WindowManager#registerBatchedSurfaceControlInputReceiver(
  * InputTransferToken, SurfaceControl, Choreographer, SurfaceControlInputReceiver)} and the
  * returned token can be used to call
  * {@link android.view.WindowManager#transferTouchGesture(InputTransferToken, InputTransferToken)}
diff --git a/core/jni/Android.bp b/core/jni/Android.bp
index ac961ee..444288c 100644
--- a/core/jni/Android.bp
+++ b/core/jni/Android.bp
@@ -164,6 +164,7 @@
                 "android_view_Surface.cpp",
                 "android_view_SurfaceControl.cpp",
                 "android_view_SurfaceControlHdrLayerInfoListener.cpp",
+                "android_view_WindowManagerGlobal.cpp",
                 "android_graphics_BLASTBufferQueue.cpp",
                 "android_view_SurfaceSession.cpp",
                 "android_view_TextureView.cpp",
diff --git a/core/jni/AndroidRuntime.cpp b/core/jni/AndroidRuntime.cpp
index 9bbd191..71d041c 100644
--- a/core/jni/AndroidRuntime.cpp
+++ b/core/jni/AndroidRuntime.cpp
@@ -222,6 +222,7 @@
 extern int register_android_tracing_PerfettoDataSourceInstance(JNIEnv* env);
 extern int register_android_tracing_PerfettoProducer(JNIEnv* env);
 extern int register_android_window_InputTransferToken(JNIEnv* env);
+extern int register_android_view_WindowManagerGlobal(JNIEnv* env);
 
 // Namespace for Android Runtime flags applied during boot time.
 static const char* RUNTIME_NATIVE_BOOT_NAMESPACE = "runtime_native_boot";
@@ -1680,6 +1681,7 @@
         REG_JNI(register_android_tracing_PerfettoDataSourceInstance),
         REG_JNI(register_android_tracing_PerfettoProducer),
         REG_JNI(register_android_window_InputTransferToken),
+        REG_JNI(register_android_view_WindowManagerGlobal),
 };
 
 /*
diff --git a/core/jni/android_view_SurfaceControl.cpp b/core/jni/android_view_SurfaceControl.cpp
index 6fec527a..1eab991 100644
--- a/core/jni/android_view_SurfaceControl.cpp
+++ b/core/jni/android_view_SurfaceControl.cpp
@@ -254,6 +254,8 @@
 static struct {
     jclass clazz;
     jfieldID mNativeObject;
+    jfieldID mName;
+    jmethodID ctor;
     jmethodID invokeReleaseCallback;
 } gSurfaceControlClassInfo;
 
@@ -2177,6 +2179,20 @@
     }
 }
 
+jobject android_view_SurfaceControl_getJavaSurfaceControl(JNIEnv* env,
+                                                          const SurfaceControl& surfaceControl) {
+    jobject surfaceControlObj =
+            env->NewObject(gSurfaceControlClassInfo.clazz, gSurfaceControlClassInfo.ctor);
+    env->SetLongField(surfaceControlObj, gSurfaceControlClassInfo.mNativeObject,
+                      reinterpret_cast<jlong>(&surfaceControl));
+    env->SetObjectField(surfaceControlObj, gSurfaceControlClassInfo.mName,
+                        ScopedLocalRef<jobject>(env,
+                                                env->NewStringUTF(surfaceControl.getName().c_str()))
+                                .get());
+    surfaceControl.incStrong((void*)nativeCreate);
+    return surfaceControlObj;
+}
+
 SurfaceComposerClient::Transaction* android_view_SurfaceTransaction_getNativeSurfaceTransaction(
         JNIEnv* env, jobject surfaceTransactionObj) {
     if (!!surfaceTransactionObj &&
@@ -2652,6 +2668,9 @@
     gSurfaceControlClassInfo.clazz = MakeGlobalRefOrDie(env, surfaceControlClazz);
     gSurfaceControlClassInfo.mNativeObject =
             GetFieldIDOrDie(env, gSurfaceControlClassInfo.clazz, "mNativeObject", "J");
+    gSurfaceControlClassInfo.mName =
+            GetFieldIDOrDie(env, gSurfaceControlClassInfo.clazz, "mName", "Ljava/lang/String;");
+    gSurfaceControlClassInfo.ctor = GetMethodIDOrDie(env, surfaceControlClazz, "<init>", "()V");
     gSurfaceControlClassInfo.invokeReleaseCallback =
             GetStaticMethodIDOrDie(env, surfaceControlClazz, "invokeReleaseCallback",
                                    "(Ljava/util/function/Consumer;J)V");
diff --git a/core/jni/android_view_WindowManagerGlobal.cpp b/core/jni/android_view_WindowManagerGlobal.cpp
new file mode 100644
index 0000000..b03ac88
--- /dev/null
+++ b/core/jni/android_view_WindowManagerGlobal.cpp
@@ -0,0 +1,91 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+
+#define LOG_TAG "WindowManagerGlobal"
+
+#include "android_view_WindowManagerGlobal.h"
+
+#include <android_runtime/AndroidRuntime.h>
+#include <android_runtime/android_view_SurfaceControl.h>
+#include <android_runtime/android_window_InputTransferToken.h>
+#include <jni.h>
+#include <nativehelper/ScopedLocalRef.h>
+
+#include "android_util_Binder.h"
+#include "android_view_InputChannel.h"
+#include "jni_wrappers.h"
+
+namespace android {
+
+static struct {
+    jclass clazz;
+    jmethodID createInputChannel;
+    jmethodID removeInputChannel;
+} gWindowManagerGlobal;
+
+std::shared_ptr<InputChannel> createInputChannel(
+        const sp<IBinder>& clientToken, const InputTransferToken& hostInputTransferToken,
+        const SurfaceControl& surfaceControl, const InputTransferToken& clientInputTransferToken) {
+    JNIEnv* env = AndroidRuntime::getJNIEnv();
+    ScopedLocalRef<jobject> hostInputTransferTokenObj(
+            env,
+            android_window_InputTransferToken_getJavaInputTransferToken(env,
+                                                                        hostInputTransferToken));
+    ScopedLocalRef<jobject>
+            surfaceControlObj(env,
+                              android_view_SurfaceControl_getJavaSurfaceControl(env,
+                                                                                surfaceControl));
+    jobject clientTokenObj = javaObjectForIBinder(env, clientToken);
+    ScopedLocalRef<jobject> clientInputTransferTokenObj(
+            env,
+            android_window_InputTransferToken_getJavaInputTransferToken(env,
+                                                                        clientInputTransferToken));
+    ScopedLocalRef<jobject>
+            inputChannelObj(env,
+                            env->CallStaticObjectMethod(gWindowManagerGlobal.clazz,
+                                                        gWindowManagerGlobal.createInputChannel,
+                                                        clientTokenObj,
+                                                        hostInputTransferTokenObj.get(),
+                                                        surfaceControlObj.get(),
+                                                        clientInputTransferTokenObj.get()));
+
+    return android_view_InputChannel_getInputChannel(env, inputChannelObj.get());
+}
+
+void removeInputChannel(const sp<IBinder>& clientToken) {
+    JNIEnv* env = AndroidRuntime::getJNIEnv();
+
+    jobject clientTokenObj(javaObjectForIBinder(env, clientToken));
+    env->CallStaticObjectMethod(gWindowManagerGlobal.clazz, gWindowManagerGlobal.removeInputChannel,
+                                clientTokenObj);
+}
+
+int register_android_view_WindowManagerGlobal(JNIEnv* env) {
+    jclass windowManagerGlobalClass = FindClassOrDie(env, "android/view/WindowManagerGlobal");
+    gWindowManagerGlobal.clazz = MakeGlobalRefOrDie(env, windowManagerGlobalClass);
+    gWindowManagerGlobal.createInputChannel =
+            GetStaticMethodIDOrDie(env, windowManagerGlobalClass, "createInputChannel",
+                                   "(Landroid/os/IBinder;Landroid/window/"
+                                   "InputTransferToken;Landroid/view/SurfaceControl;Landroid/"
+                                   "window/InputTransferToken;)Landroid/view/InputChannel;");
+    gWindowManagerGlobal.removeInputChannel =
+            GetStaticMethodIDOrDie(env, windowManagerGlobalClass, "removeInputChannel",
+                                   "(Landroid/os/IBinder;)V");
+
+    return NO_ERROR;
+}
+
+} // namespace android
\ No newline at end of file
diff --git a/core/jni/android_view_WindowManagerGlobal.h b/core/jni/android_view_WindowManagerGlobal.h
new file mode 100644
index 0000000..fcdeaa1
--- /dev/null
+++ b/core/jni/android_view_WindowManagerGlobal.h
@@ -0,0 +1,29 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+
+#include <binder/IBinder.h>
+#include <gui/InputTransferToken.h>
+#include <gui/SurfaceControl.h>
+#include <input/InputTransport.h>
+
+namespace android {
+extern std::shared_ptr<InputChannel> createInputChannel(
+        const sp<IBinder>& clientToken, const InputTransferToken& hostInputTransferToken,
+        const SurfaceControl& surfaceControl,
+        const InputTransferToken& clientInputTransferTokenObj);
+extern void removeInputChannel(const sp<IBinder>& clientToken);
+
+} // namespace android
\ No newline at end of file
diff --git a/core/jni/android_window_InputTransferToken.cpp b/core/jni/android_window_InputTransferToken.cpp
index 60568e30..8fb668d 100644
--- a/core/jni/android_window_InputTransferToken.cpp
+++ b/core/jni/android_window_InputTransferToken.cpp
@@ -83,14 +83,10 @@
 }
 
 jobject android_window_InputTransferToken_getJavaInputTransferToken(
-        JNIEnv* env, const InputTransferToken* inputTransferToken) {
-    if (inputTransferToken == nullptr || env == nullptr) {
-        return nullptr;
-    }
-
-    inputTransferToken->incStrong((void*)nativeCreate);
+        JNIEnv* env, const InputTransferToken& inputTransferToken) {
+    inputTransferToken.incStrong((void*)nativeCreate);
     return env->NewObject(gInputTransferTokenClassInfo.clazz, gInputTransferTokenClassInfo.ctor,
-                          reinterpret_cast<jlong>(inputTransferToken));
+                          reinterpret_cast<jlong>(&inputTransferToken));
 }
 
 static void release(InputTransferToken* inputTransferToken) {
diff --git a/core/jni/include/android_runtime/android_view_SurfaceControl.h b/core/jni/include/android_runtime/android_view_SurfaceControl.h
index 10a7549..543deb8 100644
--- a/core/jni/include/android_runtime/android_view_SurfaceControl.h
+++ b/core/jni/include/android_runtime/android_view_SurfaceControl.h
@@ -28,6 +28,9 @@
 extern SurfaceControl* android_view_SurfaceControl_getNativeSurfaceControl(
         JNIEnv* env, jobject surfaceControlObj);
 
+extern jobject android_view_SurfaceControl_getJavaSurfaceControl(
+        JNIEnv* env, const SurfaceControl& surfaceControl);
+
 /* Gets the underlying native SurfaceControl for a java SurfaceControl. */
 extern SurfaceComposerClient::Transaction*
 android_view_SurfaceTransaction_getNativeSurfaceTransaction(JNIEnv* env,
diff --git a/core/jni/include/android_runtime/android_window_InputTransferToken.h b/core/jni/include/android_runtime/android_window_InputTransferToken.h
index 75dbe37..5ac48f3 100644
--- a/core/jni/include/android_runtime/android_window_InputTransferToken.h
+++ b/core/jni/include/android_runtime/android_window_InputTransferToken.h
@@ -26,7 +26,7 @@
         JNIEnv* env, jobject inputTransferTokenObj);
 
 extern jobject android_window_InputTransferToken_getJavaInputTransferToken(
-        JNIEnv* env, const InputTransferToken* inputTransferToken);
+        JNIEnv* env, const InputTransferToken& inputTransferToken);
 
 } // namespace android
 
diff --git a/native/android/Android.bp b/native/android/Android.bp
index 4812685..db3a67a 100644
--- a/native/android/Android.bp
+++ b/native/android/Android.bp
@@ -54,6 +54,7 @@
     srcs: [
         "activity_manager.cpp",
         "asset_manager.cpp",
+        "surface_control_input_receiver.cpp",
         "choreographer.cpp",
         "configuration.cpp",
         "hardware_buffer_jni.cpp",
diff --git a/native/android/input_transfer_token.cpp b/native/android/input_transfer_token.cpp
index 501e1d3..2e74aa3 100644
--- a/native/android/input_transfer_token.cpp
+++ b/native/android/input_transfer_token.cpp
@@ -25,7 +25,7 @@
 #define CHECK_NOT_NULL(name) \
     LOG_ALWAYS_FATAL_IF(name == nullptr, "nullptr passed as " #name " argument");
 
-void InputTransferToken_acquire(InputTransferToken* inputTransferToken) {
+extern void InputTransferToken_acquire(InputTransferToken* inputTransferToken) {
     // incStrong/decStrong token must be the same, doesn't matter what it is
     inputTransferToken->incStrong((void*)InputTransferToken_acquire);
 }
@@ -52,7 +52,7 @@
     CHECK_NOT_NULL(aInputTransferToken);
     const InputTransferToken* inputTransferToken =
             reinterpret_cast<const InputTransferToken*>(aInputTransferToken);
-    return android_window_InputTransferToken_getJavaInputTransferToken(env, inputTransferToken);
+    return android_window_InputTransferToken_getJavaInputTransferToken(env, *inputTransferToken);
 }
 
 void AInputTransferToken_release(AInputTransferToken* aInputTransferToken) {
diff --git a/native/android/libandroid.map.txt b/native/android/libandroid.map.txt
index b2925bf..1c203e3 100644
--- a/native/android/libandroid.map.txt
+++ b/native/android/libandroid.map.txt
@@ -98,6 +98,14 @@
     AInputQueue_getEvent;
     AInputQueue_hasEvents;
     AInputQueue_preDispatchEvent;
+    AInputReceiver_createBatchedInputReceiver; # introduced=35
+    AInputReceiver_createUnbatchedInputReceiver; # introduced=35
+    AInputReceiver_release; # introduced=35
+    AInputReceiver_getInputTransferToken; # introduced=35
+    AInputReceiverCallbacks_create; # introduced=35
+    AInputReceiverCallbacks_release; # introduced=35
+    AInputReceiverCallbacks_setKeyEventCallback; # introduced=35
+    AInputReceiverCallbacks_setMotionEventCallback; # introduced=35
     AInputTransferToken_fromJava; # introduced=35
     AInputTransferToken_release; # introduced=35
     AInputTransferToken_toJava; # introduced=35
diff --git a/native/android/surface_control_input_receiver.cpp b/native/android/surface_control_input_receiver.cpp
new file mode 100644
index 0000000..da0defd
--- /dev/null
+++ b/native/android/surface_control_input_receiver.cpp
@@ -0,0 +1,213 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+
+#include <android/choreographer.h>
+#include <android/surface_control_input_receiver.h>
+#include <binder/Binder.h>
+#include <gui/Choreographer.h>
+#include <gui/InputTransferToken.h>
+#include <input/Input.h>
+#include <input/InputConsumerNoResampling.h>
+
+#include "android_view_WindowManagerGlobal.h"
+
+using namespace android;
+
+extern void InputTransferToken_acquire(InputTransferToken* inputTransferToken);
+
+struct AInputReceiverCallbacks {
+    AInputReceiverCallbacks(void* context) : context(context) {}
+    void* context;
+    AInputReceiver_onMotionEvent onMotionEvent = nullptr;
+    AInputReceiver_onKeyEvent onKeyEvent = nullptr;
+};
+
+class InputReceiver : public InputConsumerCallbacks {
+public:
+    InputReceiver(const sp<Looper>& looper, const std::shared_ptr<InputChannel>& inputChannel,
+                  const sp<IBinder>& clientToken, const sp<InputTransferToken>& inputTransferToken,
+                  AInputReceiverCallbacks* callbacks)
+          : mCallbacks(callbacks),
+            mInputConsumer(inputChannel, looper, *this),
+            mClientToken(clientToken),
+            mInputTransferToken(inputTransferToken) {}
+
+    ~InputReceiver() {
+        remove();
+    }
+
+    void onKeyEvent(std::unique_ptr<KeyEvent> event, uint32_t seq) override {
+        if (mCallbacks->onKeyEvent != nullptr) {
+            const bool handled = mCallbacks->onKeyEvent(mCallbacks->context,
+                                                        static_cast<AInputEvent*>(event.release()));
+            mInputConsumer.finishInputEvent(seq, handled);
+        }
+    }
+
+    void onMotionEvent(std::unique_ptr<MotionEvent> event, uint32_t seq) override {
+        if (mCallbacks->onMotionEvent != nullptr) {
+            const bool handled =
+                    mCallbacks->onMotionEvent(mCallbacks->context,
+                                              static_cast<AInputEvent*>(event.release()));
+            mInputConsumer.finishInputEvent(seq, handled);
+        }
+    }
+
+    void onFocusEvent(std::unique_ptr<FocusEvent>, uint32_t seq) override {
+        mInputConsumer.finishInputEvent(seq, false);
+    }
+    void onCaptureEvent(std::unique_ptr<CaptureEvent>, uint32_t seq) override {
+        mInputConsumer.finishInputEvent(seq, false);
+    }
+    void onDragEvent(std::unique_ptr<DragEvent>, uint32_t seq) override {
+        mInputConsumer.finishInputEvent(seq, false);
+    }
+    void onTouchModeEvent(std::unique_ptr<TouchModeEvent>, uint32_t seq) override {
+        mInputConsumer.finishInputEvent(seq, false);
+    }
+
+    virtual void onBatchedInputEventPending(int32_t) override {
+        mInputConsumer.consumeBatchedInputEvents(std::nullopt);
+    }
+
+    const AInputTransferToken* getInputTransferToken() {
+        InputTransferToken_acquire(mInputTransferToken.get());
+        return reinterpret_cast<const AInputTransferToken*>(mInputTransferToken.get());
+    }
+
+    void remove() {
+        removeInputChannel(mClientToken);
+    }
+
+    AInputReceiverCallbacks* mCallbacks;
+
+protected:
+    InputConsumerNoResampling mInputConsumer;
+
+private:
+    const sp<IBinder> mClientToken;
+    const sp<InputTransferToken> mInputTransferToken;
+};
+
+class BatchedInputReceiver : public InputReceiver {
+public:
+    BatchedInputReceiver(Choreographer& choreographer,
+                         const std::shared_ptr<InputChannel>& inputChannel,
+                         const sp<IBinder>& clientToken,
+                         const sp<InputTransferToken>& inputTransferToken,
+                         AInputReceiverCallbacks* callbacks)
+          : InputReceiver(choreographer.getLooper(), inputChannel, clientToken, inputTransferToken,
+                          callbacks),
+            mChoreographer(choreographer) {}
+
+    static void vsyncCallback(const AChoreographerFrameCallbackData* callbackData, void* data) {
+        BatchedInputReceiver* receiver = static_cast<BatchedInputReceiver*>(data);
+        receiver->onVsyncCallback(callbackData);
+    }
+
+    void onVsyncCallback(const AChoreographerFrameCallbackData* callbackData) {
+        int64_t frameTimeNanos = AChoreographerFrameCallbackData_getFrameTimeNanos(callbackData);
+        mInputConsumer.consumeBatchedInputEvents(frameTimeNanos);
+        mBatchedInputScheduled = false;
+    }
+
+    void onBatchedInputEventPending(int32_t) override {
+        scheduleBatchedInput();
+    }
+
+private:
+    Choreographer& mChoreographer;
+    bool mBatchedInputScheduled = false;
+
+    void scheduleBatchedInput() {
+        if (!mBatchedInputScheduled) {
+            mBatchedInputScheduled = true;
+            mChoreographer.postFrameCallbackDelayed(nullptr, nullptr, vsyncCallback, this, 0,
+                                                    CallbackType::CALLBACK_INPUT);
+        }
+    }
+};
+
+static inline AInputReceiver* InputReceiver_to_AInputReceiver(InputReceiver* inputReceiver) {
+    return reinterpret_cast<AInputReceiver*>(inputReceiver);
+}
+
+static inline InputReceiver* AInputReceiver_to_InputReceiver(AInputReceiver* aInputReceiver) {
+    return reinterpret_cast<InputReceiver*>(aInputReceiver);
+}
+
+AInputReceiver* AInputReceiver_createBatchedInputReceiver(AChoreographer* aChoreographer,
+                                                          const AInputTransferToken* hostToken,
+                                                          const ASurfaceControl* aSurfaceControl,
+                                                          AInputReceiverCallbacks* callbacks) {
+    // create input channel here through WMS
+    sp<IBinder> clientToken = sp<BBinder>::make();
+    sp<InputTransferToken> clientInputTransferToken = sp<InputTransferToken>::make();
+
+    std::shared_ptr<InputChannel> inputChannel =
+            createInputChannel(clientToken, reinterpret_cast<const InputTransferToken&>(*hostToken),
+                               reinterpret_cast<const SurfaceControl&>(*aSurfaceControl),
+                               *clientInputTransferToken);
+    return InputReceiver_to_AInputReceiver(
+            new BatchedInputReceiver(reinterpret_cast<Choreographer&>(*aChoreographer),
+                                     inputChannel, clientToken, clientInputTransferToken,
+                                     callbacks));
+}
+
+AInputReceiver* AInputReceiver_createUnbatchedInputReceiver(ALooper* aLooper,
+                                                            const AInputTransferToken* hostToken,
+                                                            const ASurfaceControl* aSurfaceControl,
+                                                            AInputReceiverCallbacks* callbacks) {
+    // create input channel here through WMS
+    sp<IBinder> clientToken = sp<BBinder>::make();
+    sp<InputTransferToken> clientInputTransferToken = sp<InputTransferToken>::make();
+
+    std::shared_ptr<InputChannel> inputChannel =
+            createInputChannel(clientToken, reinterpret_cast<const InputTransferToken&>(*hostToken),
+                               reinterpret_cast<const SurfaceControl&>(*aSurfaceControl),
+                               *clientInputTransferToken);
+    return InputReceiver_to_AInputReceiver(new InputReceiver(reinterpret_cast<Looper*>(aLooper),
+                                                             inputChannel, clientToken,
+                                                             clientInputTransferToken, callbacks));
+}
+
+const AInputTransferToken* AInputReceiver_getInputTransferToken(AInputReceiver* aInputReceiver) {
+    return AInputReceiver_to_InputReceiver(aInputReceiver)->getInputTransferToken();
+}
+
+void AInputReceiver_release(AInputReceiver* aInputReceiver) {
+    InputReceiver* inputReceiver = AInputReceiver_to_InputReceiver(aInputReceiver);
+    inputReceiver->remove();
+    delete inputReceiver;
+}
+
+void AInputReceiverCallbacks_setMotionEventCallback(AInputReceiverCallbacks* callbacks,
+                                                    AInputReceiver_onMotionEvent onMotionEvent) {
+    callbacks->onMotionEvent = onMotionEvent;
+}
+
+void AInputReceiverCallbacks_setKeyEventCallback(AInputReceiverCallbacks* callbacks,
+                                                 AInputReceiver_onKeyEvent onKeyEvent) {
+    callbacks->onKeyEvent = onKeyEvent;
+}
+
+AInputReceiverCallbacks* AInputReceiverCallbacks_create(void* context) {
+    return new AInputReceiverCallbacks(context);
+}
+
+void AInputReceiverCallbacks_release(AInputReceiverCallbacks* callbacks) {
+    delete callbacks;
+}
\ No newline at end of file
diff --git a/tests/SurfaceControlViewHostTest/src/com/android/test/viewembed/EmbeddedWindowService.java b/tests/SurfaceControlViewHostTest/src/com/android/test/viewembed/EmbeddedWindowService.java
index 0fb4f90..56fb30c 100644
--- a/tests/SurfaceControlViewHostTest/src/com/android/test/viewembed/EmbeddedWindowService.java
+++ b/tests/SurfaceControlViewHostTest/src/com/android/test/viewembed/EmbeddedWindowService.java
@@ -135,7 +135,7 @@
                 c.drawText("Remote", 250, 250, paint);
                 surface.unlockCanvasAndPost(c);
                 WindowManager wm = getSystemService(WindowManager.class);
-                wm.registerBatchedSurfaceControlInputReceiver(displayId, inputTransferToken,
+                wm.registerBatchedSurfaceControlInputReceiver(inputTransferToken,
                         mSurfaceControl,
                         Choreographer.getInstance(), event -> {
                             Log.d(TAG, "onInputEvent-remote " + event);
diff --git a/tests/SurfaceControlViewHostTest/src/com/android/test/viewembed/SurfaceInputTestActivity.java b/tests/SurfaceControlViewHostTest/src/com/android/test/viewembed/SurfaceInputTestActivity.java
index e700bc2..ac7dc9e 100644
--- a/tests/SurfaceControlViewHostTest/src/com/android/test/viewembed/SurfaceInputTestActivity.java
+++ b/tests/SurfaceControlViewHostTest/src/com/android/test/viewembed/SurfaceInputTestActivity.java
@@ -138,7 +138,7 @@
         c.drawText("Local SC", 0, 0, paint);
         surface.unlockCanvasAndPost(c);
         WindowManager wm = getSystemService(WindowManager.class);
-        wm.registerBatchedSurfaceControlInputReceiver(getDisplayId(),
+        wm.registerBatchedSurfaceControlInputReceiver(
                 attachedSurfaceControl.getInputTransferToken(), mLocalSurfaceControl,
                 Choreographer.getInstance(), event -> {
                     Log.d(TAG, "onInputEvent-sc " + event);
@@ -159,7 +159,7 @@
             holder.unlockCanvasAndPost(c);
 
             WindowManager wm = getSystemService(WindowManager.class);
-            wm.registerBatchedSurfaceControlInputReceiver(getDisplayId(),
+            wm.registerBatchedSurfaceControlInputReceiver(
                     mLocalSurfaceView.getRootSurfaceControl().getInputTransferToken(),
                     mLocalSurfaceView.getSurfaceControl(),
                     Choreographer.getInstance(), event -> {