Disallow click to partial screenshot right after screenshot is taken

This change prevents triggering partial screenshot when trying to take
full screenshot of protected window.

Change-Id: I42a67ed1ee34edffccc6dacd17ba1dffc8d194b3

2 files changed, 24 insertions, 4 deletions

diff --git a/services/core/java/com/android/server/policy/PhoneWindowManager.java b/services/core/java/com/android/server/policy/PhoneWindowManager.java
index 2e58c98da57e..817758017c3e 100644
--- a/services/core/java/com/android/server/policy/PhoneWindowManager.java
+++ b/services/core/java/com/android/server/policy/PhoneWindowManager.java
@@ -539,6 +539,7 @@ public class PhoneWindowManager implements WindowManagerPolicy {
 
     // Click volume down + power for partial screenshot
     boolean mClickPartialScreenshot;
+    boolean mClickPartialScreenshotAllowed = true;
 
     private boolean mPendingKeyguardOccluded;
     private boolean mKeyguardOccludedChanged;
@@ -1639,7 +1640,8 @@ public class PhoneWindowManager implements WindowManagerPolicy {
 
         @Override
         public void run() {
-            mDefaultDisplayPolicy.takeScreenshot(mScreenshotType, mScreenshotSource);
+            mDefaultDisplayPolicy.takeScreenshot(mScreenshotType, mScreenshotSource,
+                    uri -> { mClickPartialScreenshotAllowed = false; });
         }
     }
 
@@ -4322,10 +4324,13 @@ public class PhoneWindowManager implements WindowManagerPolicy {
                         cancelPendingScreenshotChordAction();
                         cancelPendingAccessibilityShortcutAction();
 
-                        if (mClickPartialScreenshot && mScreenshotChordVolumeDownKeyConsumed) {
+                        if (mClickPartialScreenshot && mClickPartialScreenshotAllowed &&
+                                mScreenshotChordVolumeDownKeyConsumed) {
                             mScreenshotRunnable.setScreenshotType(TAKE_SCREENSHOT_SELECTED_REGION);
                             mHandler.post(mScreenshotRunnable);
                         }
+
+                        mClickPartialScreenshotAllowed = true;
                     }
                 } else if (keyCode == KeyEvent.KEYCODE_VOLUME_UP) {
                     if (down) {
diff --git a/services/core/java/com/android/server/wm/DisplayPolicy.java b/services/core/java/com/android/server/wm/DisplayPolicy.java
index 97e4a9595da3..fb2906342bca 100644
--- a/services/core/java/com/android/server/wm/DisplayPolicy.java
+++ b/services/core/java/com/android/server/wm/DisplayPolicy.java
@@ -151,6 +151,7 @@ import android.graphics.Rect;
 import android.graphics.Region;
 import android.hardware.input.InputManager;
 import android.hardware.power.V1_0.PowerHint;
+import android.net.Uri;
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
@@ -4121,16 +4122,30 @@ public class DisplayPolicy {
      *                       {@link WindowManager#TAKE_SCREENSHOT_FULLSCREEN} or
      *                       {@link WindowManager#TAKE_SCREENSHOT_SELECTED_REGION}
      * @param source Where the screenshot originated from (see WindowManager.ScreenshotSource)
+     * @param completionConsumer Consumes `false` if a screenshot was not taken, and `true` if the
+     *                       screenshot was taken.
      */
-    public void takeScreenshot(int screenshotType, int source) {
+    public void takeScreenshot(int screenshotType, int source, Consumer<Uri> completionConsumer) {
         if (mScreenshotHelper != null) {
             mScreenshotHelper.takeScreenshot(screenshotType,
                     getStatusBar() != null && getStatusBar().isVisibleLw(),
                     getNavigationBar() != null && getNavigationBar().isVisibleLw(),
-                    source, mHandler, null /* completionConsumer */);
+                    source, mHandler, completionConsumer);
         }
     }
 
+    /**
+     * Request a screenshot be taken.
+     *
+     * @param screenshotType The type of screenshot, for example either
+     *                       {@link WindowManager#TAKE_SCREENSHOT_FULLSCREEN} or
+     *                       {@link WindowManager#TAKE_SCREENSHOT_SELECTED_REGION}
+     * @param source Where the screenshot originated from (see WindowManager.ScreenshotSource)
+     */
+    public void takeScreenshot(int screenshotType, int source) {
+        takeScreenshot(screenshotType, source, null /* completionConsumer */);
+    }
+
     RefreshRatePolicy getRefreshRatePolicy() {
         return mRefreshRatePolicy;
     }