Merge "Update IME SECURE flag if Window updates SECURE flag" into tm-dev

services/core/java/com/android/server/wm/DisplayContent.java

@@ -4048,12 +4048,20 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
     @VisibleForTesting
     void setImeInputTarget(InputTarget target) {
         mImeInputTarget = target;
-        boolean canScreenshot = mImeInputTarget == null || mImeInputTarget.canScreenshotIme();
+        if (refreshImeSecureFlag(getPendingTransaction())) {
-        if (mImeWindowsContainer.setCanScreenshot(canScreenshot)) {
             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;

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;
     }
 

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");

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 {