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/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) {