From 90d0e50b6f0dd85fb3feb6d14cfc59c14d39ec28 Mon Sep 17 00:00:00 2001
From: chaviw <chaviw@google.com>
Date: Tue, 5 Apr 2022 11:59:00 -0500
Subject: [PATCH] Update IME SECURE flag if Window updates SECURE flag

If the IME target is set before the IME target updates its SECURE flag,
the IME  may not get the SECURE flag. This is because we only update the
IME flag when the new target is being set.

This change ensures we also update the IME secure flag if needed when a
Window updates its SECURE flag since it could already be the target.

Fixes: 215005011
Test: testImeSecureFlagGetUpdatedAfterImeInputTarget
Test: Add SECURE flag after IME target is set.
Change-Id: Ie59042c8efa2527208e4020107b7aa15c885755f
---
 .../java/com/android/server/wm/DisplayContent.java | 12 ++++++++++--
 .../com/android/server/wm/WindowContainer.java     |  4 ++--
 .../android/server/wm/WindowSurfaceController.java |  5 +++++
 .../com/android/server/wm/DisplayContentTests.java | 14 ++++++++++++++
 4 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 81560d4f8676e..714da3d6e3536 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -4047,12 +4047,20 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
     @VisibleForTesting
     void setImeInputTarget(InputTarget target) {
         mImeInputTarget = target;
-        boolean canScreenshot = mImeInputTarget == null || mImeInputTarget.canScreenshotIme();
-        if (mImeWindowsContainer.setCanScreenshot(canScreenshot)) {
+        if (refreshImeSecureFlag(getPendingTransaction())) {
             mWmService.requestTraversal();
         }
     }
 
+    /**
+     * Re-check the IME target's SECURE flag since it's possible to have changed after the target
+     * was set.
+     */
+    boolean refreshImeSecureFlag(Transaction t) {
+        boolean canScreenshot = mImeInputTarget == null || mImeInputTarget.canScreenshotIme();
+        return mImeWindowsContainer.setCanScreenshot(t, canScreenshot);
+    }
+
     @VisibleForTesting
     void setImeControlTarget(InsetsControlTarget target) {
         mImeControlTarget = target;
diff --git a/services/core/java/com/android/server/wm/WindowContainer.java b/services/core/java/com/android/server/wm/WindowContainer.java
index 214524c2f42cc..2900a5df61987 100644
--- a/services/core/java/com/android/server/wm/WindowContainer.java
+++ b/services/core/java/com/android/server/wm/WindowContainer.java
@@ -3782,11 +3782,11 @@ class WindowContainer<E extends WindowContainer> extends ConfigurationContainer<
         return INVALID_WINDOW_TYPE;
     }
 
-    boolean setCanScreenshot(boolean canScreenshot) {
+    boolean setCanScreenshot(Transaction t, boolean canScreenshot) {
         if (mSurfaceControl == null) {
             return false;
         }
-        getPendingTransaction().setSecure(mSurfaceControl, !canScreenshot);
+        t.setSecure(mSurfaceControl, !canScreenshot);
         return true;
     }
 
diff --git a/services/core/java/com/android/server/wm/WindowSurfaceController.java b/services/core/java/com/android/server/wm/WindowSurfaceController.java
index c2414e72abbd2..5f43800bd9d55 100644
--- a/services/core/java/com/android/server/wm/WindowSurfaceController.java
+++ b/services/core/java/com/android/server/wm/WindowSurfaceController.java
@@ -217,6 +217,11 @@ class WindowSurfaceController {
         mService.openSurfaceTransaction();
         try {
             getGlobalTransaction().setSecure(mSurfaceControl, isSecure);
+
+            final DisplayContent dc = mAnimator.mWin.mDisplayContent;
+            if (dc != null) {
+                dc.refreshImeSecureFlag(getGlobalTransaction());
+            }
         } finally {
             mService.closeSurfaceTransaction("setSecure");
             if (SHOW_LIGHT_TRANSACTIONS) Slog.i(TAG, "<<< CLOSE TRANSACTION setSecureLocked");
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
index 605726c0e8049..118f159bee7bc 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
@@ -1201,6 +1201,20 @@ public class DisplayContentTests extends WindowTestsBase {
                 dc.computeImeControlTarget());
     }
 
+    @UseTestDisplay(addWindows = W_INPUT_METHOD)
+    @Test
+    public void testImeSecureFlagGetUpdatedAfterImeInputTarget() {
+        // Verify IME window can get up-to-date secure flag update when the IME input target
+        // set before setCanScreenshot called.
+        final WindowState app = createWindow(null, TYPE_APPLICATION, "app");
+        SurfaceControl.Transaction t = mDisplayContent.mInputMethodWindow.getPendingTransaction();
+        spyOn(t);
+        mDisplayContent.setImeInputTarget(app);
+        mDisplayContent.mInputMethodWindow.setCanScreenshot(t, false /* canScreenshot */);
+
+        verify(t).setSecure(eq(mDisplayContent.mInputMethodWindow.mSurfaceControl), eq(true));
+    }
+
     @UseTestDisplay(addWindows = W_ACTIVITY)
     @Test
     public void testComputeImeControlTarget_notMatchParentBounds() throws Exception {