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

5 files changed, 191 insertions, 101 deletions

diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 59bece051810..7b5e8b8139b4 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.SystemClock;
 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.internal.view.IInputMethodClient;
 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 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
      * 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;
-            }
-
-            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;
+    Bitmap screenshotDisplayLocked(Bitmap.Config config) {
+        if (!mService.mPolicy.isScreenOn()) {
+            if (DEBUG_SCREENSHOT) {
+                Slog.i(TAG_WM, "Attempted to take screenshot while display was off.");
             }
+            return null;
+        }
 
-            // 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;
-            }
+        int dw = mDisplayInfo.logicalWidth;
+        int dh = mDisplayInfo.logicalHeight;
 
-            // 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;
+        if (dw <= 0 || dh <= 0) {
+            return null;
         }
-    }
 
-    private boolean shouldScreenshotWallpaper() {
-        MutableBoolean screenshotReady = new MutableBoolean(false);
+        final Rect frame = new Rect(0, 0, dw, dh);
 
-        forAllWindows(w -> {
-            if (!w.mIsWallpaper) {
-                return false;
-            }
+        // The screenshot API does not apply the current screen rotation.
+        int rot = mDisplay.getRotation();
 
-            // Found the wallpaper window
-            final WindowStateAnimator winAnim = w.mWinAnimator;
+        if (rot == ROTATION_90 || rot == ROTATION_270) {
+            rot = (rot == ROTATION_90) ? ROTATION_270 : ROTATION_90;
+        }
 
-            if (winAnim.getShown() && winAnim.mLastAlpha > 0f) {
-                screenshotReady.value = true;
-            }
+        // SurfaceFlinger is not aware of orientation, so convert our logical
+        // crop to SurfaceFlinger's portrait orientation.
+        convertCropForSurfaceFlinger(frame, rot, dw, dh);
 
-            return true;
-        }, true /* traverseTopToBottom */);
+        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;
+        }
 
-        return screenshotReady.value;
+        // 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 2873b6db1a20..c509980e1e9e 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 android.view.WindowManager.LayoutParams.TYPE_WALLPAPER;
 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.os.SystemClock;
 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 @@ class WallpaperController {
     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 @@ class WallpaperController {
         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 78c04e8e6a2c..bdd64d504607 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.app.admin.DevicePolicyManager.ACTION_DEVICE_POLICY_MANAGER
 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 @@ public class WindowManagerService extends IWindowManager.Stub
 
     @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 @@ public class WindowManagerService extends IWindowManager.Stub
      */
     @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 @@ public class WindowManagerService extends IWindowManager.Stub
     }
 
     /**
-     * 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
diff --git a/services/tests/servicestests/src/com/android/server/wm/WallpaperControllerTests.java b/services/tests/servicestests/src/com/android/server/wm/WallpaperControllerTests.java
new file mode 100644
index 000000000000..71ead204c9df
--- /dev/null
+++ b/services/tests/servicestests/src/com/android/server/wm/WallpaperControllerTests.java
@@ -0,0 +1,67 @@
+package com.android.server.wm;
+
+import static android.view.WindowManager.LayoutParams.TYPE_WALLPAPER;
+
+import static junit.framework.TestCase.assertNotNull;
+
+import static org.junit.Assert.assertNull;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import android.graphics.Bitmap;
+import android.os.IBinder;
+import android.platform.test.annotations.Presubmit;
+import android.support.test.filters.SmallTest;
+import android.support.test.runner.AndroidJUnit4;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/**
+ * Tests for the {@link WallpaperController} class.
+ *
+ * Build/Install/Run:
+ *  atest com.android.server.wm.WallpaperControllerTests
+ */
+@SmallTest
+@Presubmit
+@RunWith(AndroidJUnit4.class)
+public class WallpaperControllerTests extends WindowTestsBase {
+    @Test
+    public void testWallpaperScreenshot() {
+        WindowSurfaceController windowSurfaceController = mock(WindowSurfaceController.class);
+
+        synchronized (sWm.mWindowMap) {
+            // No wallpaper
+            final DisplayContent dc = createNewDisplay();
+            Bitmap wallpaperBitmap = sWm.mRoot.mWallpaperController.screenshotWallpaperLocked();
+            assertNull(wallpaperBitmap);
+
+            // No wallpaper WSA Surface
+            WindowToken wallpaperWindowToken = new WallpaperWindowToken(sWm, mock(IBinder.class),
+                    true, dc, true /* ownerCanManageAppTokens */);
+            WindowState wallpaperWindow = createWindow(null /* parent */, TYPE_WALLPAPER,
+                    wallpaperWindowToken, "wallpaperWindow");
+            wallpaperBitmap = mWallpaperController.screenshotWallpaperLocked();
+            assertNull(wallpaperBitmap);
+
+            // Wallpaper with not visible WSA surface.
+            wallpaperWindow.mWinAnimator.mSurfaceController = windowSurfaceController;
+            wallpaperWindow.mWinAnimator.mLastAlpha = 1;
+            wallpaperBitmap = mWallpaperController.screenshotWallpaperLocked();
+            assertNull(wallpaperBitmap);
+
+            when(windowSurfaceController.getShown()).thenReturn(true);
+
+            // Wallpaper with WSA alpha set to 0.
+            wallpaperWindow.mWinAnimator.mLastAlpha = 0;
+            wallpaperBitmap = mWallpaperController.screenshotWallpaperLocked();
+            assertNull(wallpaperBitmap);
+
+            // Wallpaper window with WSA Surface
+            wallpaperWindow.mWinAnimator.mLastAlpha = 1;
+            wallpaperBitmap = mWallpaperController.screenshotWallpaperLocked();
+            assertNotNull(wallpaperBitmap);
+        }
+    }
+}
diff --git a/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java b/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
index 91d5ea463bbe..2890cbd5b6f5 100644
--- a/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
+++ b/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
@@ -84,6 +84,7 @@ class WindowTestsBase {
     WindowState mChildAppWindowAbove;
     WindowState mChildAppWindowBelow;
     HashSet<WindowState> mCommonWindows;
+    WallpaperController mWallpaperController;
 
     @Mock
     static WindowState.PowerManagerWrapper mPowerManagerWrapper;
@@ -105,6 +106,8 @@ class WindowTestsBase {
         sWm = TestWindowManagerPolicy.getWindowManagerService(context);
         beforeCreateDisplay();
 
+        mWallpaperController = new WallpaperController(sWm);
+
         context.getDisplay().getDisplayInfo(mDisplayInfo);
         mDisplayContent = createNewDisplay();
         sWm.mDisplayEnabled = true;
@@ -288,7 +291,7 @@ class WindowTestsBase {
         final int displayId = sNextDisplayId++;
         final Display display = new Display(DisplayManagerGlobal.getInstance(), displayId,
                 mDisplayInfo, DEFAULT_DISPLAY_ADJUSTMENTS);
-        return new DisplayContent(display, sWm, new WallpaperController(sWm),
+        return new DisplayContent(display, sWm, mWallpaperController,
                 mock(DisplayWindowController.class));
     }