Fix wallpaper screenshot

Wallpaper screenshot was broken since it would just screenshot the
entire screen. Updated wallpaper screenshot code to use the new
captureLayers API so the wallpaper layer can be specified for the
screenshot.

Change-Id: I594870583ddc2fb29c7eeafe003f20e4ee392a3a
Fixes: 69562019
Test: testWallpaperScreenshot
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 59bece0..7b5e8b8 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -137,7 +137,6 @@
 import android.os.Trace;
 import android.util.ArraySet;
 import android.util.DisplayMetrics;
-import android.util.MutableBoolean;
 import android.util.Slog;
 import android.util.proto.ProtoOutputStream;
 import android.view.Display;
@@ -156,7 +155,6 @@
 import com.android.server.policy.WindowManagerPolicy;
 import com.android.server.wm.utils.RotationCache;
 
-import java.io.FileDescriptor;
 import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.Comparator;
@@ -2960,83 +2958,55 @@
      * In portrait mode, it grabs the full screenshot.
      *
      * @param config of the output bitmap
-     * @param wallpaperOnly true if only the wallpaper layer should be included in the screenshot
      */
-    Bitmap screenshotDisplay(Bitmap.Config config, boolean wallpaperOnly) {
-        synchronized (mService.mWindowMap) {
-            if (!mService.mPolicy.isScreenOn()) {
-                if (DEBUG_SCREENSHOT) {
-                    Slog.i(TAG_WM, "Attempted to take screenshot while display was off.");
-                }
-                return null;
+    Bitmap screenshotDisplayLocked(Bitmap.Config config) {
+        if (!mService.mPolicy.isScreenOn()) {
+            if (DEBUG_SCREENSHOT) {
+                Slog.i(TAG_WM, "Attempted to take screenshot while display was off.");
             }
-
-            if (wallpaperOnly && !shouldScreenshotWallpaper()) {
-                return null;
-            }
-
-            int dw = mDisplayInfo.logicalWidth;
-            int dh = mDisplayInfo.logicalHeight;
-
-            if (dw <= 0 || dh <= 0) {
-                return null;
-            }
-
-            final Rect frame = new Rect(0, 0, dw, dh);
-
-            // The screenshot API does not apply the current screen rotation.
-            int rot = mDisplay.getRotation();
-
-            if (rot == ROTATION_90 || rot == ROTATION_270) {
-                rot = (rot == ROTATION_90) ? ROTATION_270 : ROTATION_90;
-            }
-
-            // SurfaceFlinger is not aware of orientation, so convert our logical
-            // crop to SurfaceFlinger's portrait orientation.
-            convertCropForSurfaceFlinger(frame, rot, dw, dh);
-
-            final ScreenRotationAnimation screenRotationAnimation =
-                    mService.mAnimator.getScreenRotationAnimationLocked(DEFAULT_DISPLAY);
-            final boolean inRotation = screenRotationAnimation != null &&
-                    screenRotationAnimation.isAnimating();
-            if (DEBUG_SCREENSHOT && inRotation) Slog.v(TAG_WM, "Taking screenshot while rotating");
-
-            // TODO(b/68392460): We should screenshot Task controls directly
-            // but it's difficult at the moment as the Task doesn't have the
-            // correct size set.
-            final Bitmap bitmap = SurfaceControl.screenshot(frame, dw, dh, 0, 1, inRotation, rot);
-            if (bitmap == null) {
-                Slog.w(TAG_WM, "Failed to take screenshot");
-                return null;
-            }
-
-            // Create a copy of the screenshot that is immutable and backed in ashmem.
-            // This greatly reduces the overhead of passing the bitmap between processes.
-            final Bitmap ret = bitmap.createAshmemBitmap(config);
-            bitmap.recycle();
-            return ret;
+            return null;
         }
-    }
 
-    private boolean shouldScreenshotWallpaper() {
-        MutableBoolean screenshotReady = new MutableBoolean(false);
+        int dw = mDisplayInfo.logicalWidth;
+        int dh = mDisplayInfo.logicalHeight;
 
-        forAllWindows(w -> {
-            if (!w.mIsWallpaper) {
-                return false;
-            }
+        if (dw <= 0 || dh <= 0) {
+            return null;
+        }
 
-            // Found the wallpaper window
-            final WindowStateAnimator winAnim = w.mWinAnimator;
+        final Rect frame = new Rect(0, 0, dw, dh);
 
-            if (winAnim.getShown() && winAnim.mLastAlpha > 0f) {
-                screenshotReady.value = true;
-            }
+        // The screenshot API does not apply the current screen rotation.
+        int rot = mDisplay.getRotation();
 
-            return true;
-        }, true /* traverseTopToBottom */);
+        if (rot == ROTATION_90 || rot == ROTATION_270) {
+            rot = (rot == ROTATION_90) ? ROTATION_270 : ROTATION_90;
+        }
 
-        return screenshotReady.value;
+        // SurfaceFlinger is not aware of orientation, so convert our logical
+        // crop to SurfaceFlinger's portrait orientation.
+        convertCropForSurfaceFlinger(frame, rot, dw, dh);
+
+        final ScreenRotationAnimation screenRotationAnimation =
+                mService.mAnimator.getScreenRotationAnimationLocked(DEFAULT_DISPLAY);
+        final boolean inRotation = screenRotationAnimation != null &&
+                screenRotationAnimation.isAnimating();
+        if (DEBUG_SCREENSHOT && inRotation) Slog.v(TAG_WM, "Taking screenshot while rotating");
+
+        // TODO(b/68392460): We should screenshot Task controls directly
+        // but it's difficult at the moment as the Task doesn't have the
+        // correct size set.
+        final Bitmap bitmap = SurfaceControl.screenshot(frame, dw, dh, 0, 1, inRotation, rot);
+        if (bitmap == null) {
+            Slog.w(TAG_WM, "Failed to take screenshot");
+            return null;
+        }
+
+        // Create a copy of the screenshot that is immutable and backed in ashmem.
+        // This greatly reduces the overhead of passing the bitmap between processes.
+        final Bitmap ret = bitmap.createAshmemBitmap(config);
+        bitmap.recycle();
+        return ret;
     }
 
     // TODO: Can this use createRotationMatrix()?
diff --git a/services/core/java/com/android/server/wm/WallpaperController.java b/services/core/java/com/android/server/wm/WallpaperController.java
index 2873b6d..c509980 100644
--- a/services/core/java/com/android/server/wm/WallpaperController.java
+++ b/services/core/java/com/android/server/wm/WallpaperController.java
@@ -27,12 +27,16 @@
 import static com.android.server.policy.WindowManagerPolicy.FINISH_LAYOUT_REDO_WALLPAPER;
 import static android.view.WindowManager.TRANSIT_FLAG_KEYGUARD_GOING_AWAY_WITH_WALLPAPER;
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_APP_TRANSITIONS;
+import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_SCREENSHOT;
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_WALLPAPER;
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_WALLPAPER_LIGHT;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WITH_CLASS_NAME;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
 import static com.android.server.wm.WindowManagerService.H.WALLPAPER_DRAW_PENDING_TIMEOUT;
 
+import android.graphics.Bitmap;
+import android.graphics.GraphicBuffer;
+import android.graphics.Rect;
 import android.os.Bundle;
 import android.os.Debug;
 import android.os.IBinder;
@@ -41,6 +45,7 @@
 import android.util.ArraySet;
 import android.util.Slog;
 import android.view.DisplayInfo;
+import android.view.SurfaceControl;
 import android.view.WindowManager;
 import android.view.animation.Animation;
 
@@ -95,6 +100,12 @@
     private static final int WALLPAPER_DRAW_TIMEOUT = 2;
     private int mWallpaperDrawState = WALLPAPER_DRAW_NORMAL;
 
+    /**
+     * Temporary storage for taking a screenshot of the wallpaper.
+     * @see #screenshotWallpaperLocked()
+     */
+    private WindowState mTmpTopWallpaper;
+
     private final FindWallpaperTargetResult mFindResults = new FindWallpaperTargetResult();
 
     private final ToBooleanFunction<WindowState> mFindWallpaperTargetFunction = w -> {
@@ -679,6 +690,58 @@
         mWallpaperTokens.remove(token);
     }
 
+    /**
+     * Take a screenshot of the wallpaper if it's visible.
+     *
+     * @return Bitmap of the wallpaper
+     */
+    Bitmap screenshotWallpaperLocked() {
+        if (!mService.mPolicy.isScreenOn()) {
+            if (DEBUG_SCREENSHOT) {
+                Slog.i(TAG_WM, "Attempted to take screenshot while display was off.");
+            }
+            return null;
+        }
+
+        final WindowState wallpaperWindowState = getTopVisibleWallpaper();
+        if (wallpaperWindowState == null) {
+            if (DEBUG_SCREENSHOT) {
+                Slog.i(TAG_WM, "No visible wallpaper to screenshot");
+            }
+            return null;
+        }
+
+        final Rect bounds = wallpaperWindowState.getBounds();
+        bounds.offsetTo(0, 0);
+
+        GraphicBuffer wallpaperBuffer = SurfaceControl.captureLayers(
+                wallpaperWindowState.getSurfaceControl().getHandle(), bounds, 1 /* frameScale */);
+
+        if (wallpaperBuffer == null) {
+            Slog.w(TAG_WM, "Failed to screenshot wallpaper");
+            return null;
+        }
+        return Bitmap.createHardwareBitmap(wallpaperBuffer);
+    }
+
+    private WindowState getTopVisibleWallpaper() {
+        mTmpTopWallpaper = null;
+
+        for (int curTokenNdx = mWallpaperTokens.size() - 1; curTokenNdx >= 0; curTokenNdx--) {
+            final WallpaperWindowToken token = mWallpaperTokens.get(curTokenNdx);
+            token.forAllWindows(w -> {
+                final WindowStateAnimator winAnim = w.mWinAnimator;
+                if (winAnim != null && winAnim.getShown() && winAnim.mLastAlpha > 0f) {
+                    mTmpTopWallpaper = w;
+                    return true;
+                }
+                return false;
+            }, true /* traverseTopToBottom */);
+        }
+
+        return mTmpTopWallpaper;
+    }
+
     void dump(PrintWriter pw, String prefix) {
         pw.print(prefix); pw.print("mWallpaperTarget="); pw.println(mWallpaperTarget);
         if (mPrevWallpaperTarget != null) {
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 78c04e8..bdd64d5 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -28,8 +28,6 @@
 import static android.content.Intent.ACTION_USER_REMOVED;
 import static android.content.Intent.EXTRA_USER_HANDLE;
 import static android.content.pm.PackageManager.PERMISSION_GRANTED;
-import static android.os.Process.ROOT_UID;
-import static android.os.Process.SHELL_UID;
 import static android.os.Process.SYSTEM_UID;
 import static android.os.Process.myPid;
 import static android.os.Trace.TRACE_TAG_WINDOW_MANAGER;
@@ -3607,14 +3605,14 @@
 
     @Override
     public Bitmap screenshotWallpaper() {
-        if (!checkCallingPermission(READ_FRAME_BUFFER,
-                "screenshotWallpaper()")) {
+        if (!checkCallingPermission(READ_FRAME_BUFFER, "screenshotWallpaper()")) {
             throw new SecurityException("Requires READ_FRAME_BUFFER permission");
         }
         try {
             Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "screenshotWallpaper");
-            return screenshotApplications(DEFAULT_DISPLAY, Bitmap.Config.ARGB_8888,
-                    true /* wallpaperOnly */);
+            synchronized (mWindowMap) {
+                return mRoot.mWallpaperController.screenshotWallpaperLocked();
+            }
         } finally {
             Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
         }
@@ -3627,14 +3625,25 @@
      */
     @Override
     public boolean requestAssistScreenshot(final IAssistDataReceiver receiver) {
-        if (!checkCallingPermission(READ_FRAME_BUFFER,
-                "requestAssistScreenshot()")) {
+        if (!checkCallingPermission(READ_FRAME_BUFFER, "requestAssistScreenshot()")) {
             throw new SecurityException("Requires READ_FRAME_BUFFER permission");
         }
 
+        final Bitmap bm;
+        synchronized (mWindowMap) {
+            final DisplayContent displayContent = mRoot.getDisplayContent(DEFAULT_DISPLAY);
+            if (displayContent == null) {
+                if (DEBUG_SCREENSHOT) {
+                    Slog.i(TAG_WM, "Screenshot returning null. No Display for displayId="
+                            + DEFAULT_DISPLAY);
+                }
+                bm = null;
+            } else {
+                bm = displayContent.screenshotDisplayLocked(Bitmap.Config.ARGB_8888);
+            }
+        }
+
         FgThread.getHandler().post(() -> {
-            Bitmap bm = screenshotApplications(DEFAULT_DISPLAY, Bitmap.Config.ARGB_8888,
-                    false /* wallpaperOnly */);
             try {
                 receiver.onHandleAssistScreenshot(bm);
             } catch (RemoteException e) {
@@ -3664,28 +3673,6 @@
     }
 
     /**
-     * Takes a snapshot of the screen.  In landscape mode this grabs the whole screen.
-     * In portrait mode, it grabs the full screenshot.
-     *
-     * @param displayId the Display to take a screenshot of.
-     * @param config of the output bitmap
-     * @param wallpaperOnly true if only the wallpaper layer should be included in the screenshot
-     */
-    private Bitmap screenshotApplications(int displayId, Bitmap.Config config,
-            boolean wallpaperOnly) {
-        final DisplayContent displayContent;
-        synchronized(mWindowMap) {
-            displayContent = mRoot.getDisplayContent(displayId);
-            if (displayContent == null) {
-                if (DEBUG_SCREENSHOT) Slog.i(TAG_WM, "Screenshot returning null. No Display for "
-                        + "displayId=" + displayId);
-                return null;
-            }
-        }
-        return displayContent.screenshotDisplay(config, wallpaperOnly);
-    }
-
-    /**
      * Freeze rotation changes.  (Enable "rotation lock".)
      * Persists across reboots.
      * @param rotation The desired rotation to freeze to, or -1 to use the