Remove rotation and use flag useIdentityTransform for screenshots.

There's a lot of confusing logic where 90 and 270 rotation values need
to be flipped to ensure the screenshot is taken the correct orientation.
There's also confusion what useIdentityTransform means, especially if a
non 0 rotation value is sent.

The cases screenshot cares about is the following:
1. Take screenshot in current display orientation
2. Take screenshot with 0 rotation so the caller can handle rotating the
screenshot themselves.

With these two cases in mind, remove the rotation value passed in for
screenshots. If useIdentityTransform is true, it will rotate the
screenshot so it's in the 0 orientation. If useIdentityTransform is
false, it will use the current display rotation.

This simplifies the caller logic since they no longer have to find the
current display rotation to ensure the screenshot is taken in the
current rotation. The callers can just request the screenshot with
useIdentityTransform set to false.

Test: adb shell screencap
Test: Power + volume screenshot
Test: Screen rotation
Fixes: 135942984
Change-Id: I3435ee8b5dac05e910ec1e695f398c5dcdcff9e9

9 files changed, 81 insertions, 201 deletions

diff --git a/core/java/android/app/IUiAutomationConnection.aidl b/core/java/android/app/IUiAutomationConnection.aidl
index 8c3180b400ef..4c9e400681ee 100644
--- a/core/java/android/app/IUiAutomationConnection.aidl
+++ b/core/java/android/app/IUiAutomationConnection.aidl
@@ -39,7 +39,7 @@ interface IUiAutomationConnection {
     boolean injectInputEvent(in InputEvent event, boolean sync);
     void syncInputTransactions();
     boolean setRotation(int rotation);
-    Bitmap takeScreenshot(in Rect crop, int rotation);
+    Bitmap takeScreenshot(in Rect crop);
     boolean clearWindowContentFrameStats(int windowId);
     WindowContentFrameStats getWindowContentFrameStats(int windowId);
     void clearWindowAnimationFrameStats();
diff --git a/core/java/android/app/UiAutomation.java b/core/java/android/app/UiAutomation.java
index e0951bf3f4d2..109205fadf18 100644
--- a/core/java/android/app/UiAutomation.java
+++ b/core/java/android/app/UiAutomation.java
@@ -903,7 +903,7 @@ public final class UiAutomation {
         try {
             // Calling out without a lock held.
             screenShot = mUiAutomationConnection.takeScreenshot(
-                    new Rect(0, 0, displaySize.x, displaySize.y), rotation);
+                    new Rect(0, 0, displaySize.x, displaySize.y));
             if (screenShot == null) {
                 return null;
             }
diff --git a/core/java/android/app/UiAutomationConnection.java b/core/java/android/app/UiAutomationConnection.java
index ce51dba76780..70d520176ca1 100644
--- a/core/java/android/app/UiAutomationConnection.java
+++ b/core/java/android/app/UiAutomationConnection.java
@@ -180,7 +180,7 @@ public final class UiAutomationConnection extends IUiAutomationConnection.Stub {
     }
 
     @Override
-    public Bitmap takeScreenshot(Rect crop, int rotation) {
+    public Bitmap takeScreenshot(Rect crop) {
         synchronized (mLock) {
             throwIfCalledByNotTrustedUidLocked();
             throwIfShutdownLocked();
@@ -190,7 +190,15 @@ public final class UiAutomationConnection extends IUiAutomationConnection.Stub {
         try {
             int width = crop.width();
             int height = crop.height();
-            return SurfaceControl.screenshot(crop, width, height, rotation);
+            final IBinder displayToken = SurfaceControl.getInternalDisplayToken();
+            final SurfaceControl.DisplayCaptureArgs captureArgs =
+                    new SurfaceControl.DisplayCaptureArgs.Builder(displayToken)
+                            .setSourceCrop(crop)
+                            .setSize(width, height)
+                            .build();
+            final SurfaceControl.ScreenshotHardwareBuffer screenshotBuffer =
+                    SurfaceControl.captureDisplay(captureArgs);
+            return screenshotBuffer == null ? null : screenshotBuffer.asBitmap();
         } finally {
             Binder.restoreCallingIdentity(identity);
         }
diff --git a/core/java/android/view/SurfaceControl.java b/core/java/android/view/SurfaceControl.java
index 50ed00cd0aa7..8bfa39105c7f 100644
--- a/core/java/android/view/SurfaceControl.java
+++ b/core/java/android/view/SurfaceControl.java
@@ -22,8 +22,6 @@ import static android.graphics.Matrix.MSKEW_X;
 import static android.graphics.Matrix.MSKEW_Y;
 import static android.graphics.Matrix.MTRANS_X;
 import static android.graphics.Matrix.MTRANS_Y;
-import static android.view.Surface.ROTATION_270;
-import static android.view.Surface.ROTATION_90;
 import static android.view.SurfaceControlProto.HASH_CODE;
 import static android.view.SurfaceControlProto.NAME;
 
@@ -590,6 +588,26 @@ public final class SurfaceControl implements Parcelable {
         public boolean containsSecureLayers() {
             return mContainsSecureLayers;
         }
+
+        /**
+         * Copy content of ScreenshotHardwareBuffer into a hardware bitmap and return it.
+         * 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)}
+         *
+         * CAVEAT: This can be extremely slow; avoid use unless absolutely necessary; prefer to
+         * directly
+         * use the {@link HardwareBuffer} directly.
+         *
+         * @return Bitmap generated from the {@link HardwareBuffer}
+         */
+        public Bitmap asBitmap() {
+            if (mHardwareBuffer == null) {
+                Log.w(TAG, "Failed to take screenshot. Null screenshot object");
+                return null;
+            }
+            return Bitmap.wrapHardwareBuffer(mHardwareBuffer, mColorSpace);
+        }
     }
 
     /**
@@ -597,7 +615,7 @@ public final class SurfaceControl implements Parcelable {
      * are shared between {@link DisplayCaptureArgs} and {@link LayerCaptureArgs}
      * @hide
      */
-    public abstract static class CaptureArgs {
+    private abstract static class CaptureArgs {
         private final int mPixelFormat;
         private final Rect mSourceCrop = new Rect();
         private final float mFrameScale;
@@ -615,7 +633,7 @@ public final class SurfaceControl implements Parcelable {
          *
          * @param <T> A builder that extends {@link Builder}
          */
-        public abstract static class Builder<T extends Builder<T>> {
+        abstract static class Builder<T extends Builder<T>> {
             private int mPixelFormat = PixelFormat.RGBA_8888;
             private final Rect mSourceCrop = new Rect();
             private float mFrameScale = 1;
@@ -675,7 +693,6 @@ public final class SurfaceControl implements Parcelable {
         private final int mWidth;
         private final int mHeight;
         private final boolean mUseIdentityTransform;
-        private final int mRotation;
 
         private DisplayCaptureArgs(Builder builder) {
             super(builder);
@@ -683,7 +700,6 @@ public final class SurfaceControl implements Parcelable {
             mWidth = builder.mWidth;
             mHeight = builder.mHeight;
             mUseIdentityTransform = builder.mUseIdentityTransform;
-            mRotation = builder.mRotation;
         }
 
         /**
@@ -694,7 +710,6 @@ public final class SurfaceControl implements Parcelable {
             private int mWidth;
             private int mHeight;
             private boolean mUseIdentityTransform;
-            private @Surface.Rotation int mRotation = Surface.ROTATION_0;
 
             /**
              * Construct a new {@link LayerCaptureArgs} with the set parameters. The builder
@@ -736,26 +751,16 @@ public final class SurfaceControl implements Parcelable {
             }
 
             /**
-             * Replace whatever transformation (rotation, scaling, translation) the surface
-             * layers are currently using with the identity transformation while taking the
-             * screenshot.
+             * Replace the rotation transform of the display with the identity transformation while
+             * taking the screenshot. This ensures the screenshot is taken in the ROTATION_0
+             * orientation. Set this value to false if the screenshot should be taken in the
+             * current screen orientation.
              */
             public Builder setUseIdentityTransform(boolean useIdentityTransform) {
                 mUseIdentityTransform = useIdentityTransform;
                 return this;
             }
 
-            /**
-             * Apply a custom clockwise rotation to the screenshot, i.e.
-             * Surface.ROTATION_0,90,180,270. SurfaceFlinger will always take screenshots in its
-             * native portrait orientation by default, so this is useful for returning screenshots
-             * that are independent of device orientation.
-             */
-            public Builder setRotation(@Surface.Rotation int rotation) {
-                mRotation = rotation;
-                return this;
-            }
-
             @Override
             Builder getThis() {
                 return this;
@@ -2221,130 +2226,16 @@ public final class SurfaceControl implements Parcelable {
     }
 
     /**
-     * @see SurfaceControl#screenshot(Rect, int, int, boolean, int)}
-     * @hide
-     */
-    @UnsupportedAppUsage
-    public static Bitmap screenshot(Rect sourceCrop, int width, int height, int rotation) {
-        return screenshot(sourceCrop, width, height, false, rotation);
-    }
-
-    /**
-     * Copy the current screen contents into a hardware bitmap and return it.
-     * 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)}
+     * Captures all the surfaces in a display and returns a {@link ScreenshotHardwareBuffer} with
+     * the content.
      *
-     * CAVEAT: Versions of screenshot that return a {@link Bitmap} can be extremely slow; avoid use
-     * unless absolutely necessary; prefer the versions that use a {@link HardwareBuffer} such as
-     * {@link SurfaceControl#screenshotToBuffer(IBinder, Rect, int, int, boolean, int)}.
-     *
-     * @see SurfaceControl#screenshotToBuffer(IBinder, Rect, int, int, boolean, int)}
      * @hide
      */
-    @UnsupportedAppUsage
-    public static Bitmap screenshot(Rect sourceCrop, int width, int height,
-            boolean useIdentityTransform, int rotation) {
-        // TODO: should take the display as a parameter
-        final IBinder displayToken = SurfaceControl.getInternalDisplayToken();
-        if (displayToken == null) {
-            Log.w(TAG, "Failed to take screenshot because internal display is disconnected");
-            return null;
-        }
-
-        if (rotation == ROTATION_90 || rotation == ROTATION_270) {
-            rotation = (rotation == ROTATION_90) ? ROTATION_270 : ROTATION_90;
-        }
-
-        SurfaceControl.rotateCropForSF(sourceCrop, rotation);
-        final ScreenshotHardwareBuffer buffer = screenshotToBuffer(displayToken, sourceCrop, width,
-                height, useIdentityTransform, rotation);
-
-        if (buffer == null) {
-            Log.w(TAG, "Failed to take screenshot");
-            return null;
-        }
-        return Bitmap.wrapHardwareBuffer(buffer.getHardwareBuffer(), buffer.getColorSpace());
-    }
-
-    /**
-     * Captures all the surfaces in a display and returns a {@link HardwareBuffer} with the content.
-     *
-     * @param display              The display to take the screenshot of.
-     * @param sourceCrop           The portion of the screen to capture into the Bitmap; caller may
-     *                             pass in 'new Rect()' if no cropping is desired.
-     * @param width                The desired width of the returned bitmap; the raw screen will be
-     *                             scaled down to this size; caller may pass in 0 if no scaling is
-     *                             desired.
-     * @param height               The desired height of the returned bitmap; the raw screen will
-     *                             be scaled down to this size; caller may pass in 0 if no scaling
-     *                             is desired.
-     * @param useIdentityTransform Replace whatever transformation (rotation, scaling, translation)
-     *                             the surface layers are currently using with the identity
-     *                             transformation while taking the screenshot.
-     * @param rotation             Apply a custom clockwise rotation to the screenshot, i.e.
-     *                             Surface.ROTATION_0,90,180,270. SurfaceFlinger will always take
-     *                             screenshots in its native portrait orientation by default, so
-     *                             this is useful for returning screenshots that are independent of
-     *                             device orientation.
-     * @return Returns a HardwareBuffer that contains the captured content.
-     * @hide
-     */
-    public static ScreenshotHardwareBuffer screenshotToBuffer(IBinder display, Rect sourceCrop,
-            int width, int height, boolean useIdentityTransform, int rotation) {
-        if (display == null) {
-            throw new IllegalArgumentException("displayToken must not be null");
-        }
-
-        DisplayCaptureArgs captureArgs = new DisplayCaptureArgs.Builder(display)
-                .setSourceCrop(sourceCrop)
-                .setSize(width, height)
-                .setUseIdentityTransform(useIdentityTransform)
-                .setRotation(rotation)
-                .build();
-
+    public static ScreenshotHardwareBuffer captureDisplay(DisplayCaptureArgs captureArgs) {
         return nativeCaptureDisplay(captureArgs);
     }
 
     /**
-     * Like screenshotToBuffer, but if the caller is AID_SYSTEM, allows
-     * for the capture of secure layers. This is used for the screen rotation
-     * animation where the system server takes screenshots but does
-     * not persist them or allow them to leave the server. However in other
-     * cases in the system server, we mostly want to omit secure layers
-     * like when we take a screenshot on behalf of the assistant.
-     *
-     * @hide
-     */
-    public static ScreenshotHardwareBuffer screenshotToBufferWithSecureLayersUnsafe(IBinder display,
-            Rect sourceCrop, int width, int height, boolean useIdentityTransform,
-            int rotation) {
-        if (display == null) {
-            throw new IllegalArgumentException("displayToken must not be null");
-        }
-
-        DisplayCaptureArgs captureArgs = new DisplayCaptureArgs.Builder(display)
-                .setSourceCrop(sourceCrop)
-                .setSize(width, height)
-                .setUseIdentityTransform(useIdentityTransform)
-                .setRotation(rotation)
-                .setCaptureSecureLayers(true)
-                .build();
-
-        return nativeCaptureDisplay(captureArgs);
-    }
-
-    private static void rotateCropForSF(Rect crop, int rot) {
-        if (rot == Surface.ROTATION_90 || rot == Surface.ROTATION_270) {
-            int tmp = crop.top;
-            crop.top = crop.left;
-            crop.left = tmp;
-            tmp = crop.right;
-            crop.right = crop.bottom;
-            crop.bottom = tmp;
-        }
-    }
-
-    /**
      * Captures a layer and its children and returns a {@link HardwareBuffer} with the content.
      *
      * @param layer            The root layer to capture.
diff --git a/core/jni/android_view_SurfaceControl.cpp b/core/jni/android_view_SurfaceControl.cpp
index 814a07e7f2df..f47f5e75421a 100644
--- a/core/jni/android_view_SurfaceControl.cpp
+++ b/core/jni/android_view_SurfaceControl.cpp
@@ -116,7 +116,6 @@ static struct {
     jfieldID width;
     jfieldID height;
     jfieldID useIdentityTransform;
-    jfieldID rotation;
 } gDisplayCaptureArgsClassInfo;
 
 static struct {
@@ -325,8 +324,6 @@ static DisplayCaptureArgs displayCaptureArgsFromObject(JNIEnv* env,
     captureArgs.useIdentityTransform =
             env->GetBooleanField(displayCaptureArgsObject,
                                  gDisplayCaptureArgsClassInfo.useIdentityTransform);
-    captureArgs.rotation = ui::toRotation(
-            env->GetIntField(displayCaptureArgsObject, gDisplayCaptureArgsClassInfo.rotation));
     return captureArgs;
 }
 
@@ -1848,8 +1845,6 @@ int register_android_view_SurfaceControl(JNIEnv* env)
             GetFieldIDOrDie(env, displayCaptureArgsClazz, "mHeight", "I");
     gDisplayCaptureArgsClassInfo.useIdentityTransform =
             GetFieldIDOrDie(env, displayCaptureArgsClazz, "mUseIdentityTransform", "Z");
-    gDisplayCaptureArgsClassInfo.rotation =
-            GetFieldIDOrDie(env, displayCaptureArgsClazz, "mRotation", "I");
 
     jclass layerCaptureArgsClazz =
             FindClassOrDie(env, "android/view/SurfaceControl$LayerCaptureArgs");
diff --git a/packages/Shell/src/com/android/shell/Screenshooter.java b/packages/Shell/src/com/android/shell/Screenshooter.java
index 8e0161961a49..85f25528f07e 100644
--- a/packages/Shell/src/com/android/shell/Screenshooter.java
+++ b/packages/Shell/src/com/android/shell/Screenshooter.java
@@ -17,11 +17,8 @@
 package com.android.shell;
 
 import android.graphics.Bitmap;
-import android.graphics.Point;
-import android.graphics.Rect;
-import android.hardware.display.DisplayManagerGlobal;
+import android.os.IBinder;
 import android.util.Log;
-import android.view.Display;
 import android.view.SurfaceControl;
 
 /**
@@ -40,22 +37,17 @@ final class Screenshooter {
      * @return The screenshot bitmap on success, null otherwise.
      */
     static Bitmap takeScreenshot() {
-        Display display = DisplayManagerGlobal.getInstance()
-                .getRealDisplay(Display.DEFAULT_DISPLAY);
-        Point displaySize = new Point();
-        display.getRealSize(displaySize);
-        final int displayWidth = displaySize.x;
-        final int displayHeight = displaySize.y;
-
-        int rotation = display.getRotation();
-        Rect crop = new Rect(0, 0, displayWidth, displayHeight);
-        Log.d(TAG, "Taking screenshot of dimensions " + displayWidth + " x " + displayHeight);
+        Log.d(TAG, "Taking fullscreen screenshot");
         // Take the screenshot
-        Bitmap screenShot =
-                SurfaceControl.screenshot(crop, displayWidth, displayHeight, rotation);
+        final IBinder displayToken = SurfaceControl.getInternalDisplayToken();
+        final SurfaceControl.DisplayCaptureArgs captureArgs =
+                new SurfaceControl.DisplayCaptureArgs.Builder(displayToken)
+                        .build();
+        final SurfaceControl.ScreenshotHardwareBuffer screenshotBuffer =
+                SurfaceControl.captureDisplay(captureArgs);
+        final Bitmap screenShot = screenshotBuffer == null ? null : screenshotBuffer.asBitmap();
         if (screenShot == null) {
-            Log.e(TAG, "Failed to take screenshot of dimensions " + displayWidth + " x "
-                    + displayHeight);
+            Log.e(TAG, "Failed to take fullscreen screenshot");
             return null;
         }
 
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/GlobalScreenshot.java b/packages/SystemUI/src/com/android/systemui/screenshot/GlobalScreenshot.java
index c53523032353..6747281ac437 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/GlobalScreenshot.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/GlobalScreenshot.java
@@ -50,6 +50,7 @@ import android.graphics.drawable.LayerDrawable;
 import android.media.MediaActionSound;
 import android.net.Uri;
 import android.os.Handler;
+import android.os.IBinder;
 import android.os.Looper;
 import android.os.Message;
 import android.os.RemoteException;
@@ -556,11 +557,18 @@ public class GlobalScreenshot implements ViewTreeObserver.OnComputeInternalInset
     private void takeScreenshotInternal(Consumer<Uri> finisher, Rect crop) {
         // copy the input Rect, since SurfaceControl.screenshot can mutate it
         Rect screenRect = new Rect(crop);
-        int rot = mDisplay.getRotation();
         int width = crop.width();
         int height = crop.height();
-        saveScreenshot(SurfaceControl.screenshot(crop, width, height, rot), finisher, screenRect,
-                Insets.NONE, true);
+        final IBinder displayToken = SurfaceControl.getInternalDisplayToken();
+        final SurfaceControl.DisplayCaptureArgs captureArgs =
+                new SurfaceControl.DisplayCaptureArgs.Builder(displayToken)
+                        .setSourceCrop(crop)
+                        .setSize(width, height)
+                        .build();
+        final SurfaceControl.ScreenshotHardwareBuffer screenshotBuffer =
+                SurfaceControl.captureDisplay(captureArgs);
+        final Bitmap screenshot = screenshotBuffer == null ? null : screenshotBuffer.asBitmap();
+        saveScreenshot(screenshot, finisher, screenRect, Insets.NONE, true);
     }
 
     private void saveScreenshot(Bitmap screenshot, Consumer<Uri> finisher, Rect screenRect,
diff --git a/services/core/java/com/android/server/display/DisplayManagerService.java b/services/core/java/com/android/server/display/DisplayManagerService.java
index 0979ad67a8cd..6638594bbbfa 100644
--- a/services/core/java/com/android/server/display/DisplayManagerService.java
+++ b/services/core/java/com/android/server/display/DisplayManagerService.java
@@ -30,8 +30,6 @@ import static android.hardware.display.DisplayManager.VIRTUAL_DISPLAY_FLAG_TRUST
 import static android.hardware.display.DisplayViewport.VIEWPORT_EXTERNAL;
 import static android.hardware.display.DisplayViewport.VIEWPORT_INTERNAL;
 import static android.hardware.display.DisplayViewport.VIEWPORT_VIRTUAL;
-import static android.view.Surface.ROTATION_270;
-import static android.view.Surface.ROTATION_90;
 
 import android.Manifest;
 import android.annotation.NonNull;
@@ -46,7 +44,6 @@ import android.content.res.TypedArray;
 import android.database.ContentObserver;
 import android.graphics.ColorSpace;
 import android.graphics.Point;
-import android.graphics.Rect;
 import android.hardware.SensorManager;
 import android.hardware.display.AmbientBrightnessDayStats;
 import android.hardware.display.BrightnessChangeEvent;
@@ -104,7 +101,6 @@ import com.android.server.AnimationThread;
 import com.android.server.DisplayThread;
 import com.android.server.LocalServices;
 import com.android.server.SystemService;
-import com.android.server.SystemService.TargetUser;
 import com.android.server.UiThread;
 import com.android.server.wm.SurfaceAnimationThread;
 import com.android.server.wm.WindowManagerInternal;
@@ -1394,9 +1390,13 @@ public final class DisplayManagerService extends SystemService {
             }
 
             final DisplayInfo displayInfo = logicalDisplay.getDisplayInfoLocked();
-            return SurfaceControl.screenshotToBufferWithSecureLayersUnsafe(token, new Rect(),
-                    displayInfo.getNaturalWidth(), displayInfo.getNaturalHeight(),
-                    false /* useIdentityTransform */, 0 /* rotation */);
+            final SurfaceControl.DisplayCaptureArgs captureArgs =
+                    new SurfaceControl.DisplayCaptureArgs.Builder(token)
+                            .setSize(displayInfo.getNaturalWidth(), displayInfo.getNaturalHeight())
+                            .setUseIdentityTransform(true)
+                            .setCaptureSecureLayers(true)
+                            .build();
+            return SurfaceControl.captureDisplay(captureArgs);
         }
     }
 
@@ -1406,30 +1406,11 @@ public final class DisplayManagerService extends SystemService {
             if (token == null) {
                 return null;
             }
-            final LogicalDisplay logicalDisplay = mLogicalDisplays.get(displayId);
-            if (logicalDisplay == null) {
-                return null;
-            }
-
-            final DisplayInfo displayInfo = logicalDisplay.getDisplayInfoLocked();
-            // Takes screenshot based on current device orientation.
-            final Display display = DisplayManagerGlobal.getInstance()
-                    .getRealDisplay(displayId);
-            if (display == null) {
-                return null;
-            }
-            final Point displaySize = new Point();
-            display.getRealSize(displaySize);
-
-            int rotation = displayInfo.rotation;
-            // TODO (b/153382624) : This workaround solution would be removed after
-            // SurfaceFlinger fixes the inconsistency with rotation direction issue.
-            if (rotation == ROTATION_90 || rotation == ROTATION_270) {
-                rotation = (rotation == ROTATION_90) ? ROTATION_270 : ROTATION_90;
-            }
 
-            return SurfaceControl.screenshotToBuffer(token, new Rect(), displaySize.x,
-                    displaySize.y, false /* useIdentityTransform */, rotation /* rotation */);
+            final SurfaceControl.DisplayCaptureArgs captureArgs =
+                    new SurfaceControl.DisplayCaptureArgs.Builder(token)
+                            .build();
+            return SurfaceControl.captureDisplay(captureArgs);
         }
     }
 
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index fc170538994c..c4fb174d1c32 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -4063,9 +4063,14 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
         if (DEBUG_SCREENSHOT && inRotation) Slog.v(TAG_WM, "Taking screenshot while rotating");
 
         // Send invalid rect and no width and height since it will screenshot the entire display.
-        Rect frame = new Rect(0, 0, -1, -1);
-        final Bitmap bitmap = SurfaceControl.screenshot(frame, 0, 0, inRotation,
-                mDisplay.getRotation());
+        final IBinder displayToken = SurfaceControl.getInternalDisplayToken();
+        final SurfaceControl.DisplayCaptureArgs captureArgs =
+                new SurfaceControl.DisplayCaptureArgs.Builder(displayToken)
+                        .setUseIdentityTransform(inRotation)
+                        .build();
+        final SurfaceControl.ScreenshotHardwareBuffer screenshotBuffer =
+                SurfaceControl.captureDisplay(captureArgs);
+        final Bitmap bitmap = screenshotBuffer == null ? null : screenshotBuffer.asBitmap();
         if (bitmap == null) {
             Slog.w(TAG_WM, "Failed to take screenshot");
             return null;