Implementing new API : takeScreenshot()

Implementing the new API for A11y services to take the screenshot
of the specified display.
Bitmap takeScreenshot(int displayId)
1. The main codes is moved from the UiAutomationConnection class and
let the UiAutomation change to use this API.
2. Add a capability at metadata to check whether the A11y services could use this
API or not.
3. This API is a async one for A11y services, but is a sync one for
UiAutomation.

Bug: 10931661
Test: a11y CTS & unit tests
Change-Id: I478bd93c60d4742bef20ae0e423ca9de1bda55d2
diff --git a/api/current.txt b/api/current.txt
index e166c14..bc0099a 100644
--- a/api/current.txt
+++ b/api/current.txt
@@ -386,6 +386,7 @@
     field public static final int canRequestFingerprintGestures = 16844109; // 0x101054d
     field public static final int canRequestTouchExplorationMode = 16843735; // 0x10103d7
     field public static final int canRetrieveWindowContent = 16843653; // 0x1010385
+    field public static final int canTakeScreenshot = 16844304; // 0x1010610
     field public static final int candidatesTextStyleSpans = 16843312; // 0x1010230
     field public static final int cantSaveState = 16844142; // 0x101056e
     field @Deprecated public static final int capitalize = 16843113; // 0x1010169
@@ -2864,6 +2865,7 @@
     method protected void onServiceConnected();
     method public final boolean performGlobalAction(int);
     method public final void setServiceInfo(android.accessibilityservice.AccessibilityServiceInfo);
+    method public boolean takeScreenshot(int, @NonNull java.util.concurrent.Executor, @NonNull java.util.function.Consumer<android.graphics.Bitmap>);
     field public static final int GESTURE_SWIPE_DOWN = 2; // 0x2
     field public static final int GESTURE_SWIPE_DOWN_AND_LEFT = 15; // 0xf
     field public static final int GESTURE_SWIPE_DOWN_AND_RIGHT = 16; // 0x10
@@ -2960,6 +2962,7 @@
     field public static final int CAPABILITY_CAN_REQUEST_FINGERPRINT_GESTURES = 64; // 0x40
     field public static final int CAPABILITY_CAN_REQUEST_TOUCH_EXPLORATION = 2; // 0x2
     field public static final int CAPABILITY_CAN_RETRIEVE_WINDOW_CONTENT = 1; // 0x1
+    field public static final int CAPABILITY_CAN_TAKE_SCREENSHOT = 128; // 0x80
     field @NonNull public static final android.os.Parcelable.Creator<android.accessibilityservice.AccessibilityServiceInfo> CREATOR;
     field public static final int DEFAULT = 1; // 0x1
     field public static final int FEEDBACK_ALL_MASK = -1; // 0xffffffff
diff --git a/core/java/android/accessibilityservice/AccessibilityService.java b/core/java/android/accessibilityservice/AccessibilityService.java
index e46840c..0bd8ce6 100644
--- a/core/java/android/accessibilityservice/AccessibilityService.java
+++ b/core/java/android/accessibilityservice/AccessibilityService.java
@@ -17,6 +17,7 @@
 package android.accessibilityservice;
 
 import android.accessibilityservice.GestureDescription.MotionEventGenerator;
+import android.annotation.CallbackExecutor;
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -26,12 +27,15 @@
 import android.content.Context;
 import android.content.Intent;
 import android.content.pm.ParceledListSlice;
+import android.graphics.Bitmap;
 import android.graphics.Region;
+import android.os.Binder;
 import android.os.Build;
 import android.os.Handler;
 import android.os.IBinder;
 import android.os.Looper;
 import android.os.Message;
+import android.os.RemoteCallback;
 import android.os.RemoteException;
 import android.util.ArrayMap;
 import android.util.Log;
@@ -48,10 +52,13 @@
 
 import com.android.internal.os.HandlerCaller;
 import com.android.internal.os.SomeArgs;
+import com.android.internal.util.Preconditions;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.List;
+import java.util.concurrent.Executor;
+import java.util.function.Consumer;
 
 /**
  * Accessibility services should only be used to assist users with disabilities in using
@@ -483,6 +490,9 @@
 
     private FingerprintGestureController mFingerprintGestureController;
 
+    /** @hide */
+    public static final String KEY_ACCESSIBILITY_SCREENSHOT = "screenshot";
+
     /**
      * Callback for {@link android.view.accessibility.AccessibilityEvent}s.
      *
@@ -1761,6 +1771,51 @@
     }
 
     /**
+     * Takes a screenshot of the specified display and returns it by {@link Bitmap.Config#HARDWARE}
+     * format.
+     * <p>
+     * <strong>Note:</strong> In order to take screenshot your service has
+     * 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}.
+     * Besides, This API is only supported for default display now
+     * {@link Display#DEFAULT_DISPLAY}.
+     * </p>
+     *
+     * @param displayId The logic display id, must be {@link Display#DEFAULT_DISPLAY} for
+     *                  default display.
+     * @param executor Executor on which to run the callback.
+     * @param callback The callback invoked when the taking screenshot is done.
+     *
+     * @return {@code true} if the taking screenshot accepted, {@code false} if not.
+     */
+    public boolean takeScreenshot(int displayId, @NonNull @CallbackExecutor Executor executor,
+            @NonNull Consumer<Bitmap> callback) {
+        Preconditions.checkNotNull(executor, "executor cannot be null");
+        Preconditions.checkNotNull(callback, "callback cannot be null");
+        final IAccessibilityServiceConnection connection =
+                AccessibilityInteractionClient.getInstance().getConnection(
+                        mConnectionId);
+        if (connection == null) {
+            return false;
+        }
+        try {
+            connection.takeScreenshotWithCallback(displayId, new RemoteCallback((result) -> {
+                final Bitmap screenshot = result.getParcelable(KEY_ACCESSIBILITY_SCREENSHOT);
+                final long identity = Binder.clearCallingIdentity();
+                try {
+                    executor.execute(() -> callback.accept(screenshot));
+                } finally {
+                    Binder.restoreCallingIdentity(identity);
+                }
+            }));
+        } catch (RemoteException re) {
+            throw new RuntimeException(re);
+        }
+        return true;
+    }
+
+    /**
      * Implement to return the implementation of the internal accessibility
      * service interface.
      */
diff --git a/core/java/android/accessibilityservice/AccessibilityServiceInfo.java b/core/java/android/accessibilityservice/AccessibilityServiceInfo.java
index 5e2c1fa..12f2c3b 100644
--- a/core/java/android/accessibilityservice/AccessibilityServiceInfo.java
+++ b/core/java/android/accessibilityservice/AccessibilityServiceInfo.java
@@ -86,6 +86,7 @@
  * @attr ref android.R.styleable#AccessibilityService_settingsActivity
  * @attr ref android.R.styleable#AccessibilityService_nonInteractiveUiTimeout
  * @attr ref android.R.styleable#AccessibilityService_interactiveUiTimeout
+ * @attr ref android.R.styleable#AccessibilityService_canTakeScreenshot
  * @see AccessibilityService
  * @see android.view.accessibility.AccessibilityEvent
  * @see android.view.accessibility.AccessibilityManager
@@ -136,6 +137,12 @@
      */
     public static final int CAPABILITY_CAN_REQUEST_FINGERPRINT_GESTURES = 0x00000040;
 
+    /**
+     * Capability: This accessibility service can take screenshot.
+     * @see android.R.styleable#AccessibilityService_canTakeScreenshot
+     */
+    public static final int CAPABILITY_CAN_TAKE_SCREENSHOT = 0x00000080;
+
     private static SparseArray<CapabilityInfo> sAvailableCapabilityInfos;
 
     /**
@@ -625,6 +632,10 @@
                     .AccessibilityService_canRequestFingerprintGestures, false)) {
                 mCapabilities |= CAPABILITY_CAN_REQUEST_FINGERPRINT_GESTURES;
             }
+            if (asAttributes.getBoolean(com.android.internal.R.styleable
+                    .AccessibilityService_canTakeScreenshot, false)) {
+                mCapabilities |= CAPABILITY_CAN_TAKE_SCREENSHOT;
+            }
             TypedValue peekedValue = asAttributes.peekValue(
                     com.android.internal.R.styleable.AccessibilityService_description);
             if (peekedValue != null) {
@@ -794,6 +805,7 @@
      * @see #CAPABILITY_CAN_REQUEST_FILTER_KEY_EVENTS
      * @see #CAPABILITY_CAN_CONTROL_MAGNIFICATION
      * @see #CAPABILITY_CAN_PERFORM_GESTURES
+     * @see #CAPABILITY_CAN_TAKE_SCREENSHOT
      */
     public int getCapabilities() {
         return mCapabilities;
@@ -810,6 +822,7 @@
      * @see #CAPABILITY_CAN_REQUEST_FILTER_KEY_EVENTS
      * @see #CAPABILITY_CAN_CONTROL_MAGNIFICATION
      * @see #CAPABILITY_CAN_PERFORM_GESTURES
+     * @see #CAPABILITY_CAN_TAKE_SCREENSHOT
      *
      * @hide
      */
@@ -1253,6 +1266,8 @@
                 return "CAPABILITY_CAN_PERFORM_GESTURES";
             case CAPABILITY_CAN_REQUEST_FINGERPRINT_GESTURES:
                 return "CAPABILITY_CAN_REQUEST_FINGERPRINT_GESTURES";
+            case CAPABILITY_CAN_TAKE_SCREENSHOT:
+                return "CAPABILITY_CAN_TAKE_SCREENSHOT";
             default:
                 return "UNKNOWN";
         }
@@ -1314,6 +1329,10 @@
                     new CapabilityInfo(CAPABILITY_CAN_PERFORM_GESTURES,
                             R.string.capability_title_canPerformGestures,
                             R.string.capability_desc_canPerformGestures));
+            sAvailableCapabilityInfos.put(CAPABILITY_CAN_TAKE_SCREENSHOT,
+                    new CapabilityInfo(CAPABILITY_CAN_TAKE_SCREENSHOT,
+                            R.string.capability_title_canTakeScreenshot,
+                            R.string.capability_desc_canTakeScreenshot));
             if ((context == null) || fingerprintAvailable(context)) {
                 sAvailableCapabilityInfos.put(CAPABILITY_CAN_REQUEST_FINGERPRINT_GESTURES,
                         new CapabilityInfo(CAPABILITY_CAN_REQUEST_FINGERPRINT_GESTURES,
diff --git a/core/java/android/accessibilityservice/IAccessibilityServiceConnection.aidl b/core/java/android/accessibilityservice/IAccessibilityServiceConnection.aidl
index 932d0a4..4ea5c62 100644
--- a/core/java/android/accessibilityservice/IAccessibilityServiceConnection.aidl
+++ b/core/java/android/accessibilityservice/IAccessibilityServiceConnection.aidl
@@ -18,8 +18,10 @@
 
 import android.accessibilityservice.AccessibilityServiceInfo;
 import android.content.pm.ParceledListSlice;
+import android.graphics.Bitmap;
 import android.graphics.Region;
 import android.os.Bundle;
+import android.os.RemoteCallback;
 import android.view.MagnificationSpec;
 import android.view.MotionEvent;
 import android.view.accessibility.AccessibilityNodeInfo;
@@ -104,4 +106,8 @@
     IBinder getOverlayWindowToken(int displayid);
 
     int getWindowIdForLeashToken(IBinder token);
+
+    Bitmap takeScreenshot(int displayId);
+
+    void takeScreenshotWithCallback(int displayId, in RemoteCallback callback);
 }
diff --git a/core/java/android/app/IUiAutomationConnection.aidl b/core/java/android/app/IUiAutomationConnection.aidl
index 8c3180b..80ba464 100644
--- a/core/java/android/app/IUiAutomationConnection.aidl
+++ b/core/java/android/app/IUiAutomationConnection.aidl
@@ -39,7 +39,6 @@
     boolean injectInputEvent(in InputEvent event, boolean sync);
     void syncInputTransactions();
     boolean setRotation(int rotation);
-    Bitmap takeScreenshot(in Rect crop, int rotation);
     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 18a3e6e..2579bd1 100644
--- a/core/java/android/app/UiAutomation.java
+++ b/core/java/android/app/UiAutomation.java
@@ -27,10 +27,7 @@
 import android.annotation.TestApi;
 import android.compat.annotation.UnsupportedAppUsage;
 import android.graphics.Bitmap;
-import android.graphics.Point;
-import android.graphics.Rect;
 import android.graphics.Region;
-import android.hardware.display.DisplayManagerGlobal;
 import android.os.Build;
 import android.os.Handler;
 import android.os.HandlerThread;
@@ -831,39 +828,20 @@
     }
 
     /**
-     * Takes a screenshot.
+     * Takes a screenshot of the default display and returns it by {@link Bitmap.Config#HARDWARE}
+     * format.
      *
      * @return The screenshot bitmap on success, null otherwise.
      */
     public Bitmap takeScreenshot() {
+        final int connectionId;
         synchronized (mLock) {
             throwIfNotConnectedLocked();
+            connectionId = mConnectionId;
         }
-        Display display = DisplayManagerGlobal.getInstance()
-                .getRealDisplay(Display.DEFAULT_DISPLAY);
-        Point displaySize = new Point();
-        display.getRealSize(displaySize);
-
-        int rotation = display.getRotation();
-
-        // Take the screenshot
-        Bitmap screenShot = null;
-        try {
-            // Calling out without a lock held.
-            screenShot = mUiAutomationConnection.takeScreenshot(
-                    new Rect(0, 0, displaySize.x, displaySize.y), rotation);
-            if (screenShot == null) {
-                return null;
-            }
-        } catch (RemoteException re) {
-            Log.e(LOG_TAG, "Error while taking screnshot!", re);
-            return null;
-        }
-
-        // Optimization
-        screenShot.setHasAlpha(false);
-
-        return screenShot;
+        // Calling out without a lock held.
+        return AccessibilityInteractionClient.getInstance()
+                .takeScreenshot(connectionId, Display.DEFAULT_DISPLAY);
     }
 
     /**
diff --git a/core/java/android/app/UiAutomationConnection.java b/core/java/android/app/UiAutomationConnection.java
index 82e9881..4fb9743 100644
--- a/core/java/android/app/UiAutomationConnection.java
+++ b/core/java/android/app/UiAutomationConnection.java
@@ -21,8 +21,6 @@
 import android.annotation.Nullable;
 import android.compat.annotation.UnsupportedAppUsage;
 import android.content.Context;
-import android.graphics.Bitmap;
-import android.graphics.Rect;
 import android.hardware.input.InputManager;
 import android.os.Binder;
 import android.os.IBinder;
@@ -180,23 +178,6 @@
     }
 
     @Override
-    public Bitmap takeScreenshot(Rect crop, int rotation) {
-        synchronized (mLock) {
-            throwIfCalledByNotTrustedUidLocked();
-            throwIfShutdownLocked();
-            throwIfNotConnectedLocked();
-        }
-        final long identity = Binder.clearCallingIdentity();
-        try {
-            int width = crop.width();
-            int height = crop.height();
-            return SurfaceControl.screenshot(crop, width, height, rotation);
-        } finally {
-            Binder.restoreCallingIdentity(identity);
-        }
-    }
-
-    @Override
     public boolean clearWindowContentFrameStats(int windowId) throws RemoteException {
         synchronized (mLock) {
             throwIfCalledByNotTrustedUidLocked();
@@ -449,7 +430,8 @@
         info.setCapabilities(AccessibilityServiceInfo.CAPABILITY_CAN_RETRIEVE_WINDOW_CONTENT
                 | AccessibilityServiceInfo.CAPABILITY_CAN_REQUEST_TOUCH_EXPLORATION
                 | AccessibilityServiceInfo.CAPABILITY_CAN_REQUEST_ENHANCED_WEB_ACCESSIBILITY
-                | AccessibilityServiceInfo.CAPABILITY_CAN_REQUEST_FILTER_KEY_EVENTS);
+                | AccessibilityServiceInfo.CAPABILITY_CAN_REQUEST_FILTER_KEY_EVENTS
+                | AccessibilityServiceInfo.CAPABILITY_CAN_TAKE_SCREENSHOT);
         try {
             // Calling out with a lock held is fine since if the system
             // process is gone the client calling in will be killed.
diff --git a/core/java/android/view/accessibility/AccessibilityInteractionClient.java b/core/java/android/view/accessibility/AccessibilityInteractionClient.java
index b4c8795..b9f08ad 100644
--- a/core/java/android/view/accessibility/AccessibilityInteractionClient.java
+++ b/core/java/android/view/accessibility/AccessibilityInteractionClient.java
@@ -20,6 +20,7 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.compat.annotation.UnsupportedAppUsage;
+import android.graphics.Bitmap;
 import android.os.Binder;
 import android.os.Build;
 import android.os.Bundle;
@@ -828,6 +829,31 @@
     }
 
     /**
+     * Takes the screenshot of the specified display and returns it by bitmap format.
+     *
+     * @param connectionId The id of a connection for interacting with the system.
+     * @param displayId The logic display id, use {@link Display#DEFAULT_DISPLAY} for
+     *                  default display.
+     * @return The screenshot bitmap on success, null otherwise.
+     */
+    public Bitmap takeScreenshot(int connectionId, int displayId) {
+        Bitmap screenShot = null;
+        try {
+            IAccessibilityServiceConnection connection = getConnection(connectionId);
+            if (connection != null) {
+                screenShot = connection.takeScreenshot(displayId);
+            } else {
+                if (DEBUG) {
+                    Log.w(LOG_TAG, "No connection for connection id: " + connectionId);
+                }
+            }
+        } catch (RemoteException re) {
+            Log.w(LOG_TAG, "Error while calling remote takeScreenshot", re);
+        }
+        return screenShot;
+    }
+
+    /**
      * Clears the result state.
      */
     private void clearResultLocked() {
diff --git a/core/res/res/values/attrs.xml b/core/res/res/values/attrs.xml
index 9c08728..4475415 100644
--- a/core/res/res/values/attrs.xml
+++ b/core/res/res/values/attrs.xml
@@ -3752,6 +3752,8 @@
              </p>
          -->
         <attr name="canRequestFingerprintGestures" format="boolean" />
+        <!-- Attribute whether the accessibility service wants to be able to take screenshot. -->
+        <attr name="canTakeScreenshot" format="boolean" />
 
         <!-- Animated image of the accessibility service purpose or behavior, to help users
              understand how the service can help them.-->
diff --git a/core/res/res/values/public.xml b/core/res/res/values/public.xml
index 6cf6a68..9860830 100644
--- a/core/res/res/values/public.xml
+++ b/core/res/res/values/public.xml
@@ -3007,6 +3007,7 @@
       <public name="featureId" />
       <public name="supportsInlineSuggestions" />
       <public name="crossProfile" />
+      <public name="canTakeScreenshot"/>
     </public-group>
 
     <public-group type="drawable" first-id="0x010800b5">
diff --git a/core/res/res/values/strings.xml b/core/res/res/values/strings.xml
index a0e4064..b9fac09 100644
--- a/core/res/res/values/strings.xml
+++ b/core/res/res/values/strings.xml
@@ -794,6 +794,11 @@
     <string name="capability_desc_canCaptureFingerprintGestures">Can capture gestures performed on
         the device\'s fingerprint sensor.</string>
 
+    <!-- Title for the capability of an accessibility service to take screenshot. [CHAR LIMIT=32] -->
+    <string name="capability_title_canTakeScreenshot">Take screenshot</string>
+    <!-- Description for the capability of an accessibility service to take screenshot. [CHAR LIMIT=NONE] -->
+    <string name="capability_desc_canTakeScreenshot">Can take a screenshot of the display.</string>
+
     <!--  Permissions -->
 
     <!-- Title of an application permission, listed so the user can choose whether they want to allow the application to do this. -->
diff --git a/core/res/res/values/symbols.xml b/core/res/res/values/symbols.xml
index f01c2e8..bf97bb3 100644
--- a/core/res/res/values/symbols.xml
+++ b/core/res/res/values/symbols.xml
@@ -3806,4 +3806,9 @@
   <java-symbol type="dimen" name="waterfall_display_top_edge_size" />
   <java-symbol type="dimen" name="waterfall_display_right_edge_size" />
   <java-symbol type="dimen" name="waterfall_display_bottom_edge_size" />
+
+  <!-- Accessibility take screenshot -->
+  <java-symbol type="string" name="capability_desc_canTakeScreenshot" />
+  <java-symbol type="string" name="capability_title_canTakeScreenshot" />
+
 </resources>
diff --git a/core/tests/coretests/src/android/view/accessibility/AccessibilityServiceConnectionImpl.java b/core/tests/coretests/src/android/view/accessibility/AccessibilityServiceConnectionImpl.java
index 836cb95..f151b81 100644
--- a/core/tests/coretests/src/android/view/accessibility/AccessibilityServiceConnectionImpl.java
+++ b/core/tests/coretests/src/android/view/accessibility/AccessibilityServiceConnectionImpl.java
@@ -19,9 +19,11 @@
 import android.accessibilityservice.AccessibilityServiceInfo;
 import android.accessibilityservice.IAccessibilityServiceConnection;
 import android.content.pm.ParceledListSlice;
+import android.graphics.Bitmap;
 import android.graphics.Region;
 import android.os.Bundle;
 import android.os.IBinder;
+import android.os.RemoteCallback;
 
 /**
  * Stub implementation of IAccessibilityServiceConnection so each test doesn't need to implement
@@ -147,4 +149,10 @@
     public int getWindowIdForLeashToken(IBinder token) {
         return -1;
     }
+
+    public Bitmap takeScreenshot(int displayId) {
+        return null;
+    }
+
+    public void takeScreenshotWithCallback(int displayId, RemoteCallback callback) {}
 }
diff --git a/services/accessibility/java/com/android/server/accessibility/AbstractAccessibilityServiceConnection.java b/services/accessibility/java/com/android/server/accessibility/AbstractAccessibilityServiceConnection.java
index 0160286..c3965c4 100644
--- a/services/accessibility/java/com/android/server/accessibility/AbstractAccessibilityServiceConnection.java
+++ b/services/accessibility/java/com/android/server/accessibility/AbstractAccessibilityServiceConnection.java
@@ -36,8 +36,12 @@
 import android.content.ServiceConnection;
 import android.content.pm.PackageManager;
 import android.content.pm.ParceledListSlice;
+import android.graphics.Bitmap;
+import android.graphics.Point;
+import android.graphics.Rect;
 import android.graphics.Region;
 import android.hardware.display.DisplayManager;
+import android.hardware.display.DisplayManagerGlobal;
 import android.os.Binder;
 import android.os.Build;
 import android.os.Bundle;
@@ -54,6 +58,7 @@
 import android.view.Display;
 import android.view.KeyEvent;
 import android.view.MagnificationSpec;
+import android.view.SurfaceControl;
 import android.view.View;
 import android.view.WindowInfo;
 import android.view.accessibility.AccessibilityCache;
@@ -90,6 +95,7 @@
     private static final String LOG_TAG = "AbstractAccessibilityServiceConnection";
     private static final int WAIT_WINDOWS_TIMEOUT_MILLIS = 5000;
 
+    protected static final String TAKE_SCREENSHOT = "takeScreenshot";
     protected final Context mContext;
     protected final SystemSupport mSystemSupport;
     protected final WindowManagerInternal mWindowManagerService;
@@ -934,6 +940,54 @@
         mInvocationHandler.setSoftKeyboardCallbackEnabled(enabled);
     }
 
+    @Nullable
+    @Override
+    public Bitmap takeScreenshot(int displayId) {
+        synchronized (mLock) {
+            if (!hasRightsToCurrentUserLocked()) {
+                return null;
+            }
+
+            if (!mSecurityPolicy.canTakeScreenshotLocked(this)) {
+                return null;
+            }
+        }
+
+        if (!mSecurityPolicy.checkAccessibilityAccess(this)) {
+            return null;
+        }
+
+        final Display display = DisplayManagerGlobal.getInstance()
+                .getRealDisplay(displayId);
+        if (display == null) {
+            return null;
+        }
+        final Point displaySize = new Point();
+        display.getRealSize(displaySize);
+
+        final int rotation = display.getRotation();
+        Bitmap screenShot = null;
+
+        final long identity = Binder.clearCallingIdentity();
+        try {
+            final Rect crop = new Rect(0, 0, displaySize.x, displaySize.y);
+            // TODO (b/145893483): calling new API with the display as a parameter
+            // when surface control supported.
+            screenShot = SurfaceControl.screenshot(crop, displaySize.x, displaySize.y,
+                    rotation);
+            if (screenShot != null) {
+                // Optimization for telling the bitmap that all of the pixels are known to be
+                // opaque (false). This is meant as a drawing hint, as in some cases a bitmap
+                // that is known to be opaque can take a faster drawing case than one that may
+                // have non-opaque per-pixel alpha values.
+                screenShot.setHasAlpha(false);
+            }
+        } finally {
+            Binder.restoreCallingIdentity(identity);
+        }
+        return screenShot;
+    }
+
     @Override
     public void dump(FileDescriptor fd, final PrintWriter pw, String[] args) {
         if (!DumpUtils.checkDumpPermission(mContext, LOG_TAG, pw)) return;
diff --git a/services/accessibility/java/com/android/server/accessibility/AccessibilitySecurityPolicy.java b/services/accessibility/java/com/android/server/accessibility/AccessibilitySecurityPolicy.java
index 0151a52..7dbec7c 100644
--- a/services/accessibility/java/com/android/server/accessibility/AccessibilitySecurityPolicy.java
+++ b/services/accessibility/java/com/android/server/accessibility/AccessibilitySecurityPolicy.java
@@ -326,6 +326,19 @@
     }
 
     /**
+     * Checks if a service can take screenshot.
+     *
+     * @param service The service requesting access
+     *
+     * @return Whether ot not the service may take screenshot
+     */
+    public boolean canTakeScreenshotLocked(
+            @NonNull AbstractAccessibilityServiceConnection service) {
+        return (service.getCapabilities()
+                & AccessibilityServiceInfo.CAPABILITY_CAN_TAKE_SCREENSHOT) != 0;
+    }
+
+    /**
      * Returns the parent userId of the profile according to the specified userId.
      *
      * @param userId The userId to check
diff --git a/services/accessibility/java/com/android/server/accessibility/AccessibilityServiceConnection.java b/services/accessibility/java/com/android/server/accessibility/AccessibilityServiceConnection.java
index 37ac3ec..76cf98c 100644
--- a/services/accessibility/java/com/android/server/accessibility/AccessibilityServiceConnection.java
+++ b/services/accessibility/java/com/android/server/accessibility/AccessibilityServiceConnection.java
@@ -16,6 +16,8 @@
 
 package com.android.server.accessibility;
 
+import static android.accessibilityservice.AccessibilityService.KEY_ACCESSIBILITY_SCREENSHOT;
+
 import static com.android.internal.util.function.pooled.PooledLambda.obtainMessage;
 
 import android.Manifest;
@@ -25,16 +27,20 @@
 import android.content.Context;
 import android.content.Intent;
 import android.content.pm.ParceledListSlice;
+import android.graphics.Bitmap;
 import android.os.Binder;
+import android.os.Bundle;
 import android.os.Handler;
 import android.os.IBinder;
 import android.os.Process;
+import android.os.RemoteCallback;
 import android.os.RemoteException;
 import android.os.UserHandle;
 import android.provider.Settings;
 import android.util.Slog;
 import android.view.Display;
 
+import com.android.internal.util.function.pooled.PooledLambda;
 import com.android.server.inputmethod.InputMethodManagerInternal;
 import com.android.server.wm.ActivityTaskManagerInternal;
 import com.android.server.wm.WindowManagerInternal;
@@ -419,4 +425,15 @@
             }
         }
     }
+
+    @Override
+    public void takeScreenshotWithCallback(int displayId, RemoteCallback callback) {
+        mMainHandler.post(PooledLambda.obtainRunnable((nonArg) -> {
+            final Bitmap screenshot = super.takeScreenshot(displayId);
+            // Send back the result.
+            final Bundle payload = new Bundle();
+            payload.putParcelable(KEY_ACCESSIBILITY_SCREENSHOT, screenshot);
+            callback.sendResult(payload);
+        }, null).recycleOnUse());
+    }
 }
diff --git a/services/accessibility/java/com/android/server/accessibility/UiAutomationManager.java b/services/accessibility/java/com/android/server/accessibility/UiAutomationManager.java
index 7a8a112..5d9af26 100644
--- a/services/accessibility/java/com/android/server/accessibility/UiAutomationManager.java
+++ b/services/accessibility/java/com/android/server/accessibility/UiAutomationManager.java
@@ -25,6 +25,7 @@
 import android.os.Handler;
 import android.os.IBinder;
 import android.os.IBinder.DeathRecipient;
+import android.os.RemoteCallback;
 import android.os.RemoteException;
 import android.util.Slog;
 import android.view.Display;
@@ -325,5 +326,8 @@
 
         @Override
         public void onFingerprintGesture(int gesture) {}
+
+        @Override
+        public void takeScreenshotWithCallback(int displayId, RemoteCallback callback) {}
     }
 }
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 c223f13..e1e9b7e 100644
--- a/services/tests/servicestests/src/com/android/server/accessibility/AbstractAccessibilityServiceConnectionTest.java
+++ b/services/tests/servicestests/src/com/android/server/accessibility/AbstractAccessibilityServiceConnectionTest.java
@@ -75,6 +75,7 @@
 import android.os.IPowerManager;
 import android.os.PowerManager;
 import android.os.Process;
+import android.os.RemoteCallback;
 import android.os.RemoteException;
 import android.testing.DexmakerShareClassLoaderRule;
 import android.view.Display;
@@ -693,6 +694,18 @@
         assertThat(result, is(false));
     }
 
+    @Test
+    public void takeScreenshot_returnNull() {
+        // no canTakeScreenshot, should return null.
+        when(mMockSecurityPolicy.canTakeScreenshotLocked(mServiceConnection)).thenReturn(false);
+        assertThat(mServiceConnection.takeScreenshot(Display.DEFAULT_DISPLAY), is(nullValue()));
+
+        // no checkAccessibilityAccess, should return null.
+        when(mMockSecurityPolicy.canTakeScreenshotLocked(mServiceConnection)).thenReturn(true);
+        when(mMockSecurityPolicy.checkAccessibilityAccess(mServiceConnection)).thenReturn(false);
+        assertThat(mServiceConnection.takeScreenshot(Display.DEFAULT_DISPLAY), is(nullValue()));
+    }
+
     private void updateServiceInfo(AccessibilityServiceInfo serviceInfo, int eventType,
             int feedbackType, int flags, String[] packageNames, int notificationTimeout) {
         serviceInfo.eventTypes = eventType;
@@ -832,5 +845,8 @@
 
         @Override
         public void onFingerprintGesture(int gesture) {}
+
+        @Override
+        public void takeScreenshotWithCallback(int displayId, RemoteCallback callback) {}
     }
 }
diff --git a/services/tests/servicestests/src/com/android/server/accessibility/AccessibilitySecurityPolicyTest.java b/services/tests/servicestests/src/com/android/server/accessibility/AccessibilitySecurityPolicyTest.java
index 04ac7fe..1504097 100644
--- a/services/tests/servicestests/src/com/android/server/accessibility/AccessibilitySecurityPolicyTest.java
+++ b/services/tests/servicestests/src/com/android/server/accessibility/AccessibilitySecurityPolicyTest.java
@@ -346,6 +346,14 @@
     }
 
     @Test
+    public void canTakeScreenshot_hasCapability_returnTrue() {
+        when(mMockA11yServiceConnection.getCapabilities())
+                .thenReturn(AccessibilityServiceInfo.CAPABILITY_CAN_TAKE_SCREENSHOT);
+
+        assertTrue(mA11ySecurityPolicy.canTakeScreenshotLocked(mMockA11yServiceConnection));
+    }
+
+    @Test
     public void resolveProfileParent_userIdIsCurrentUser_returnCurrentUser() {
         final int currentUserId = 10;
         final int userId = currentUserId;