Add captureDisplay API in WMS.

Allow clients to request a display screen capture using a specified
displayId. The caller can pass in their own required arguments and they
will get the screenshot via an async callback directly from SF.

Test: ScreenshotTests
Test: WindowManagerServiceTests#testCaptureDisplay
Bug: 242714168
Change-Id: If56bdf7aae07120e94cdb11bc6620302a886e585

9 files changed, 547 insertions, 215 deletions

diff --git a/core/java/android/view/IWindowManager.aidl b/core/java/android/view/IWindowManager.aidl
index b2a26fa665f8..067946e90361 100644
--- a/core/java/android/view/IWindowManager.aidl
+++ b/core/java/android/view/IWindowManager.aidl
@@ -67,6 +67,7 @@ import android.view.SurfaceControl;
 import android.view.displayhash.DisplayHash;
 import android.view.displayhash.VerifiedDisplayHash;
 import android.window.ITaskFpsCallback;
+import android.window.ScreenCapture;
 
 /**
  * System private interface to the window manager.
@@ -968,4 +969,11 @@ interface IWindowManager
      * treatment.
      */
     boolean isLetterboxBackgroundMultiColored();
+
+    /**
+     * Captures the entire display specified by the displayId using the args provided. If the args
+     * are null or if the sourceCrop is invalid or null, the entire display bounds will be captured.
+     */
+    oneway void captureDisplay(int displayId, in @nullable ScreenCapture.CaptureArgs captureArgs,
+            in ScreenCapture.ScreenCaptureListener listener);
 }
diff --git a/core/java/android/window/ScreenCapture.aidl b/core/java/android/window/ScreenCapture.aidl
new file mode 100644
index 000000000000..267a7c60b60d
--- /dev/null
+++ b/core/java/android/window/ScreenCapture.aidl
@@ -0,0 +1,26 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.window;
+
+/** @hide */
+parcelable ScreenCapture.CaptureArgs;
+
+/** @hide */
+parcelable ScreenCapture.ScreenshotHardwareBuffer;
+
+/** @hide */
+parcelable ScreenCapture.ScreenCaptureListener;
\ No newline at end of file
diff --git a/core/java/android/window/ScreenCapture.java b/core/java/android/window/ScreenCapture.java
index 887d027d26a8..8a7efb93d961 100644
--- a/core/java/android/window/ScreenCapture.java
+++ b/core/java/android/window/ScreenCapture.java
@@ -24,11 +24,17 @@ import android.graphics.PixelFormat;
 import android.graphics.Rect;
 import android.hardware.HardwareBuffer;
 import android.os.IBinder;
+import android.os.Parcel;
+import android.os.Parcelable;
 import android.util.Log;
+import android.util.Pair;
 import android.view.SurfaceControl;
 
+import libcore.util.NativeAllocationRegistry;
+
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
 
 /**
  * Handles display and layer captures for the system.
@@ -37,20 +43,26 @@ import java.util.concurrent.TimeUnit;
  */
 public class ScreenCapture {
     private static final String TAG = "ScreenCapture";
+    private static final int SCREENSHOT_WAIT_TIME_S = 1;
 
     private static native int nativeCaptureDisplay(DisplayCaptureArgs captureArgs,
-            ScreenCaptureListener captureListener);
+            long captureListener);
     private static native int nativeCaptureLayers(LayerCaptureArgs captureArgs,
-            ScreenCaptureListener captureListener);
+            long captureListener);
+    private static native long nativeCreateScreenCaptureListener(
+            Consumer<ScreenshotHardwareBuffer> consumer);
+    private static native void nativeWriteListenerToParcel(long nativeObject, Parcel out);
+    private static native long nativeReadListenerFromParcel(Parcel in);
+    private static native long getNativeListenerFinalizer();
 
     /**
-     * @param captureArgs Arguments about how to take the screenshot
+     * @param captureArgs     Arguments about how to take the screenshot
      * @param captureListener A listener to receive the screenshot callback
      * @hide
      */
     public static int captureDisplay(@NonNull DisplayCaptureArgs captureArgs,
             @NonNull ScreenCaptureListener captureListener) {
-        return nativeCaptureDisplay(captureArgs, captureListener);
+        return nativeCaptureDisplay(captureArgs, captureListener.mNativeObject);
     }
 
     /**
@@ -61,28 +73,29 @@ public class ScreenCapture {
      */
     public static ScreenshotHardwareBuffer captureDisplay(
             DisplayCaptureArgs captureArgs) {
-        SyncScreenCaptureListener
-                screenCaptureListener = new SyncScreenCaptureListener();
-
-        int status = captureDisplay(captureArgs, screenCaptureListener);
+        Pair<ScreenCaptureListener, ScreenshotSync> syncScreenCapture = createSyncCaptureListener();
+        int status = captureDisplay(captureArgs, syncScreenCapture.first);
         if (status != 0) {
             return null;
         }
 
-        return screenCaptureListener.waitForScreenshot();
+        try {
+            return syncScreenCapture.second.get();
+        } catch (Exception e) {
+            return null;
+        }
     }
 
     /**
      * Captures a layer and its children and returns a {@link HardwareBuffer} with the content.
      *
-     * @param layer            The root layer to capture.
-     * @param sourceCrop       The portion of the root surface to capture; caller may pass in 'new
-     *                         Rect()' or null if no cropping is desired. If the root layer does not
-     *                         have a buffer or a crop set, then a non-empty source crop must be
-     *                         specified.
-     * @param frameScale       The desired scale of the returned buffer; the raw
-     *                         screen will be scaled up/down.
-     *
+     * @param layer      The root layer to capture.
+     * @param sourceCrop The portion of the root surface to capture; caller may pass in 'new
+     *                   Rect()' or null if no cropping is desired. If the root layer does not
+     *                   have a buffer or a crop set, then a non-empty source crop must be
+     *                   specified.
+     * @param frameScale The desired scale of the returned buffer; the raw screen will be scaled
+     *                   up/down.
      * @return Returns a HardwareBuffer that contains the layer capture.
      * @hide
      */
@@ -94,15 +107,14 @@ public class ScreenCapture {
     /**
      * Captures a layer and its children and returns a {@link HardwareBuffer} with the content.
      *
-     * @param layer            The root layer to capture.
-     * @param sourceCrop       The portion of the root surface to capture; caller may pass in 'new
-     *                         Rect()' or null if no cropping is desired. If the root layer does not
-     *                         have a buffer or a crop set, then a non-empty source crop must be
-     *                         specified.
-     * @param frameScale       The desired scale of the returned buffer; the raw
-     *                         screen will be scaled up/down.
-     * @param format           The desired pixel format of the returned buffer.
-     *
+     * @param layer      The root layer to capture.
+     * @param sourceCrop The portion of the root surface to capture; caller may pass in 'new
+     *                   Rect()' or null if no cropping is desired. If the root layer does not
+     *                   have a buffer or a crop set, then a non-empty source crop must be
+     *                   specified.
+     * @param frameScale The desired scale of the returned buffer; the raw screen will be scaled
+     *                   up/down.
+     * @param format     The desired pixel format of the returned buffer.
      * @return Returns a HardwareBuffer that contains the layer capture.
      * @hide
      */
@@ -120,21 +132,24 @@ public class ScreenCapture {
     /**
      * @hide
      */
-    public static ScreenshotHardwareBuffer captureLayers(
-            LayerCaptureArgs captureArgs) {
-        SyncScreenCaptureListener screenCaptureListener = new SyncScreenCaptureListener();
-
-        int status = captureLayers(captureArgs, screenCaptureListener);
+    public static ScreenshotHardwareBuffer captureLayers(LayerCaptureArgs captureArgs) {
+        Pair<ScreenCaptureListener, ScreenshotSync> syncScreenCapture = createSyncCaptureListener();
+        int status = captureLayers(captureArgs, syncScreenCapture.first);
         if (status != 0) {
             return null;
         }
 
-        return screenCaptureListener.waitForScreenshot();
+        try {
+            return syncScreenCapture.second.get();
+        } catch (Exception e) {
+            return null;
+        }
     }
 
     /**
      * Like {@link #captureLayers(SurfaceControl, Rect, float, int)} but with an array of layer
      * handles to exclude.
+     *
      * @hide
      */
     public static ScreenshotHardwareBuffer captureLayersExcluding(SurfaceControl layer,
@@ -150,24 +165,13 @@ public class ScreenCapture {
     }
 
     /**
-     * @param captureArgs Arguments about how to take the screenshot
+     * @param captureArgs     Arguments about how to take the screenshot
      * @param captureListener A listener to receive the screenshot callback
      * @hide
      */
     public static int captureLayers(@NonNull LayerCaptureArgs captureArgs,
             @NonNull ScreenCaptureListener captureListener) {
-        return nativeCaptureLayers(captureArgs, captureListener);
-    }
-
-    /**
-     * @hide
-     */
-    public interface ScreenCaptureListener {
-        /**
-         * The callback invoked when the screen capture is complete.
-         * @param hardwareBuffer Data containing info about the screen capture.
-         */
-        void onScreenCaptureComplete(ScreenshotHardwareBuffer hardwareBuffer);
+        return nativeCaptureLayers(captureArgs, captureListener.mNativeObject);
     }
 
     /**
@@ -190,15 +194,16 @@ public class ScreenCapture {
             mContainsHdrLayers = containsHdrLayers;
         }
 
-       /**
-        * Create ScreenshotHardwareBuffer from an existing HardwareBuffer object.
-        * @param hardwareBuffer The existing HardwareBuffer object
-        * @param namedColorSpace Integer value of a named color space {@link ColorSpace.Named}
-        * @param containsSecureLayers Indicates whether this graphic buffer contains captured
-        *                             contents of secure layers, in which case the screenshot
-        *                             should not be persisted.
-        * @param containsHdrLayers Indicates whether this graphic buffer contains HDR content.
-        */
+        /**
+         * Create ScreenshotHardwareBuffer from an existing HardwareBuffer object.
+         *
+         * @param hardwareBuffer       The existing HardwareBuffer object
+         * @param namedColorSpace      Integer value of a named color space {@link ColorSpace.Named}
+         * @param containsSecureLayers Indicates whether this graphic buffer contains captured
+         *                             contents of secure layers, in which case the screenshot
+         *                             should not be persisted.
+         * @param containsHdrLayers    Indicates whether this graphic buffer contains HDR content.
+         */
         private static ScreenshotHardwareBuffer createFromNative(HardwareBuffer hardwareBuffer,
                 int namedColorSpace, boolean containsSecureLayers, boolean containsHdrLayers) {
             ColorSpace colorSpace = ColorSpace.get(ColorSpace.Named.values()[namedColorSpace]);
@@ -220,6 +225,7 @@ public class ScreenCapture {
         public boolean containsSecureLayers() {
             return mContainsSecureLayers;
         }
+
         /**
          * Returns whether the screenshot contains at least one HDR layer.
          * This information may be useful for informing the display whether this screenshot
@@ -234,7 +240,7 @@ public class ScreenCapture {
          * Note: If you want to modify the Bitmap in software, you will need to copy the Bitmap
          * into
          * a software Bitmap using {@link Bitmap#copy(Bitmap.Config, boolean)}
-         *
+         * <p>
          * CAVEAT: This can be extremely slow; avoid use unless absolutely necessary; prefer to
          * directly
          * use the {@link HardwareBuffer} directly.
@@ -250,44 +256,23 @@ public class ScreenCapture {
         }
     }
 
-    private static class SyncScreenCaptureListener implements ScreenCaptureListener {
-        private static final int SCREENSHOT_WAIT_TIME_S = 1;
-        private ScreenshotHardwareBuffer mScreenshotHardwareBuffer;
-        private final CountDownLatch mCountDownLatch = new CountDownLatch(1);
-
-        @Override
-        public void onScreenCaptureComplete(ScreenshotHardwareBuffer hardwareBuffer) {
-            mScreenshotHardwareBuffer = hardwareBuffer;
-            mCountDownLatch.countDown();
-        }
-
-        private ScreenshotHardwareBuffer waitForScreenshot() {
-            try {
-                mCountDownLatch.await(SCREENSHOT_WAIT_TIME_S, TimeUnit.SECONDS);
-            } catch (Exception e) {
-                Log.e(TAG, "Failed to wait for screen capture result", e);
-            }
-
-            return mScreenshotHardwareBuffer;
-        }
-    }
-
     /**
      * A common arguments class used for various screenshot requests. This contains arguments that
      * are shared between {@link DisplayCaptureArgs} and {@link LayerCaptureArgs}
+     *
      * @hide
      */
-    private abstract static class CaptureArgs {
-        private final int mPixelFormat;
-        private final Rect mSourceCrop = new Rect();
-        private final float mFrameScaleX;
-        private final float mFrameScaleY;
-        private final boolean mCaptureSecureLayers;
-        private final boolean mAllowProtected;
-        private final long mUid;
-        private final boolean mGrayscale;
-
-        private CaptureArgs(Builder<? extends Builder<?>> builder) {
+    public static class CaptureArgs implements Parcelable {
+        public final int mPixelFormat;
+        public final Rect mSourceCrop = new Rect();
+        public final float mFrameScaleX;
+        public final float mFrameScaleY;
+        public final boolean mCaptureSecureLayers;
+        public final boolean mAllowProtected;
+        public final long mUid;
+        public final boolean mGrayscale;
+
+        private CaptureArgs(CaptureArgs.Builder<? extends CaptureArgs.Builder<?>> builder) {
             mPixelFormat = builder.mPixelFormat;
             mSourceCrop.set(builder.mSourceCrop);
             mFrameScaleX = builder.mFrameScaleX;
@@ -298,12 +283,23 @@ public class ScreenCapture {
             mGrayscale = builder.mGrayscale;
         }
 
+        private CaptureArgs(Parcel in) {
+            mPixelFormat = in.readInt();
+            mSourceCrop.readFromParcel(in);
+            mFrameScaleX = in.readFloat();
+            mFrameScaleY = in.readFloat();
+            mCaptureSecureLayers = in.readBoolean();
+            mAllowProtected = in.readBoolean();
+            mUid = in.readLong();
+            mGrayscale = in.readBoolean();
+        }
+
         /**
          * The Builder class used to construct {@link CaptureArgs}
          *
-         * @param <T> A builder that extends {@link Builder}
+         * @param <T> A builder that extends {@link CaptureArgs.Builder}
          */
-        abstract static class Builder<T extends Builder<T>> {
+        public static class Builder<T extends CaptureArgs.Builder<T>> {
             private int mPixelFormat = PixelFormat.RGBA_8888;
             private final Rect mSourceCrop = new Rect();
             private float mFrameScaleX = 1;
@@ -314,6 +310,14 @@ public class ScreenCapture {
             private boolean mGrayscale;
 
             /**
+             * Construct a new {@link CaptureArgs} with the set parameters. The builder remains
+             * valid.
+             */
+            public CaptureArgs build() {
+                return new CaptureArgs(this);
+            }
+
+            /**
              * The desired pixel format of the returned buffer.
              */
             public T setPixelFormat(int pixelFormat) {
@@ -395,15 +399,47 @@ public class ScreenCapture {
             /**
              * Each sub class should return itself to allow the builder to chain properly
              */
-            abstract T getThis();
+            T getThis() {
+                return (T) this;
+            }
+        }
+
+        @Override
+        public int describeContents() {
+            return 0;
         }
+
+        @Override
+        public void writeToParcel(@NonNull Parcel dest, int flags) {
+            dest.writeInt(mPixelFormat);
+            mSourceCrop.writeToParcel(dest, flags);
+            dest.writeFloat(mFrameScaleX);
+            dest.writeFloat(mFrameScaleY);
+            dest.writeBoolean(mCaptureSecureLayers);
+            dest.writeBoolean(mAllowProtected);
+            dest.writeLong(mUid);
+            dest.writeBoolean(mGrayscale);
+        }
+
+        public static final Parcelable.Creator<CaptureArgs> CREATOR =
+                new Parcelable.Creator<CaptureArgs>() {
+                    @Override
+                    public CaptureArgs createFromParcel(Parcel in) {
+                        return new CaptureArgs(in);
+                    }
+
+                    @Override
+                    public CaptureArgs[] newArray(int size) {
+                        return new CaptureArgs[size];
+                    }
+                };
     }
 
     /**
      * The arguments class used to make display capture requests.
      *
-     * @see #nativeCaptureDisplay(DisplayCaptureArgs, ScreenCaptureListener)
      * @hide
+     * @see #nativeCaptureDisplay(DisplayCaptureArgs, long)
      */
     public static class DisplayCaptureArgs extends CaptureArgs {
         private final IBinder mDisplayToken;
@@ -488,8 +524,8 @@ public class ScreenCapture {
     /**
      * The arguments class used to make layer capture requests.
      *
-     * @see #nativeCaptureLayers(LayerCaptureArgs, ScreenCaptureListener)
      * @hide
+     * @see #nativeCaptureLayers(LayerCaptureArgs, long)
      */
     public static class LayerCaptureArgs extends CaptureArgs {
         private final long mNativeLayer;
@@ -530,6 +566,17 @@ public class ScreenCapture {
                 return new LayerCaptureArgs(this);
             }
 
+            public Builder(SurfaceControl layer, CaptureArgs args) {
+                setLayer(layer);
+                setPixelFormat(args.mPixelFormat);
+                setSourceCrop(args.mSourceCrop);
+                setFrameScale(args.mFrameScaleX, args.mFrameScaleY);
+                setCaptureSecureLayers(args.mCaptureSecureLayers);
+                setAllowProtected(args.mAllowProtected);
+                setUid(args.mUid);
+                setGrayscale(args.mGrayscale);
+            }
+
             public Builder(SurfaceControl layer) {
                 setLayer(layer);
             }
@@ -542,7 +589,6 @@ public class ScreenCapture {
                 return this;
             }
 
-
             /**
              * An array of layer handles to exclude.
              */
@@ -564,8 +610,109 @@ public class ScreenCapture {
             Builder getThis() {
                 return this;
             }
+        }
+    }
+
+    /**
+     * The object used to receive the results when invoking screen capture requests via
+     * {@link #captureDisplay(DisplayCaptureArgs, ScreenCaptureListener)} or
+     * {@link #captureLayers(LayerCaptureArgs, ScreenCaptureListener)}
+     *
+     * This listener can only be used for a single call to capture content call.
+     */
+    public static class ScreenCaptureListener implements Parcelable {
+        private final long mNativeObject;
+        private static final NativeAllocationRegistry sRegistry =
+                NativeAllocationRegistry.createMalloced(
+                        ScreenCaptureListener.class.getClassLoader(), getNativeListenerFinalizer());
 
+        /**
+         * @param consumer The callback invoked when the screen capture is complete.
+         */
+        public ScreenCaptureListener(Consumer<ScreenshotHardwareBuffer> consumer) {
+            mNativeObject = nativeCreateScreenCaptureListener(consumer);
+            sRegistry.registerNativeAllocation(this, mNativeObject);
+        }
+
+        private ScreenCaptureListener(Parcel in) {
+            if (in.readBoolean()) {
+                mNativeObject = nativeReadListenerFromParcel(in);
+                sRegistry.registerNativeAllocation(this, mNativeObject);
+            } else {
+                mNativeObject = 0;
+            }
+        }
+
+        @Override
+        public int describeContents() {
+            return 0;
+        }
+
+        @Override
+        public void writeToParcel(@NonNull Parcel dest, int flags) {
+            if (mNativeObject == 0) {
+                dest.writeBoolean(false);
+            } else {
+                dest.writeBoolean(true);
+                nativeWriteListenerToParcel(mNativeObject, dest);
+            }
         }
+
+        public static final Parcelable.Creator<ScreenCaptureListener> CREATOR =
+                new Parcelable.Creator<ScreenCaptureListener>() {
+                    @Override
+                    public ScreenCaptureListener createFromParcel(Parcel in) {
+                        return new ScreenCaptureListener(in);
+                    }
+
+                    @Override
+                    public ScreenCaptureListener[] newArray(int size) {
+                        return new ScreenCaptureListener[0];
+                    }
+                };
+    }
+
+    /**
+     * A helper method to handle the async screencapture callbacks synchronously. This should only
+     * be used if the screencapture caller doesn't care that it blocks waiting for a screenshot.
+     *
+     * @return a Pair that holds the {@link ScreenCaptureListener} that should be used for capture
+     * calls into SurfaceFlinger and a {@link ScreenshotSync} object to retrieve the results.
+     */
+    public static Pair<ScreenCaptureListener, ScreenshotSync> createSyncCaptureListener() {
+        final ScreenshotSync screenshotSync = new ScreenshotSync();
+        final ScreenCaptureListener screenCaptureListener = new ScreenCaptureListener(
+                screenshotSync::setScreenshotHardwareBuffer);
+        return new Pair<>(screenCaptureListener, screenshotSync);
     }
 
+    /**
+     * Helper class to synchronously get the {@link ScreenshotHardwareBuffer} when calling
+     * {@link #captureLayers(LayerCaptureArgs, ScreenCaptureListener)} or
+     * {@link #captureDisplay(DisplayCaptureArgs, ScreenCaptureListener)}
+     */
+    public static class ScreenshotSync {
+        private final CountDownLatch mCountDownLatch = new CountDownLatch(1);
+        private ScreenshotHardwareBuffer mScreenshotHardwareBuffer;
+
+        private void setScreenshotHardwareBuffer(
+                ScreenshotHardwareBuffer screenshotHardwareBuffer) {
+            mScreenshotHardwareBuffer = screenshotHardwareBuffer;
+            mCountDownLatch.countDown();
+        }
+
+        /**
+         * Get the {@link ScreenshotHardwareBuffer} synchronously. This can be null if the
+         * screenshot failed or if there was no callback in {@link #SCREENSHOT_WAIT_TIME_S} seconds.
+         */
+        public ScreenshotHardwareBuffer get() {
+            try {
+                mCountDownLatch.await(SCREENSHOT_WAIT_TIME_S, TimeUnit.SECONDS);
+                return mScreenshotHardwareBuffer;
+            } catch (Exception e) {
+                Log.e(TAG, "Failed to wait for screen capture result", e);
+                return null;
+            }
+        }
+    }
 }
diff --git a/core/jni/android_window_ScreenCapture.cpp b/core/jni/android_window_ScreenCapture.cpp
index 3bada15aa834..c1929c6535fb 100644
--- a/core/jni/android_window_ScreenCapture.cpp
+++ b/core/jni/android_window_ScreenCapture.cpp
@@ -61,9 +61,8 @@ static struct {
 } gLayerCaptureArgsClassInfo;
 
 static struct {
-    jclass clazz;
-    jmethodID onScreenCaptureComplete;
-} gScreenCaptureListenerClassInfo;
+    jmethodID accept;
+} gConsumerClassInfo;
 
 static struct {
     jclass clazz;
@@ -98,14 +97,14 @@ class ScreenCaptureListenerWrapper : public gui::BnScreenCaptureListener {
 public:
     explicit ScreenCaptureListenerWrapper(JNIEnv* env, jobject jobject) {
         env->GetJavaVM(&mVm);
-        mScreenCaptureListenerObject = env->NewGlobalRef(jobject);
-        LOG_ALWAYS_FATAL_IF(!mScreenCaptureListenerObject, "Failed to make global ref");
+        mConsumerObject = env->NewGlobalRef(jobject);
+        LOG_ALWAYS_FATAL_IF(!mConsumerObject, "Failed to make global ref");
     }
 
     ~ScreenCaptureListenerWrapper() {
-        if (mScreenCaptureListenerObject) {
-            getenv()->DeleteGlobalRef(mScreenCaptureListenerObject);
-            mScreenCaptureListenerObject = nullptr;
+        if (mConsumerObject) {
+            getenv()->DeleteGlobalRef(mConsumerObject);
+            mConsumerObject = nullptr;
         }
     }
 
@@ -113,9 +112,8 @@ public:
             const gui::ScreenCaptureResults& captureResults) override {
         JNIEnv* env = getenv();
         if (!captureResults.fenceResult.ok() || captureResults.buffer == nullptr) {
-            env->CallVoidMethod(mScreenCaptureListenerObject,
-                                gScreenCaptureListenerClassInfo.onScreenCaptureComplete, nullptr);
-            checkAndClearException(env, "onScreenCaptureComplete");
+            env->CallVoidMethod(mConsumerObject, gConsumerClassInfo.accept, nullptr);
+            checkAndClearException(env, "accept");
             return binder::Status::ok();
         }
         captureResults.fenceResult.value()->waitForever(LOG_TAG);
@@ -130,17 +128,15 @@ public:
                                             captureResults.capturedSecureLayers,
                                             captureResults.capturedHdrLayers);
         checkAndClearException(env, "builder");
-        env->CallVoidMethod(mScreenCaptureListenerObject,
-                            gScreenCaptureListenerClassInfo.onScreenCaptureComplete,
-                            screenshotHardwareBuffer);
-        checkAndClearException(env, "onScreenCaptureComplete");
+        env->CallVoidMethod(mConsumerObject, gConsumerClassInfo.accept, screenshotHardwareBuffer);
+        checkAndClearException(env, "accept");
         env->DeleteLocalRef(jhardwareBuffer);
         env->DeleteLocalRef(screenshotHardwareBuffer);
         return binder::Status::ok();
     }
 
 private:
-    jobject mScreenCaptureListenerObject;
+    jobject mConsumerObject;
     JavaVM* mVm;
 
     JNIEnv* getenv() {
@@ -194,7 +190,7 @@ static DisplayCaptureArgs displayCaptureArgsFromObject(JNIEnv* env,
 }
 
 static jint nativeCaptureDisplay(JNIEnv* env, jclass clazz, jobject displayCaptureArgsObject,
-                                 jobject screenCaptureListenerObject) {
+                                 jlong screenCaptureListenerObject) {
     const DisplayCaptureArgs captureArgs =
             displayCaptureArgsFromObject(env, displayCaptureArgsObject);
 
@@ -202,13 +198,13 @@ static jint nativeCaptureDisplay(JNIEnv* env, jclass clazz, jobject displayCaptu
         return BAD_VALUE;
     }
 
-    sp<IScreenCaptureListener> captureListener =
-            sp<ScreenCaptureListenerWrapper>::make(env, screenCaptureListenerObject);
+    sp<gui::IScreenCaptureListener> captureListener =
+            reinterpret_cast<gui::IScreenCaptureListener*>(screenCaptureListenerObject);
     return ScreenshotClient::captureDisplay(captureArgs, captureListener);
 }
 
 static jint nativeCaptureLayers(JNIEnv* env, jclass clazz, jobject layerCaptureArgsObject,
-                                jobject screenCaptureListenerObject) {
+                                jlong screenCaptureListenerObject) {
     LayerCaptureArgs captureArgs;
     getCaptureArgs(env, layerCaptureArgsObject, captureArgs);
     SurfaceControl* layer = reinterpret_cast<SurfaceControl*>(
@@ -238,21 +234,70 @@ static jint nativeCaptureLayers(JNIEnv* env, jclass clazz, jobject layerCaptureA
         }
     }
 
-    sp<IScreenCaptureListener> captureListener =
-            sp<ScreenCaptureListenerWrapper>::make(env, screenCaptureListenerObject);
+    sp<gui::IScreenCaptureListener> captureListener =
+            reinterpret_cast<gui::IScreenCaptureListener*>(screenCaptureListenerObject);
     return ScreenshotClient::captureLayers(captureArgs, captureListener);
 }
 
+static jlong nativeCreateScreenCaptureListener(JNIEnv* env, jclass clazz, jobject consumerObj) {
+    sp<gui::IScreenCaptureListener> listener =
+            sp<ScreenCaptureListenerWrapper>::make(env, consumerObj);
+    listener->incStrong((void*)nativeCreateScreenCaptureListener);
+    return reinterpret_cast<jlong>(listener.get());
+}
+
+static void nativeWriteListenerToParcel(JNIEnv* env, jclass clazz, jlong nativeObject,
+                                        jobject parcelObj) {
+    Parcel* parcel = parcelForJavaObject(env, parcelObj);
+    if (parcel == NULL) {
+        jniThrowNullPointerException(env, NULL);
+        return;
+    }
+    ScreenCaptureListenerWrapper* const self =
+            reinterpret_cast<ScreenCaptureListenerWrapper*>(nativeObject);
+    if (self != nullptr) {
+        parcel->writeStrongBinder(IInterface::asBinder(self));
+    }
+}
+
+static jlong nativeReadListenerFromParcel(JNIEnv* env, jclass clazz, jobject parcelObj) {
+    Parcel* parcel = parcelForJavaObject(env, parcelObj);
+    if (parcel == NULL) {
+        jniThrowNullPointerException(env, NULL);
+        return 0;
+    }
+    sp<gui::IScreenCaptureListener> listener =
+            interface_cast<gui::IScreenCaptureListener>(parcel->readStrongBinder());
+    if (listener == nullptr) {
+        return 0;
+    }
+    listener->incStrong((void*)nativeCreateScreenCaptureListener);
+    return reinterpret_cast<jlong>(listener.get());
+}
+
+void destroyNativeListener(void* ptr) {
+    ScreenCaptureListenerWrapper* listener = reinterpret_cast<ScreenCaptureListenerWrapper*>(ptr);
+    listener->decStrong((void*)nativeCreateScreenCaptureListener);
+}
+
+static jlong getNativeListenerFinalizer(JNIEnv* env, jclass clazz) {
+    return static_cast<jlong>(reinterpret_cast<uintptr_t>(&destroyNativeListener));
+}
+
 // ----------------------------------------------------------------------------
 
 static const JNINativeMethod sScreenCaptureMethods[] = {
         // clang-format off
-   {"nativeCaptureDisplay",
-            "(Landroid/window/ScreenCapture$DisplayCaptureArgs;Landroid/window/ScreenCapture$ScreenCaptureListener;)I",
+    {"nativeCaptureDisplay", "(Landroid/window/ScreenCapture$DisplayCaptureArgs;J)I",
             (void*)nativeCaptureDisplay },
-    {"nativeCaptureLayers",
-            "(Landroid/window/ScreenCapture$LayerCaptureArgs;Landroid/window/ScreenCapture$ScreenCaptureListener;)I",
+    {"nativeCaptureLayers",  "(Landroid/window/ScreenCapture$LayerCaptureArgs;J)I",
             (void*)nativeCaptureLayers },
+    {"nativeCreateScreenCaptureListener", "(Ljava/util/function/Consumer;)J",
+            (void*)nativeCreateScreenCaptureListener },
+    {"nativeWriteListenerToParcel", "(JLandroid/os/Parcel;)V", (void*)nativeWriteListenerToParcel },
+    {"nativeReadListenerFromParcel", "(Landroid/os/Parcel;)J",
+            (void*)nativeReadListenerFromParcel },
+    {"getNativeListenerFinalizer", "()J", (void*)getNativeListenerFinalizer },
         // clang-format on
 };
 
@@ -293,12 +338,8 @@ int register_android_window_ScreenCapture(JNIEnv* env) {
     gLayerCaptureArgsClassInfo.childrenOnly =
             GetFieldIDOrDie(env, layerCaptureArgsClazz, "mChildrenOnly", "Z");
 
-    jclass screenCaptureListenerClazz =
-            FindClassOrDie(env, "android/window/ScreenCapture$ScreenCaptureListener");
-    gScreenCaptureListenerClassInfo.clazz = MakeGlobalRefOrDie(env, screenCaptureListenerClazz);
-    gScreenCaptureListenerClassInfo.onScreenCaptureComplete =
-            GetMethodIDOrDie(env, screenCaptureListenerClazz, "onScreenCaptureComplete",
-                             "(Landroid/window/ScreenCapture$ScreenshotHardwareBuffer;)V");
+    jclass consumer = FindClassOrDie(env, "java/util/function/Consumer");
+    gConsumerClassInfo.accept = GetMethodIDOrDie(env, consumer, "accept", "(Ljava/lang/Object;)V");
 
     jclass screenshotGraphicsBufferClazz =
             FindClassOrDie(env, "android/window/ScreenCapture$ScreenshotHardwareBuffer");
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 43ef1d4d778d..4e68c00ff23b 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -9279,4 +9279,46 @@ public class WindowManagerService extends IWindowManager.Stub
                         "Unexpected letterbox background type: " + letterboxBackgroundType);
         }
     }
+
+    @Override
+    public void captureDisplay(int displayId, @Nullable ScreenCapture.CaptureArgs captureArgs,
+            ScreenCapture.ScreenCaptureListener listener) {
+        Slog.d(TAG, "captureDisplay");
+        if (!checkCallingPermission(READ_FRAME_BUFFER, "captureDisplay()")) {
+            throw new SecurityException("Requires READ_FRAME_BUFFER permission");
+        }
+
+        ScreenCapture.captureLayers(getCaptureArgs(displayId, captureArgs), listener);
+    }
+
+    @VisibleForTesting
+    ScreenCapture.LayerCaptureArgs getCaptureArgs(int displayId,
+            @Nullable ScreenCapture.CaptureArgs captureArgs) {
+        final SurfaceControl displaySurfaceControl;
+        synchronized (mGlobalLock) {
+            DisplayContent displayContent = mRoot.getDisplayContent(displayId);
+            if (displayContent == null) {
+                throw new IllegalArgumentException("Trying to screenshot and invalid display: "
+                        + displayId);
+            }
+
+            displaySurfaceControl = displayContent.getSurfaceControl();
+
+            if (captureArgs == null) {
+                captureArgs = new ScreenCapture.CaptureArgs.Builder<>()
+                        .build();
+            }
+
+            if (captureArgs.mSourceCrop.isEmpty()) {
+                displayContent.getBounds(mTmpRect);
+                mTmpRect.offsetTo(0, 0);
+            } else {
+                mTmpRect.set(captureArgs.mSourceCrop);
+            }
+        }
+
+        return new ScreenCapture.LayerCaptureArgs.Builder(displaySurfaceControl, captureArgs)
+                        .setSourceCrop(mTmpRect)
+                        .build();
+    }
 }
diff --git a/services/tests/wmtests/AndroidManifest.xml b/services/tests/wmtests/AndroidManifest.xml
index 2918365b94c8..107bbe1c79e3 100644
--- a/services/tests/wmtests/AndroidManifest.xml
+++ b/services/tests/wmtests/AndroidManifest.xml
@@ -79,7 +79,10 @@
         <activity android:name="com.android.server.wm.ActivityOptionsTest$MainActivity"
                   android:turnScreenOn="true"
                   android:showWhenLocked="true" />
-        <activity android:name="com.android.server.wm.ScreenshotTests$ScreenshotActivity" />
+        <activity android:name="com.android.server.wm.ScreenshotTests$ScreenshotActivity"
+                  android:theme="@style/WhiteBackgroundTheme"
+                  android:turnScreenOn="true"
+                  android:showWhenLocked="true"/>
         <activity android:name="android.view.cts.surfacevalidator.CapturedActivity"/>
 
         <service android:name="android.view.cts.surfacevalidator.LocalMediaProjectionService"
diff --git a/services/tests/wmtests/res/values/styles.xml b/services/tests/wmtests/res/values/styles.xml
new file mode 100644
index 000000000000..6857ff99e9b8
--- /dev/null
+++ b/services/tests/wmtests/res/values/styles.xml
@@ -0,0 +1,24 @@
+<!--
+  ~ Copyright (C) 2022 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+
+<resources>
+    <style name="WhiteBackgroundTheme" parent="@android:style/Theme.DeviceDefault">
+        <item name="android:windowNoTitle">true</item>
+        <item name="android:windowFullscreen">true</item>
+        <item name="android:windowIsTranslucent">true</item>
+        <item name="android:windowLayoutInDisplayCutoutMode">shortEdges</item>
+    </style>
+</resources>
diff --git a/services/tests/wmtests/src/com/android/server/wm/ScreenshotTests.java b/services/tests/wmtests/src/com/android/server/wm/ScreenshotTests.java
index 0b58428e9bfb..736f8f7a5ed3 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ScreenshotTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ScreenshotTests.java
@@ -16,6 +16,10 @@
 
 package com.android.server.wm;
 
+import static android.view.Display.DEFAULT_DISPLAY;
+import static android.view.WindowInsets.Type.displayCutout;
+import static android.view.WindowInsets.Type.statusBars;
+
 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 
 import static org.junit.Assert.assertNotNull;
@@ -23,20 +27,33 @@ import static org.junit.Assert.assertTrue;
 
 import android.app.Activity;
 import android.app.Instrumentation;
+import android.content.Context;
 import android.graphics.Bitmap;
 import android.graphics.Canvas;
 import android.graphics.Color;
-import android.graphics.ColorSpace;
 import android.graphics.GraphicBuffer;
+import android.graphics.Insets;
 import android.graphics.PixelFormat;
+import android.graphics.Point;
+import android.graphics.Rect;
+import android.hardware.DataSpace;
+import android.hardware.HardwareBuffer;
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.Looper;
+import android.os.ServiceManager;
 import android.platform.test.annotations.Presubmit;
+import android.util.Pair;
+import android.view.IWindowManager;
 import android.view.PointerIcon;
 import android.view.SurfaceControl;
-import android.view.WindowManager;
+import android.view.cts.surfacevalidator.BitmapPixelChecker;
+import android.view.cts.surfacevalidator.PixelColor;
+import android.view.cts.surfacevalidator.SaveBitmapHelper;
 import android.window.ScreenCapture;
+import android.window.ScreenCapture.ScreenCaptureListener;
+import android.window.ScreenCapture.ScreenshotHardwareBuffer;
+import android.window.ScreenCapture.ScreenshotSync;
 
 import androidx.annotation.Nullable;
 import androidx.test.filters.SmallTest;
@@ -45,6 +62,7 @@ import androidx.test.rule.ActivityTestRule;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.TestName;
 
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
@@ -60,6 +78,8 @@ public class ScreenshotTests {
     private static final int BUFFER_HEIGHT = 100;
 
     private final Instrumentation mInstrumentation = getInstrumentation();
+    @Rule
+    public TestName mTestName = new TestName();
 
     @Rule
     public ActivityTestRule<ScreenshotActivity> mActivityRule =
@@ -95,8 +115,8 @@ public class ScreenshotTests {
         buffer.unlockCanvasAndPost(canvas);
 
         t.show(secureSC)
-                .setBuffer(secureSC, buffer)
-                .setColorSpace(secureSC, ColorSpace.get(ColorSpace.Named.SRGB))
+                .setBuffer(secureSC, HardwareBuffer.createFromGraphicBuffer(buffer))
+                .setDataSpace(secureSC, DataSpace.DATASPACE_SRGB)
                 .apply(true);
 
         ScreenCapture.LayerCaptureArgs args = new ScreenCapture.LayerCaptureArgs.Builder(secureSC)
@@ -112,15 +132,70 @@ public class ScreenshotTests {
         Bitmap swBitmap = screenshot.copy(Bitmap.Config.ARGB_8888, false);
         screenshot.recycle();
 
-        int numMatchingPixels = PixelChecker.getNumMatchingPixels(swBitmap,
-                new PixelColor(PixelColor.RED));
-        long sizeOfBitmap = swBitmap.getWidth() * swBitmap.getHeight();
+        BitmapPixelChecker bitmapPixelChecker = new BitmapPixelChecker(PixelColor.RED);
+        Rect bounds = new Rect(0, 0, swBitmap.getWidth(), swBitmap.getHeight());
+        int numMatchingPixels = bitmapPixelChecker.getNumMatchingPixels(swBitmap, bounds);
+        int sizeOfBitmap = bounds.width() * bounds.height();
         boolean success = numMatchingPixels == sizeOfBitmap;
         swBitmap.recycle();
 
         assertTrue(success);
     }
 
+    @Test
+    public void testCaptureDisplay() throws Exception {
+        IWindowManager windowManager = IWindowManager.Stub.asInterface(
+                ServiceManager.getService(Context.WINDOW_SERVICE));
+        SurfaceControl sc = new SurfaceControl.Builder()
+                .setName("Layer")
+                .setCallsite("testCaptureDisplay")
+                .build();
+
+        SurfaceControl.Transaction t = mActivity.addChildSc(sc);
+        mInstrumentation.waitForIdleSync();
+
+        GraphicBuffer buffer = GraphicBuffer.create(BUFFER_WIDTH, BUFFER_HEIGHT,
+                PixelFormat.RGBA_8888,
+                GraphicBuffer.USAGE_HW_TEXTURE | GraphicBuffer.USAGE_HW_COMPOSER
+                        | GraphicBuffer.USAGE_SW_WRITE_RARELY);
+
+        Canvas canvas = buffer.lockCanvas();
+        canvas.drawColor(Color.RED);
+        buffer.unlockCanvasAndPost(canvas);
+
+        Point point = mActivity.getPositionBelowStatusBar();
+        t.show(sc)
+                .setBuffer(sc, HardwareBuffer.createFromGraphicBuffer(buffer))
+                .setDataSpace(sc, DataSpace.DATASPACE_SRGB)
+                .setPosition(sc, point.x, point.y)
+                .apply(true);
+
+        Pair<ScreenCaptureListener, ScreenshotSync> syncScreenCapture =
+                ScreenCapture.createSyncCaptureListener();
+        windowManager.captureDisplay(DEFAULT_DISPLAY, null, syncScreenCapture.first);
+        ScreenshotHardwareBuffer hardwareBuffer = syncScreenCapture.second.get();
+        assertNotNull(hardwareBuffer);
+
+        Bitmap screenshot = hardwareBuffer.asBitmap();
+        assertNotNull(screenshot);
+
+        Bitmap swBitmap = screenshot.copy(Bitmap.Config.ARGB_8888, false);
+        screenshot.recycle();
+
+        BitmapPixelChecker bitmapPixelChecker = new BitmapPixelChecker(PixelColor.RED);
+        Rect bounds = new Rect(point.x, point.y, BUFFER_WIDTH + point.x, BUFFER_HEIGHT + point.y);
+        int numMatchingPixels = bitmapPixelChecker.getNumMatchingPixels(swBitmap, bounds);
+        int pixelMatchSize = bounds.width() * bounds.height();
+        boolean success = numMatchingPixels == pixelMatchSize;
+
+        if (!success) {
+            SaveBitmapHelper.saveBitmap(swBitmap, getClass(), mTestName, "failedImage");
+        }
+        swBitmap.recycle();
+        assertTrue("numMatchingPixels=" + numMatchingPixels + " pixelMatchSize=" + pixelMatchSize,
+                success);
+    }
+
     public static class ScreenshotActivity extends Activity {
         private static final long WAIT_TIMEOUT_S = 5;
         private final Handler mHandler = new Handler(Looper.getMainLooper());
@@ -130,7 +205,6 @@ public class ScreenshotTests {
             super.onCreate(savedInstanceState);
             getWindow().getDecorView().setPointerIcon(
                     PointerIcon.getSystemIcon(this, PointerIcon.TYPE_NULL));
-            getWindow().addFlags(WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON);
         }
 
         SurfaceControl.Transaction addChildSc(SurfaceControl surfaceControl) {
@@ -148,88 +222,14 @@ public class ScreenshotTests {
             }
             return t;
         }
-    }
 
-    public abstract static class PixelChecker {
-        static int getNumMatchingPixels(Bitmap bitmap, PixelColor pixelColor) {
-            int numMatchingPixels = 0;
-            for (int x = 0; x < bitmap.getWidth(); x++) {
-                for (int y = 0; y < bitmap.getHeight(); y++) {
-                    int color = bitmap.getPixel(x, y);
-                    if (matchesColor(pixelColor, color)) {
-                        numMatchingPixels++;
-                    }
-                }
-            }
-            return numMatchingPixels;
-        }
-
-        static boolean matchesColor(PixelColor expectedColor, int color) {
-            final float red = Color.red(color);
-            final float green = Color.green(color);
-            final float blue = Color.blue(color);
-            final float alpha = Color.alpha(color);
-
-            return alpha <= expectedColor.mMaxAlpha
-                    && alpha >= expectedColor.mMinAlpha
-                    && red <= expectedColor.mMaxRed
-                    && red >= expectedColor.mMinRed
-                    && green <= expectedColor.mMaxGreen
-                    && green >= expectedColor.mMinGreen
-                    && blue <= expectedColor.mMaxBlue
-                    && blue >= expectedColor.mMinBlue;
-        }
-    }
-
-    public static class PixelColor {
-        public static final int BLACK = 0xFF000000;
-        public static final int RED = 0xFF0000FF;
-        public static final int GREEN = 0xFF00FF00;
-        public static final int BLUE = 0xFFFF0000;
-        public static final int YELLOW = 0xFF00FFFF;
-        public static final int MAGENTA = 0xFFFF00FF;
-        public static final int WHITE = 0xFFFFFFFF;
-
-        public static final int TRANSPARENT_RED = 0x7F0000FF;
-        public static final int TRANSPARENT_BLUE = 0x7FFF0000;
-        public static final int TRANSPARENT = 0x00000000;
-
-        // Default to black
-        public short mMinAlpha;
-        public short mMaxAlpha;
-        public short mMinRed;
-        public short mMaxRed;
-        public short mMinBlue;
-        public short mMaxBlue;
-        public short mMinGreen;
-        public short mMaxGreen;
-
-        public PixelColor(int color) {
-            short alpha = (short) ((color >> 24) & 0xFF);
-            short blue = (short) ((color >> 16) & 0xFF);
-            short green = (short) ((color >> 8) & 0xFF);
-            short red = (short) (color & 0xFF);
-
-            mMinAlpha = (short) getMinValue(alpha);
-            mMaxAlpha = (short) getMaxValue(alpha);
-            mMinRed = (short) getMinValue(red);
-            mMaxRed = (short) getMaxValue(red);
-            mMinBlue = (short) getMinValue(blue);
-            mMaxBlue = (short) getMaxValue(blue);
-            mMinGreen = (short) getMinValue(green);
-            mMaxGreen = (short) getMaxValue(green);
-        }
-
-        public PixelColor() {
-            this(BLACK);
-        }
-
-        private int getMinValue(short color) {
-            return Math.max(color - 4, 0);
-        }
+        public Point getPositionBelowStatusBar() {
+            Insets statusBarInsets = getWindow()
+                    .getDecorView()
+                    .getRootWindowInsets()
+                    .getInsets(statusBars() | displayCutout());
 
-        private int getMaxValue(short color) {
-            return Math.min(color + 4, 0xFF);
+            return new Point(statusBarInsets.left, statusBarInsets.top);
         }
     }
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java
index 46b4b76dc12f..8b63904baf31 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java
@@ -41,6 +41,7 @@ import static com.android.server.wm.LetterboxConfiguration.LETTERBOX_BACKGROUND_
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
 import static org.mockito.ArgumentMatchers.any;
@@ -68,6 +69,7 @@ import android.view.SurfaceControl;
 import android.view.View;
 import android.view.WindowManager;
 import android.window.ClientWindowFrames;
+import android.window.ScreenCapture;
 import android.window.WindowContainerToken;
 
 import androidx.test.filters.SmallTest;
@@ -423,6 +425,45 @@ public class WindowManagerServiceTests extends WindowTestsBase {
                 LETTERBOX_BACKGROUND_SOLID_COLOR)).isFalse();
     }
 
+    @Test
+    public void testCaptureDisplay() {
+        Rect displayBounds = new Rect(0, 0, 100, 200);
+        spyOn(mDisplayContent);
+        when(mDisplayContent.getBounds()).thenReturn(displayBounds);
+
+        // Null captureArgs
+        ScreenCapture.LayerCaptureArgs resultingArgs =
+                mWm.getCaptureArgs(DEFAULT_DISPLAY, null /* captureArgs */);
+        assertEquals(displayBounds, resultingArgs.mSourceCrop);
+
+        // Non null captureArgs, didn't set rect
+        ScreenCapture.CaptureArgs captureArgs = new ScreenCapture.CaptureArgs.Builder<>().build();
+        resultingArgs = mWm.getCaptureArgs(DEFAULT_DISPLAY, captureArgs);
+        assertEquals(displayBounds, resultingArgs.mSourceCrop);
+
+        // Non null captureArgs, invalid rect
+        captureArgs = new ScreenCapture.CaptureArgs.Builder<>()
+                .setSourceCrop(new Rect(0, 0, -1, -1))
+                .build();
+        resultingArgs = mWm.getCaptureArgs(DEFAULT_DISPLAY, captureArgs);
+        assertEquals(displayBounds, resultingArgs.mSourceCrop);
+
+        // Non null captureArgs, null rect
+        captureArgs = new ScreenCapture.CaptureArgs.Builder<>()
+                .setSourceCrop(null)
+                .build();
+        resultingArgs = mWm.getCaptureArgs(DEFAULT_DISPLAY, captureArgs);
+        assertEquals(displayBounds, resultingArgs.mSourceCrop);
+
+        // Non null captureArgs, valid rect
+        Rect validRect = new Rect(0, 0, 10, 50);
+        captureArgs = new ScreenCapture.CaptureArgs.Builder<>()
+                .setSourceCrop(validRect)
+                .build();
+        resultingArgs = mWm.getCaptureArgs(DEFAULT_DISPLAY, captureArgs);
+        assertEquals(validRect, resultingArgs.mSourceCrop);
+    }
+
     private void setupActivityWithLaunchCookie(IBinder launchCookie, WindowContainerToken wct) {
         final WindowContainer.RemoteToken remoteToken = mock(WindowContainer.RemoteToken.class);
         when(remoteToken.toWindowContainerToken()).thenReturn(wct);