Added captureScreen function for displayId or layerStack

Previously, screencap shell command could accept a display id as an argument.
It would use the id to look up information about the display (displayToken,
orientation, colorMode) by making requests into SurfaceFlinger. Since those
requests only worked for physical display ids, screencap didn't work to capture
virtual displays.

Instead of adding new methods to look up display information for a virtual
display, this change just adds a new captureScreen function in SurfaceFlinger
that accepts a displayId or layerStack and handles getting the default info
for the requested display. This works for both physical and virtual displays.

Test: adb shell screencap -d <layerStack for virt display>
Fixes: 130974213
Change-Id: I24b7558c973a057414c6b4f81ab1d60152fff38d
diff --git a/libs/gui/ISurfaceComposer.cpp b/libs/gui/ISurfaceComposer.cpp
index a8b1a4c..4372e16 100644
--- a/libs/gui/ISurfaceComposer.cpp
+++ b/libs/gui/ISurfaceComposer.cpp
@@ -142,6 +142,28 @@
         return result;
     }
 
+    virtual status_t captureScreen(uint64_t displayOrLayerStack, ui::Dataspace* outDataspace,
+                                   sp<GraphicBuffer>* outBuffer) {
+        Parcel data, reply;
+        data.writeInterfaceToken(ISurfaceComposer::getInterfaceDescriptor());
+        data.writeUint64(displayOrLayerStack);
+        status_t result = remote()->transact(BnSurfaceComposer::CAPTURE_SCREEN_BY_ID, data, &reply);
+        if (result != NO_ERROR) {
+            ALOGE("captureScreen failed to transact: %d", result);
+            return result;
+        }
+        result = reply.readInt32();
+        if (result != NO_ERROR) {
+            ALOGE("captureScreen failed to readInt32: %d", result);
+            return result;
+        }
+
+        *outDataspace = static_cast<ui::Dataspace>(reply.readInt32());
+        *outBuffer = new GraphicBuffer();
+        reply.read(**outBuffer);
+        return result;
+    }
+
     virtual status_t captureLayers(
             const sp<IBinder>& layerHandleBinder, sp<GraphicBuffer>* outBuffer,
             const ui::Dataspace reqDataspace, const ui::PixelFormat reqPixelFormat,
@@ -1042,6 +1064,19 @@
             }
             return NO_ERROR;
         }
+        case CAPTURE_SCREEN_BY_ID: {
+            CHECK_INTERFACE(ISurfaceComposer, data, reply);
+            uint64_t displayOrLayerStack = data.readUint64();
+            ui::Dataspace outDataspace = ui::Dataspace::V0_SRGB;
+            sp<GraphicBuffer> outBuffer;
+            status_t res = captureScreen(displayOrLayerStack, &outDataspace, &outBuffer);
+            reply->writeInt32(res);
+            if (res == NO_ERROR) {
+                reply->writeInt32(static_cast<int32_t>(outDataspace));
+                reply->write(*outBuffer);
+            }
+            return NO_ERROR;
+        }
         case CAPTURE_LAYERS: {
             CHECK_INTERFACE(ISurfaceComposer, data, reply);
             sp<IBinder> layerHandleBinder = data.readStrongBinder();
diff --git a/libs/gui/SurfaceComposerClient.cpp b/libs/gui/SurfaceComposerClient.cpp
index 437cdd7..5d4367d 100644
--- a/libs/gui/SurfaceComposerClient.cpp
+++ b/libs/gui/SurfaceComposerClient.cpp
@@ -1572,6 +1572,13 @@
                    useIdentityTransform, rotation, false, outBuffer, ignored);
 }
 
+status_t ScreenshotClient::capture(uint64_t displayOrLayerStack, ui::Dataspace* outDataspace,
+                                   sp<GraphicBuffer>* outBuffer) {
+    sp<ISurfaceComposer> s(ComposerService::getComposerService());
+    if (s == nullptr) return NO_INIT;
+    return s->captureScreen(displayOrLayerStack, outDataspace, outBuffer);
+}
+
 status_t ScreenshotClient::captureLayers(const sp<IBinder>& layerHandle,
                                          const ui::Dataspace reqDataSpace,
                                          const ui::PixelFormat reqPixelFormat, Rect sourceCrop,
diff --git a/libs/gui/include/gui/ISurfaceComposer.h b/libs/gui/include/gui/ISurfaceComposer.h
index e8c7a39..fd67754 100644
--- a/libs/gui/include/gui/ISurfaceComposer.h
+++ b/libs/gui/include/gui/ISurfaceComposer.h
@@ -247,6 +247,9 @@
                              useIdentityTransform, rotation);
     }
 
+    virtual status_t captureScreen(uint64_t displayOrLayerStack, ui::Dataspace* outDataspace,
+                                   sp<GraphicBuffer>* outBuffer) = 0;
+
     template <class AA>
     struct SpHash {
         size_t operator()(const sp<AA>& k) const { return std::hash<AA*>()(k.get()); }
@@ -473,6 +476,7 @@
         GET_ALLOWED_DISPLAY_CONFIGS,
         GET_DISPLAY_BRIGHTNESS_SUPPORT,
         SET_DISPLAY_BRIGHTNESS,
+        CAPTURE_SCREEN_BY_ID,
         // Always append new enum to the end.
     };
 
diff --git a/libs/gui/include/gui/SurfaceComposerClient.h b/libs/gui/include/gui/SurfaceComposerClient.h
index 9d34468..a038838 100644
--- a/libs/gui/include/gui/SurfaceComposerClient.h
+++ b/libs/gui/include/gui/SurfaceComposerClient.h
@@ -515,6 +515,8 @@
                             const ui::PixelFormat reqPixelFormat, Rect sourceCrop,
                             uint32_t reqWidth, uint32_t reqHeight, bool useIdentityTransform,
                             uint32_t rotation, sp<GraphicBuffer>* outBuffer);
+    static status_t capture(uint64_t displayOrLayerStack, ui::Dataspace* outDataspace,
+                            sp<GraphicBuffer>* outBuffer);
     static status_t captureLayers(const sp<IBinder>& layerHandle, const ui::Dataspace reqDataSpace,
                                   const ui::PixelFormat reqPixelFormat, Rect sourceCrop,
                                   float frameScale, sp<GraphicBuffer>* outBuffer);
diff --git a/libs/gui/tests/Surface_test.cpp b/libs/gui/tests/Surface_test.cpp
index 014b1fa..db97fae 100644
--- a/libs/gui/tests/Surface_test.cpp
+++ b/libs/gui/tests/Surface_test.cpp
@@ -623,6 +623,10 @@
                            bool /*captureSecureLayers*/) override {
         return NO_ERROR;
     }
+    status_t captureScreen(uint64_t /*displayOrLayerStack*/, ui::Dataspace* /*outDataspace*/,
+                           sp<GraphicBuffer>* /*outBuffer*/) override {
+        return NO_ERROR;
+    }
     virtual status_t captureLayers(
             const sp<IBinder>& /*parentHandle*/, sp<GraphicBuffer>* /*outBuffer*/,
             const ui::Dataspace /*reqDataspace*/, const ui::PixelFormat /*reqPixelFormat*/,
diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index fc0d5fc..8be24b5 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -5086,6 +5086,14 @@
             ALOGE("Attempting to access SurfaceFlinger with unused code: %u", code);
             return PERMISSION_DENIED;
         }
+        case CAPTURE_SCREEN_BY_ID: {
+            IPCThreadState* ipc = IPCThreadState::self();
+            const int uid = ipc->getCallingUid();
+            if ((uid == AID_GRAPHICS) || (uid == AID_SYSTEM) || (uid == AID_SHELL)) {
+                return OK;
+            }
+            return PERMISSION_DENIED;
+        }
     }
 
     // These codes are used for the IBinder protocol to either interrogate the recipient
@@ -5500,6 +5508,66 @@
                                useIdentityTransform, outCapturedSecureLayers);
 }
 
+static Dataspace pickDataspaceFromColorMode(const ColorMode colorMode) {
+    switch (colorMode) {
+        case ColorMode::DISPLAY_P3:
+        case ColorMode::BT2100_PQ:
+        case ColorMode::BT2100_HLG:
+        case ColorMode::DISPLAY_BT2020:
+            return Dataspace::DISPLAY_P3;
+        default:
+            return Dataspace::V0_SRGB;
+    }
+}
+
+const sp<DisplayDevice> SurfaceFlinger::getDisplayByIdOrLayerStack(uint64_t displayOrLayerStack) {
+    const sp<IBinder> displayToken = getPhysicalDisplayTokenLocked(DisplayId{displayOrLayerStack});
+    if (displayToken) {
+        return getDisplayDeviceLocked(displayToken);
+    }
+    // Couldn't find display by displayId. Try to get display by layerStack since virtual displays
+    // may not have a displayId.
+    for (const auto& [token, display] : mDisplays) {
+        if (display->getLayerStack() == displayOrLayerStack) {
+            return display;
+        }
+    }
+    return nullptr;
+}
+
+status_t SurfaceFlinger::captureScreen(uint64_t displayOrLayerStack, Dataspace* outDataspace,
+                                       sp<GraphicBuffer>* outBuffer) {
+    sp<DisplayDevice> display;
+    uint32_t width;
+    uint32_t height;
+    ui::Transform::orientation_flags captureOrientation;
+    {
+        Mutex::Autolock _l(mStateLock);
+        display = getDisplayByIdOrLayerStack(displayOrLayerStack);
+        if (!display) {
+            return BAD_VALUE;
+        }
+
+        width = uint32_t(display->getViewport().width());
+        height = uint32_t(display->getViewport().height());
+
+        captureOrientation = fromSurfaceComposerRotation(
+                static_cast<ISurfaceComposer::Rotation>(display->getOrientation()));
+        *outDataspace =
+                pickDataspaceFromColorMode(display->getCompositionDisplay()->getState().colorMode);
+    }
+
+    DisplayRenderArea renderArea(display, Rect(), width, height, *outDataspace, captureOrientation,
+                                 false /* captureSecureLayers */);
+
+    auto traverseLayers = std::bind(&SurfaceFlinger::traverseLayersInDisplay, this, display,
+                                    std::placeholders::_1);
+    bool ignored = false;
+    return captureScreenCommon(renderArea, traverseLayers, outBuffer, ui::PixelFormat::RGBA_8888,
+                               false /* useIdentityTransform */,
+                               ignored /* outCapturedSecureLayers */);
+}
+
 status_t SurfaceFlinger::captureLayers(
         const sp<IBinder>& layerHandleBinder, sp<GraphicBuffer>* outBuffer,
         const Dataspace reqDataspace, const ui::PixelFormat reqPixelFormat, const Rect& sourceCrop,
diff --git a/services/surfaceflinger/SurfaceFlinger.h b/services/surfaceflinger/SurfaceFlinger.h
index 5871774..d9de07b 100644
--- a/services/surfaceflinger/SurfaceFlinger.h
+++ b/services/surfaceflinger/SurfaceFlinger.h
@@ -405,6 +405,8 @@
             const ui::PixelFormat reqPixelFormat, Rect sourceCrop,
             uint32_t reqWidth, uint32_t reqHeight,
             bool useIdentityTransform, ISurfaceComposer::Rotation rotation, bool captureSecureLayers) override;
+    status_t captureScreen(uint64_t displayOrLayerStack, ui::Dataspace* outDataspace,
+                           sp<GraphicBuffer>* outBuffer) override;
     status_t captureLayers(
             const sp<IBinder>& parentHandle, sp<GraphicBuffer>* outBuffer,
             const ui::Dataspace reqDataspace, const ui::PixelFormat reqPixelFormat,
@@ -636,6 +638,7 @@
     status_t captureScreenCommon(RenderArea& renderArea, TraverseLayersFunction traverseLayers,
                                  const sp<GraphicBuffer>& buffer, bool useIdentityTransform,
                                  bool& outCapturedSecureLayers);
+    const sp<DisplayDevice> getDisplayByIdOrLayerStack(uint64_t displayOrLayerStack);
     status_t captureScreenImplLocked(const RenderArea& renderArea,
                                      TraverseLayersFunction traverseLayers,
                                      ANativeWindowBuffer* buffer, bool useIdentityTransform,