Merge "Supporting multi-display for takeScreenshot()" into rvc-dev
diff --git a/core/java/android/accessibilityservice/AccessibilityService.java b/core/java/android/accessibilityservice/AccessibilityService.java
index f3759fd..d4f5112 100644
--- a/core/java/android/accessibilityservice/AccessibilityService.java
+++ b/core/java/android/accessibilityservice/AccessibilityService.java
@@ -47,7 +47,6 @@
 import android.util.SparseArray;
 import android.view.Display;
 import android.view.KeyEvent;
-import android.view.SurfaceControl;
 import android.view.SurfaceView;
 import android.view.WindowManager;
 import android.view.WindowManagerImpl;
@@ -1981,8 +1980,6 @@
      * to declare the capability to take screenshot by setting the
      * {@link android.R.styleable#AccessibilityService_canTakeScreenshot}
      * property in its meta-data. For details refer to {@link #SERVICE_META_DATA}.
-     * This API only will support {@link Display#DEFAULT_DISPLAY} until {@link SurfaceControl}
-     * supports non-default displays.
      * </p>
      *
      * @param displayId The logic display id, must be {@link Display#DEFAULT_DISPLAY} for
@@ -1990,18 +1987,11 @@
      * @param executor Executor on which to run the callback.
      * @param callback The callback invoked when taking screenshot has succeeded or failed.
      *                 See {@link TakeScreenshotCallback} for details.
-     *
-     * @throws IllegalArgumentException if displayId is not {@link Display#DEFAULT_DISPLAY}.
      */
     public void takeScreenshot(int displayId, @NonNull @CallbackExecutor Executor executor,
             @NonNull TakeScreenshotCallback callback) {
         Preconditions.checkNotNull(executor, "executor cannot be null");
         Preconditions.checkNotNull(callback, "callback cannot be null");
-
-        if (displayId != Display.DEFAULT_DISPLAY) {
-            throw new IllegalArgumentException("DisplayId isn't the default display");
-        }
-
         final IAccessibilityServiceConnection connection =
                 AccessibilityInteractionClient.getInstance().getConnection(
                         mConnectionId);
diff --git a/core/java/android/hardware/display/DisplayManagerInternal.java b/core/java/android/hardware/display/DisplayManagerInternal.java
index ea2b9e7..571537c 100644
--- a/core/java/android/hardware/display/DisplayManagerInternal.java
+++ b/core/java/android/hardware/display/DisplayManagerInternal.java
@@ -72,6 +72,15 @@
     public abstract SurfaceControl.ScreenshotGraphicBuffer screenshot(int displayId);
 
     /**
+     * Take a screenshot without secure layer of the specified display and return a buffer.
+     *
+     * @param displayId The display id to take the screenshot of.
+     * @return The buffer or null if we have failed.
+     */
+    public abstract SurfaceControl.ScreenshotGraphicBuffer screenshotWithoutSecureLayer(
+            int displayId);
+
+    /**
      * Returns information about the specified logical display.
      *
      * @param displayId The logical display id.
diff --git a/services/accessibility/java/com/android/server/accessibility/AbstractAccessibilityServiceConnection.java b/services/accessibility/java/com/android/server/accessibility/AbstractAccessibilityServiceConnection.java
index c8b6d8df..132b692 100644
--- a/services/accessibility/java/com/android/server/accessibility/AbstractAccessibilityServiceConnection.java
+++ b/services/accessibility/java/com/android/server/accessibility/AbstractAccessibilityServiceConnection.java
@@ -43,12 +43,10 @@
 import android.content.pm.ParceledListSlice;
 import android.graphics.GraphicBuffer;
 import android.graphics.ParcelableColorSpace;
-import android.graphics.Point;
-import android.graphics.Rect;
 import android.graphics.Region;
 import android.hardware.HardwareBuffer;
 import android.hardware.display.DisplayManager;
-import android.hardware.display.DisplayManagerGlobal;
+import android.hardware.display.DisplayManagerInternal;
 import android.os.Binder;
 import android.os.Build;
 import android.os.Bundle;
@@ -66,7 +64,6 @@
 import android.view.Display;
 import android.view.KeyEvent;
 import android.view.MagnificationSpec;
-import android.view.SurfaceControl;
 import android.view.SurfaceControl.ScreenshotGraphicBuffer;
 import android.view.View;
 import android.view.WindowInfo;
@@ -1010,53 +1007,29 @@
             return;
         }
 
-        final Display display = DisplayManagerGlobal.getInstance()
-                .getRealDisplay(displayId);
-        if (display == null) {
-            sendScreenshotFailure(AccessibilityService.ERROR_TAKE_SCREENSHOT_INVALID_DISPLAY,
-                    callback);
+        // Private virtual displays are created by the ap and is not allowed to access by other
+        // aps.  We assume the contents on this display should not be captured.
+        final DisplayManager displayManager =
+                (DisplayManager) mContext.getSystemService(Context.DISPLAY_SERVICE);
+        final Display display = displayManager.getDisplay(displayId);
+        if ((display == null) || (display.getType() == Display.TYPE_VIRTUAL
+                && (display.getFlags() & Display.FLAG_PRIVATE) != 0)) {
+            sendScreenshotFailure(
+                    AccessibilityService.ERROR_TAKE_SCREENSHOT_INVALID_DISPLAY, callback);
             return;
         }
 
-        sendScreenshotSuccess(display, callback);
-    }
-
-    private ScreenshotGraphicBuffer takeScreenshotBuffer(Display display) {
-        final Point displaySize = new Point();
-        // TODO (b/145893483): calling new API with the display as a parameter
-        // when surface control supported.
-        final IBinder token = SurfaceControl.getInternalDisplayToken();
-        final Rect crop = new Rect(0, 0, displaySize.x, displaySize.y);
-        final int rotation = display.getRotation();
-        display.getRealSize(displaySize);
-
-        return SurfaceControl.screenshotToBuffer(token, crop, displaySize.x, displaySize.y,
-                false, rotation);
-    }
-
-    private void sendScreenshotSuccess(Display display, RemoteCallback callback) {
         final long identity = Binder.clearCallingIdentity();
         try {
             mMainHandler.post(PooledLambda.obtainRunnable((nonArg) -> {
-                final ScreenshotGraphicBuffer screenshotBuffer = takeScreenshotBuffer(display);
-                final GraphicBuffer graphicBuffer = screenshotBuffer.getGraphicBuffer();
-                try (HardwareBuffer hardwareBuffer =
-                             HardwareBuffer.createFromGraphicBuffer(graphicBuffer)) {
-                    final ParcelableColorSpace colorSpace =
-                            new ParcelableColorSpace(screenshotBuffer.getColorSpace());
-
-                    final Bundle payload = new Bundle();
-                    payload.putInt(KEY_ACCESSIBILITY_SCREENSHOT_STATUS,
-                            AccessibilityService.TAKE_SCREENSHOT_SUCCESS);
-                    payload.putParcelable(KEY_ACCESSIBILITY_SCREENSHOT_HARDWAREBUFFER,
-                            hardwareBuffer);
-                    payload.putParcelable(KEY_ACCESSIBILITY_SCREENSHOT_COLORSPACE, colorSpace);
-                    payload.putLong(KEY_ACCESSIBILITY_SCREENSHOT_TIMESTAMP,
-                            SystemClock.uptimeMillis());
-
-                    // Send back the result.
-                    callback.sendResult(payload);
-                    hardwareBuffer.close();
+                final ScreenshotGraphicBuffer screenshotBuffer = LocalServices
+                        .getService(DisplayManagerInternal.class)
+                        .screenshotWithoutSecureLayer(displayId);
+                if (screenshotBuffer != null) {
+                    sendScreenshotSuccess(screenshotBuffer, callback);
+                } else {
+                    sendScreenshotFailure(
+                            AccessibilityService.ERROR_TAKE_SCREENSHOT_INVALID_DISPLAY, callback);
                 }
             }, null).recycleOnUse());
         } finally {
@@ -1064,6 +1037,29 @@
         }
     }
 
+    private void sendScreenshotSuccess(ScreenshotGraphicBuffer screenshotBuffer,
+            RemoteCallback callback) {
+        final GraphicBuffer graphicBuffer = screenshotBuffer.getGraphicBuffer();
+        try (HardwareBuffer hardwareBuffer =
+                     HardwareBuffer.createFromGraphicBuffer(graphicBuffer)) {
+            final ParcelableColorSpace colorSpace =
+                    new ParcelableColorSpace(screenshotBuffer.getColorSpace());
+
+            final Bundle payload = new Bundle();
+            payload.putInt(KEY_ACCESSIBILITY_SCREENSHOT_STATUS,
+                    AccessibilityService.TAKE_SCREENSHOT_SUCCESS);
+            payload.putParcelable(KEY_ACCESSIBILITY_SCREENSHOT_HARDWAREBUFFER,
+                    hardwareBuffer);
+            payload.putParcelable(KEY_ACCESSIBILITY_SCREENSHOT_COLORSPACE, colorSpace);
+            payload.putLong(KEY_ACCESSIBILITY_SCREENSHOT_TIMESTAMP,
+                    SystemClock.uptimeMillis());
+
+            // Send back the result.
+            callback.sendResult(payload);
+            hardwareBuffer.close();
+        }
+    }
+
     private void sendScreenshotFailure(@AccessibilityService.ScreenshotErrorCode int errorCode,
             RemoteCallback callback) {
         mMainHandler.post(PooledLambda.obtainRunnable((nonArg) -> {
diff --git a/services/core/java/com/android/server/display/DisplayManagerService.java b/services/core/java/com/android/server/display/DisplayManagerService.java
index 3afbf66..d527a80 100644
--- a/services/core/java/com/android/server/display/DisplayManagerService.java
+++ b/services/core/java/com/android/server/display/DisplayManagerService.java
@@ -1362,7 +1362,8 @@
         return null;
     }
 
-    private SurfaceControl.ScreenshotGraphicBuffer screenshotInternal(int displayId) {
+    private SurfaceControl.ScreenshotGraphicBuffer screenshotInternal(int displayId,
+            boolean captureSecureLayer) {
         synchronized (mSyncRoot) {
             final IBinder token = getDisplayToken(displayId);
             if (token == null) {
@@ -1374,9 +1375,15 @@
             }
 
             final DisplayInfo displayInfo = logicalDisplay.getDisplayInfoLocked();
-            return SurfaceControl.screenshotToBufferWithSecureLayersUnsafe(token, new Rect(),
-                    displayInfo.getNaturalWidth(), displayInfo.getNaturalHeight(),
-                    false /* useIdentityTransform */, 0 /* rotation */);
+            if (captureSecureLayer) {
+                return SurfaceControl.screenshotToBufferWithSecureLayersUnsafe(token, new Rect(),
+                        displayInfo.getNaturalWidth(), displayInfo.getNaturalHeight(),
+                        false /* useIdentityTransform */, 0 /* rotation */);
+            } else {
+                return SurfaceControl.screenshotToBuffer(token, new Rect(),
+                        displayInfo.getNaturalWidth(), displayInfo.getNaturalHeight(),
+                        false /* useIdentityTransform */, 0 /* rotation */);
+            }
         }
     }
 
@@ -2494,7 +2501,12 @@
 
         @Override
         public SurfaceControl.ScreenshotGraphicBuffer screenshot(int displayId) {
-            return screenshotInternal(displayId);
+            return screenshotInternal(displayId, true);
+        }
+
+        @Override
+        public SurfaceControl.ScreenshotGraphicBuffer screenshotWithoutSecureLayer(int displayId) {
+            return screenshotInternal(displayId, false);
         }
 
         @Override
diff --git a/services/tests/servicestests/src/com/android/server/accessibility/AbstractAccessibilityServiceConnectionTest.java b/services/tests/servicestests/src/com/android/server/accessibility/AbstractAccessibilityServiceConnectionTest.java
index 7bf1d98..0445bff 100644
--- a/services/tests/servicestests/src/com/android/server/accessibility/AbstractAccessibilityServiceConnectionTest.java
+++ b/services/tests/servicestests/src/com/android/server/accessibility/AbstractAccessibilityServiceConnectionTest.java
@@ -16,7 +16,10 @@
 
 package com.android.server.accessibility;
 
+import static android.accessibilityservice.AccessibilityService.ERROR_TAKE_SCREENSHOT_INVALID_DISPLAY;
+import static android.accessibilityservice.AccessibilityService.ERROR_TAKE_SCREENSHOT_NO_ACCESSIBILITY_ACCESS;
 import static android.accessibilityservice.AccessibilityService.GLOBAL_ACTION_HOME;
+import static android.accessibilityservice.AccessibilityService.KEY_ACCESSIBILITY_SCREENSHOT_STATUS;
 import static android.accessibilityservice.AccessibilityServiceInfo.CAPABILITY_CAN_CONTROL_MAGNIFICATION;
 import static android.accessibilityservice.AccessibilityServiceInfo.CAPABILITY_CAN_PERFORM_GESTURES;
 import static android.accessibilityservice.AccessibilityServiceInfo.CAPABILITY_CAN_REQUEST_FILTER_KEY_EVENTS;
@@ -69,6 +72,7 @@
 import android.content.pm.ResolveInfo;
 import android.content.pm.ServiceInfo;
 import android.graphics.Region;
+import android.hardware.display.DisplayManager;
 import android.os.Build;
 import android.os.Handler;
 import android.os.IBinder;
@@ -93,6 +97,7 @@
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
+import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 import org.mockito.Spy;
 
@@ -162,6 +167,7 @@
     @Mock private IAccessibilityInteractionConnectionCallback mMockCallback;
     @Mock private FingerprintGestureDispatcher mMockFingerprintGestureDispatcher;
     @Mock private MagnificationController mMockMagnificationController;
+    @Mock private RemoteCallback.OnResultListener mMockListener;
 
     @Before
     public void setup() {
@@ -705,6 +711,38 @@
         }));
     }
 
+    @Test
+    public void takeScreenshot_NoA11yAccess_returnErrorCode() throws InterruptedException {
+        // no checkAccessibilityAccess, should return error code.
+        when(mMockSecurityPolicy.canTakeScreenshotLocked(mServiceConnection)).thenReturn(true);
+        when(mMockSecurityPolicy.checkAccessibilityAccess(mServiceConnection)).thenReturn(false);
+
+        mServiceConnection.takeScreenshot(Display.DEFAULT_DISPLAY,
+                new RemoteCallback(mMockListener));
+        mHandler.sendLastMessage();
+
+        verify(mMockListener).onResult(Mockito.argThat(
+                bundle -> ERROR_TAKE_SCREENSHOT_NO_ACCESSIBILITY_ACCESS
+                        == bundle.getInt(KEY_ACCESSIBILITY_SCREENSHOT_STATUS)));
+    }
+
+    @Test
+    public void takeScreenshot_invalidDisplay_returnErrorCode() throws InterruptedException {
+        when(mMockSecurityPolicy.canTakeScreenshotLocked(mServiceConnection)).thenReturn(true);
+        when(mMockSecurityPolicy.checkAccessibilityAccess(mServiceConnection)).thenReturn(true);
+
+        final DisplayManager displayManager = new DisplayManager(mMockContext);
+        when(mMockContext.getSystemService(Context.DISPLAY_SERVICE)).thenReturn(displayManager);
+
+        mServiceConnection.takeScreenshot(Display.DEFAULT_DISPLAY + 1,
+                new RemoteCallback(mMockListener));
+        mHandler.sendLastMessage();
+
+        verify(mMockListener).onResult(Mockito.argThat(
+                bundle -> ERROR_TAKE_SCREENSHOT_INVALID_DISPLAY
+                        == bundle.getInt(KEY_ACCESSIBILITY_SCREENSHOT_STATUS)));
+    }
+
     private void updateServiceInfo(AccessibilityServiceInfo serviceInfo, int eventType,
             int feedbackType, int flags, String[] packageNames, int notificationTimeout) {
         serviceInfo.eventTypes = eventType;