Replace screenshots using internal displayToken with new API

In order to remove SurfaceControl.getInternalDisplayToken, we need to
replace all usages of it. The primary use is for screen capturing since
you could call directly into SF via ScreenCapture.captureDisplay if you
had a display token. However, this isn't scalable with multi-display
since you need to be explicity which display to capture.

The change provides a new API into WMS to capture a display using the
displayId. This is still a privilege call so only processes with
READ_FRAME_BUFFER can use it. In most cases, the replacement for now is
to use DEFAULT_DISPLAY, but they can be modified later to send the
desired display.

Test: power + volume down
Test: assistant screencapture
Test: SurfaceViewTests
Bug: 242714168

Change-Id: I9b406b699d48ae34c6ffd91c34941f1580f38d28

5 files changed, 54 insertions, 169 deletions

diff --git a/core/java/android/app/UiAutomationConnection.java b/core/java/android/app/UiAutomationConnection.java
index d17f2ba1b401..45c0072f23d4 100644
--- a/core/java/android/app/UiAutomationConnection.java
+++ b/core/java/android/app/UiAutomationConnection.java
@@ -16,6 +16,8 @@
 
 package android.app;
 
+import static android.view.Display.DEFAULT_DISPLAY;
+
 import android.accessibilityservice.AccessibilityServiceInfo;
 import android.accessibilityservice.IAccessibilityServiceClient;
 import android.annotation.NonNull;
@@ -35,6 +37,7 @@ import android.os.ServiceManager;
 import android.os.UserHandle;
 import android.permission.IPermissionManager;
 import android.util.Log;
+import android.util.Pair;
 import android.view.IWindowManager;
 import android.view.InputDevice;
 import android.view.InputEvent;
@@ -46,6 +49,8 @@ import android.view.WindowContentFrameStats;
 import android.view.accessibility.AccessibilityEvent;
 import android.view.accessibility.IAccessibilityManager;
 import android.window.ScreenCapture;
+import android.window.ScreenCapture.CaptureArgs;
+import android.window.ScreenCapture.ScreenCaptureListener;
 
 import libcore.io.IoUtils;
 
@@ -218,20 +223,22 @@ public final class UiAutomationConnection extends IUiAutomationConnection.Stub {
         }
         final long identity = Binder.clearCallingIdentity();
         try {
-            int width = crop.width();
-            int height = crop.height();
-            final IBinder displayToken = SurfaceControl.getInternalDisplayToken();
-            final ScreenCapture.DisplayCaptureArgs captureArgs =
-                    new ScreenCapture.DisplayCaptureArgs.Builder(displayToken)
-                            .setSourceCrop(crop)
-                            .setSize(width, height)
-                            .build();
+            final CaptureArgs captureArgs = new CaptureArgs.Builder<>()
+                    .setSourceCrop(crop)
+                    .build();
+            Pair<ScreenCaptureListener, ScreenCapture.ScreenshotSync> syncScreenCapture =
+                    ScreenCapture.createSyncCaptureListener();
+            mWindowManager.captureDisplay(DEFAULT_DISPLAY, captureArgs,
+                    syncScreenCapture.first);
             final ScreenCapture.ScreenshotHardwareBuffer screenshotBuffer =
-                    ScreenCapture.captureDisplay(captureArgs);
+                    syncScreenCapture.second.get();
             return screenshotBuffer == null ? null : screenshotBuffer.asBitmap();
+        } catch (RemoteException re) {
+            re.rethrowAsRuntimeException();
         } finally {
             Binder.restoreCallingIdentity(identity);
         }
+        return null;
     }
 
     @Nullable
diff --git a/packages/Shell/src/com/android/shell/Screenshooter.java b/packages/Shell/src/com/android/shell/Screenshooter.java
index d55eda0c7062..baaddf53bcab 100644
--- a/packages/Shell/src/com/android/shell/Screenshooter.java
+++ b/packages/Shell/src/com/android/shell/Screenshooter.java
@@ -16,11 +16,17 @@
 
 package com.android.shell;
 
+import static android.view.Display.DEFAULT_DISPLAY;
+
 import android.graphics.Bitmap;
-import android.os.IBinder;
+import android.os.RemoteException;
 import android.util.Log;
-import android.view.SurfaceControl;
+import android.util.Pair;
+import android.view.WindowManagerGlobal;
 import android.window.ScreenCapture;
+import android.window.ScreenCapture.ScreenCaptureListener;
+import android.window.ScreenCapture.ScreenshotHardwareBuffer;
+import android.window.ScreenCapture.ScreenshotSync;
 
 /**
  * Helper class used to take screenshots.
@@ -40,12 +46,15 @@ final class Screenshooter {
     static Bitmap takeScreenshot() {
         Log.d(TAG, "Taking fullscreen screenshot");
         // Take the screenshot
-        final IBinder displayToken = SurfaceControl.getInternalDisplayToken();
-        final ScreenCapture.DisplayCaptureArgs captureArgs =
-                new ScreenCapture.DisplayCaptureArgs.Builder(displayToken)
-                        .build();
-        final ScreenCapture.ScreenshotHardwareBuffer screenshotBuffer =
-                ScreenCapture.captureDisplay(captureArgs);
+        final Pair<ScreenCaptureListener, ScreenshotSync> syncScreenCapture =
+                ScreenCapture.createSyncCaptureListener();
+        try {
+            WindowManagerGlobal.getWindowManagerService().captureDisplay(DEFAULT_DISPLAY, null,
+                    syncScreenCapture.first);
+        } catch (RemoteException e) {
+            e.rethrowAsRuntimeException();
+        }
+        final ScreenshotHardwareBuffer screenshotBuffer = syncScreenCapture.second.get();
         final Bitmap screenShot = screenshotBuffer == null ? null : screenshotBuffer.asBitmap();
         if (screenShot == null) {
             Log.e(TAG, "Failed to take fullscreen screenshot");
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt
index ac5640b14716..67e9a87e4f18 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt
@@ -19,15 +19,9 @@ package com.android.systemui.screenshot
 import android.app.IActivityTaskManager
 import android.graphics.Bitmap
 import android.graphics.Rect
-import android.hardware.display.DisplayManager
-import android.os.IBinder
-import android.util.Log
-import android.view.DisplayAddress
-import android.view.SurfaceControl
+import android.view.IWindowManager
 import android.window.ScreenCapture
-import android.window.ScreenCapture.DisplayCaptureArgs
-import android.window.ScreenCapture.ScreenshotHardwareBuffer
-import androidx.annotation.VisibleForTesting
+import android.window.ScreenCapture.CaptureArgs
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
 import javax.inject.Inject
@@ -38,18 +32,18 @@ private const val TAG = "ImageCaptureImpl"
 
 @SysUISingleton
 open class ImageCaptureImpl @Inject constructor(
-    private val displayManager: DisplayManager,
+    private val windowManager: IWindowManager,
     private val atmService: IActivityTaskManager,
     @Background private val bgContext: CoroutineDispatcher
 ) : ImageCapture {
 
     override fun captureDisplay(displayId: Int, crop: Rect?): Bitmap? {
-        val width = crop?.width() ?: 0
-        val height = crop?.height() ?: 0
-        val sourceCrop = crop ?: Rect()
-        val displayToken = physicalDisplayToken(displayId) ?: return null
-        val buffer = captureDisplay(displayToken, width, height, sourceCrop)
-
+        val captureArgs = CaptureArgs.Builder()
+            .setSourceCrop(crop)
+            .build()
+        val syncScreenCapture = ScreenCapture.createSyncCaptureListener()
+        windowManager.captureDisplay(displayId, captureArgs, syncScreenCapture.first)
+        val buffer = syncScreenCapture.second.get()
         return buffer?.asBitmap()
     }
 
@@ -57,34 +51,4 @@ open class ImageCaptureImpl @Inject constructor(
         val snapshot = withContext(bgContext) { atmService.takeTaskSnapshot(taskId) } ?: return null
         return Bitmap.wrapHardwareBuffer(snapshot.hardwareBuffer, snapshot.colorSpace)
     }
-
-    @VisibleForTesting
-    open fun physicalDisplayToken(displayId: Int): IBinder? {
-        val display = displayManager.getDisplay(displayId)
-        if (display == null) {
-            Log.e(TAG, "No display with id: $displayId")
-            return null
-        }
-        val address = display.address
-        if (address !is DisplayAddress.Physical) {
-            Log.e(TAG, "Display does not have a physical address: $display")
-            return null
-        }
-        return SurfaceControl.getPhysicalDisplayToken(address.physicalDisplayId)
-    }
-
-    @VisibleForTesting
-    open fun captureDisplay(
-        displayToken: IBinder,
-        width: Int,
-        height: Int,
-        crop: Rect
-    ): ScreenshotHardwareBuffer? {
-        val captureArgs =
-            DisplayCaptureArgs.Builder(displayToken)
-                .setSize(width, height)
-                .setSourceCrop(crop)
-                .build()
-        return ScreenCapture.captureDisplay(captureArgs)
-    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ImageCaptureImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/ImageCaptureImplTest.kt
deleted file mode 100644
index 5a4bafcb7ded..000000000000
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ImageCaptureImplTest.kt
+++ /dev/null
@@ -1,95 +0,0 @@
-/*
- * 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 com.android.systemui.screenshot
-
-import android.app.IActivityTaskManager
-import android.graphics.Rect
-import android.hardware.display.DisplayManager
-import android.os.Binder
-import android.os.IBinder
-import android.testing.AndroidTestingRunner
-import android.view.Display
-import android.window.ScreenCapture.ScreenshotHardwareBuffer
-import com.android.systemui.SysuiTestCase
-import com.android.systemui.util.mockito.mock
-import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.CoroutineDispatcher
-import kotlinx.coroutines.Dispatchers
-import org.junit.Test
-import org.junit.runner.RunWith
-
-/**
- * Test the logic within ImageCaptureImpl
- */
-@RunWith(AndroidTestingRunner::class)
-class ImageCaptureImplTest : SysuiTestCase() {
-    private val displayManager = mock<DisplayManager>()
-    private val atmService = mock<IActivityTaskManager>()
-    private val capture = TestableImageCaptureImpl(
-        displayManager,
-        atmService,
-        Dispatchers.Unconfined)
-
-    @Test
-    fun captureDisplayWithCrop() {
-        capture.captureDisplay(Display.DEFAULT_DISPLAY, Rect(1, 2, 3, 4))
-        assertThat(capture.token).isNotNull()
-        assertThat(capture.width!!).isEqualTo(2)
-        assertThat(capture.height!!).isEqualTo(2)
-        assertThat(capture.crop!!).isEqualTo(Rect(1, 2, 3, 4))
-    }
-
-    @Test
-    fun captureDisplayWithNullCrop() {
-        capture.captureDisplay(Display.DEFAULT_DISPLAY, null)
-        assertThat(capture.token).isNotNull()
-        assertThat(capture.width!!).isEqualTo(0)
-        assertThat(capture.height!!).isEqualTo(0)
-        assertThat(capture.crop!!).isEqualTo(Rect())
-    }
-
-    class TestableImageCaptureImpl(
-        displayManager: DisplayManager,
-        atmService: IActivityTaskManager,
-        bgDispatcher: CoroutineDispatcher
-    ) :
-        ImageCaptureImpl(displayManager, atmService, bgDispatcher) {
-
-        var token: IBinder? = null
-        var width: Int? = null
-        var height: Int? = null
-        var crop: Rect? = null
-
-        override fun physicalDisplayToken(displayId: Int): IBinder {
-            return Binder()
-        }
-
-        override fun captureDisplay(displayToken: IBinder, width: Int, height: Int, crop: Rect):
-                ScreenshotHardwareBuffer {
-            this.token = displayToken
-            this.width = width
-            this.height = height
-            this.crop = crop
-            return ScreenshotHardwareBuffer(
-                null,
-                null,
-                false,
-                false
-            )
-        }
-    }
-}
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 28093074303d..607e575e7a91 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -196,6 +196,7 @@ import android.util.ArraySet;
 import android.util.DisplayMetrics;
 import android.util.DisplayUtils;
 import android.util.IntArray;
+import android.util.Pair;
 import android.util.RotationUtils;
 import android.util.Size;
 import android.util.Slog;
@@ -4898,20 +4899,19 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
             return null;
         }
 
-        final ScreenRotationAnimation screenRotationAnimation =
-                mWmService.mRoot.getDisplayContent(DEFAULT_DISPLAY).getRotationAnimation();
-        final boolean inRotation = screenRotationAnimation != null &&
-                screenRotationAnimation.isAnimating();
-        if (DEBUG_SCREENSHOT && inRotation) Slog.v(TAG_WM, "Taking screenshot while rotating");
+        Pair<ScreenCapture.ScreenCaptureListener, ScreenCapture.ScreenshotSync> syncScreenCapture =
+                ScreenCapture.createSyncCaptureListener();
+
+        getBounds(mTmpRect);
+        mTmpRect.offsetTo(0, 0);
+        ScreenCapture.LayerCaptureArgs args =
+                new ScreenCapture.LayerCaptureArgs.Builder(getSurfaceControl())
+                        .setSourceCrop(mTmpRect).build();
+
+        ScreenCapture.captureLayers(args, syncScreenCapture.first);
 
-        // Send invalid rect and no width and height since it will screenshot the entire display.
-        final IBinder displayToken = SurfaceControl.getInternalDisplayToken();
-        final ScreenCapture.DisplayCaptureArgs captureArgs =
-                new ScreenCapture.DisplayCaptureArgs.Builder(displayToken)
-                        .setUseIdentityTransform(inRotation)
-                        .build();
         final ScreenCapture.ScreenshotHardwareBuffer screenshotBuffer =
-                ScreenCapture.captureDisplay(captureArgs);
+                syncScreenCapture.second.get();
         final Bitmap bitmap = screenshotBuffer == null ? null : screenshotBuffer.asBitmap();
         if (bitmap == null) {
             Slog.w(TAG_WM, "Failed to take screenshot");