From 1934e139d942620ca786a85ef63bd677cab6d4be Mon Sep 17 00:00:00 2001
From: Shan Huang <shanh@google.com>
Date: Wed, 1 Jun 2022 11:51:36 -0700
Subject: [PATCH] Fix IME callback not unregistered when IME is up during an
 IMM client switch.

This can happen when one navigates from an activity showing IME
to another activity requesting to show IME. The following sequence would
happen:
- IMM of the current activity receives an UNBIND message with reason CLIENT_SWITCH
- A new IMM is created for the new activity requesting input and
  requests to bind.
- The new IMM requests to start input, triggering
  IMS#dispatchStartInputWithToken with a *new* ImeOnBackInvokedDispatcher instance than the one previously kept by IMS.

Previously, the ime dispatcher instance was updated before the existing
back callback was unregistered and caused the unregister call to fail.
The fix is to update the dispatcher after the other logic in
IMS#dispatcherStartInputWithToken (specifically, to after startInput(),
which calls doFinishInput() if input is already started).

Bug: 234416098
Test: atest CtsInputMethodTestCases

Change-Id: Ic7a5daab8709b808343388b095b5b6ad40adc777
---
 .../InputMethodService.java                   | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/core/java/android/inputmethodservice/InputMethodService.java b/core/java/android/inputmethodservice/InputMethodService.java
index 21ecf8b0888c1..334a659d7f239 100644
--- a/core/java/android/inputmethodservice/InputMethodService.java
+++ b/core/java/android/inputmethodservice/InputMethodService.java
@@ -806,11 +806,6 @@ public class InputMethodService extends AbstractInputMethodService {
                 @NonNull EditorInfo editorInfo, boolean restarting,
                 @NonNull IBinder startInputToken, @InputMethodNavButtonFlags int navButtonFlags,
                 @NonNull ImeOnBackInvokedDispatcher imeDispatcher) {
-            mImeDispatcher = imeDispatcher;
-            if (mWindow != null) {
-                mWindow.getOnBackInvokedDispatcher().setImeOnBackInvokedDispatcher(
-                        imeDispatcher);
-            }
             mPrivOps.reportStartInputAsync(startInputToken);
             mNavigationBarController.onNavButtonFlagsChanged(navButtonFlags);
             if (restarting) {
@@ -818,6 +813,15 @@ public class InputMethodService extends AbstractInputMethodService {
             } else {
                 startInput(inputConnection, editorInfo);
             }
+            // Update the IME dispatcher last, so that the previously registered back callback
+            // (if any) can be unregistered using the old dispatcher if {@link #doFinishInput()}
+            // is called from {@link #startInput(InputConnection, EditorInfo)} or
+            // {@link #restartInput(InputConnection, EditorInfo)}.
+            mImeDispatcher = imeDispatcher;
+            if (mWindow != null) {
+                mWindow.getOnBackInvokedDispatcher().setImeOnBackInvokedDispatcher(
+                        imeDispatcher);
+            }
         }
 
         /**
@@ -3860,6 +3864,11 @@ public class InputMethodService extends AbstractInputMethodService {
     };
 
     private void compatHandleBack() {
+        if (!mDecorViewVisible) {
+            Log.e(TAG, "Back callback invoked on a hidden IME. Removing the callback...");
+            unregisterCompatOnBackInvokedCallback();
+            return;
+        }
         final KeyEvent downEvent = createBackKeyEvent(
                 KeyEvent.ACTION_DOWN, false /* isTracking */);
         onKeyDown(KeyEvent.KEYCODE_BACK, downEvent);