Set the IME to secure if the target is secure

If the target window for the IME is secure, set the IME to be secure.
This ensures that if the window can't be screenshot, the IME shouldn't
be able to be screenshot at that time. Once the IME gets a new target,
it can be screenshot if the new target isn't secure.

Test: Secure app with IME no longer screenrecords IME
Bug: 118188362
Change-Id: I5d64e4a2262734eb8b0cba09481479231d166c4a

2 files changed, 13 insertions, 1 deletions

diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index a9d33dc29467..c9af62bcfd61 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -3750,6 +3750,10 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
     @VisibleForTesting
     void setImeInputTarget(WindowState target) {
         mImeInputTarget = target;
+        boolean canScreenshot = mImeInputTarget == null || !mImeInputTarget.isSecureLocked();
+        if (mImeWindowsContainer.setCanScreenshot(canScreenshot)) {
+            mWmService.requestTraversal();
+        }
     }
 
     @VisibleForTesting
@@ -3867,7 +3871,7 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
     void updateImeInputAndControlTarget(WindowState target) {
         if (mImeInputTarget != target) {
             ProtoLog.i(WM_DEBUG_IME, "setInputMethodInputTarget %s", target);
-            mImeInputTarget = target;
+            setImeInputTarget(target);
             updateImeControlTarget();
         }
     }
diff --git a/services/core/java/com/android/server/wm/WindowContainer.java b/services/core/java/com/android/server/wm/WindowContainer.java
index a5843d496add..070a7252a3bf 100644
--- a/services/core/java/com/android/server/wm/WindowContainer.java
+++ b/services/core/java/com/android/server/wm/WindowContainer.java
@@ -3321,4 +3321,12 @@ class WindowContainer<E extends WindowContainer> extends ConfigurationContainer<
     @WindowManager.LayoutParams.WindowType int getWindowType() {
         return INVALID_WINDOW_TYPE;
     }
+
+    boolean setCanScreenshot(boolean canScreenshot) {
+        if (mSurfaceControl == null) {
+            return false;
+        }
+        getPendingTransaction().setSecure(mSurfaceControl, !canScreenshot);
+        return true;
+    }
 }