[SurfaceFlinger] Accept data space and pixel format to take screenshot.

Previously we always take screenshot in SRGB data space. However, the current
color mode of the display can be in a different color mode, and so can HWUI
render in a different color space. Thus we want to remove the assumption that
screenshot is captured in SRGB color space, and give the ability to the upper
stack to specify the color space and pixel format it wants.

BUG: 111436479
BUG: 116112787
Test: Build, flash, boot and take screenshot, observe there's no color shift.
Test: Take screenshot, verified the color profile is correct of the png.
Change-Id: I83a0682839f39a99bff6b471a33060dd570be9b0
diff --git a/services/surfaceflinger/DisplayDevice.h b/services/surfaceflinger/DisplayDevice.h
index 38482c9..918f7de 100644
--- a/services/surfaceflinger/DisplayDevice.h
+++ b/services/surfaceflinger/DisplayDevice.h
@@ -349,10 +349,11 @@
     DisplayRenderArea(const sp<const DisplayDevice> device,
                       ui::Transform::orientation_flags rotation = ui::Transform::ROT_0)
           : DisplayRenderArea(device, device->getBounds(), device->getWidth(), device->getHeight(),
-                              rotation) {}
+                              device->getCompositionDataSpace(), rotation) {}
     DisplayRenderArea(const sp<const DisplayDevice> device, Rect sourceCrop, uint32_t reqWidth,
-                      uint32_t reqHeight, ui::Transform::orientation_flags rotation)
-          : RenderArea(reqWidth, reqHeight, CaptureFill::OPAQUE,
+                      uint32_t reqHeight, ui::Dataspace reqDataSpace,
+                      ui::Transform::orientation_flags rotation)
+          : RenderArea(reqWidth, reqHeight, CaptureFill::OPAQUE, reqDataSpace,
                        getDisplayRotation(rotation, device->getInstallOrientation())),
             mDevice(device),
             mSourceCrop(sourceCrop) {}
diff --git a/services/surfaceflinger/RenderArea.h b/services/surfaceflinger/RenderArea.h
index 3c11e73..9bad6de 100644
--- a/services/surfaceflinger/RenderArea.h
+++ b/services/surfaceflinger/RenderArea.h
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <ui/GraphicTypes.h>
 #include <ui/Transform.h>
 
 #include <functional>
@@ -19,9 +20,11 @@
     static float getCaptureFillValue(CaptureFill captureFill);
 
     RenderArea(uint32_t reqWidth, uint32_t reqHeight, CaptureFill captureFill,
+               ui::Dataspace reqDataSpace,
                ui::Transform::orientation_flags rotation = ui::Transform::ROT_0)
           : mReqWidth(reqWidth),
             mReqHeight(reqHeight),
+            mReqDataSpace(reqDataSpace),
             mCaptureFill(captureFill),
             mRotationFlags(rotation) {}
 
@@ -66,6 +69,9 @@
     int getReqWidth() const { return mReqWidth; };
     int getReqHeight() const { return mReqHeight; };
 
+    // Returns the composition data space of the render area.
+    ui::Dataspace getReqDataSpace() const { return mReqDataSpace; }
+
     // Returns the fill color of the physical render area.  Regions not
     // covered by any rendered layer should be filled with this color.
     CaptureFill getCaptureFill() const { return mCaptureFill; };
@@ -73,6 +79,7 @@
 private:
     const uint32_t mReqWidth;
     const uint32_t mReqHeight;
+    const ui::Dataspace mReqDataSpace;
     const CaptureFill mCaptureFill;
     const ui::Transform::orientation_flags mRotationFlags;
 };
diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index f0723e8..fe42480 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -5257,7 +5257,8 @@
 };
 
 status_t SurfaceFlinger::captureScreen(const sp<IBinder>& displayToken,
-                                       sp<GraphicBuffer>* outBuffer, Rect sourceCrop,
+                                       sp<GraphicBuffer>* outBuffer, const Dataspace reqDataspace,
+                                       const ui::PixelFormat reqPixelFormat, Rect sourceCrop,
                                        uint32_t reqWidth, uint32_t reqHeight,
                                        bool useIdentityTransform,
                                        ISurfaceComposer::Rotation rotation) {
@@ -5286,23 +5287,27 @@
         }
     }
 
-    DisplayRenderArea renderArea(display, sourceCrop, reqWidth, reqHeight, renderAreaRotation);
+    DisplayRenderArea renderArea(display, sourceCrop, reqWidth, reqHeight, reqDataspace,
+                                 renderAreaRotation);
 
     auto traverseLayers = std::bind(std::mem_fn(&SurfaceFlinger::traverseLayersInDisplay), this,
                                     display, std::placeholders::_1);
-    return captureScreenCommon(renderArea, traverseLayers, outBuffer, useIdentityTransform);
+    return captureScreenCommon(renderArea, traverseLayers, outBuffer, reqPixelFormat,
+                               useIdentityTransform);
 }
 
 status_t SurfaceFlinger::captureLayers(const sp<IBinder>& layerHandleBinder,
-                                       sp<GraphicBuffer>* outBuffer, const Rect& sourceCrop,
+                                       sp<GraphicBuffer>* outBuffer, const Dataspace reqDataspace,
+                                       const ui::PixelFormat reqPixelFormat, const Rect& sourceCrop,
                                        float frameScale, bool childrenOnly) {
     ATRACE_CALL();
 
     class LayerRenderArea : public RenderArea {
     public:
         LayerRenderArea(SurfaceFlinger* flinger, const sp<Layer>& layer, const Rect crop,
-                        int32_t reqWidth, int32_t reqHeight, bool childrenOnly)
-              : RenderArea(reqWidth, reqHeight, CaptureFill::CLEAR),
+                        int32_t reqWidth, int32_t reqHeight, Dataspace reqDataSpace,
+                        bool childrenOnly)
+              : RenderArea(reqWidth, reqHeight, CaptureFill::CLEAR, reqDataSpace),
                 mLayer(layer),
                 mCrop(crop),
                 mNeedsFiltering(false),
@@ -5409,7 +5414,7 @@
         reqHeight = 1;
     }
 
-    LayerRenderArea renderArea(this, parent, crop, reqWidth, reqHeight, childrenOnly);
+    LayerRenderArea renderArea(this, parent, crop, reqWidth, reqHeight, reqDataspace, childrenOnly);
 
     auto traverseLayers = [parent, childrenOnly](const LayerVector::Visitor& visitor) {
         parent->traverseChildrenInZOrder(LayerVector::StateSet::Drawing, [&](Layer* layer) {
@@ -5421,19 +5426,22 @@
             visitor(layer);
         });
     };
-    return captureScreenCommon(renderArea, traverseLayers, outBuffer, false);
+    return captureScreenCommon(renderArea, traverseLayers, outBuffer, reqPixelFormat, false);
 }
 
 status_t SurfaceFlinger::captureScreenCommon(RenderArea& renderArea,
                                              TraverseLayersFunction traverseLayers,
                                              sp<GraphicBuffer>* outBuffer,
+                                             const ui::PixelFormat reqPixelFormat,
                                              bool useIdentityTransform) {
     ATRACE_CALL();
 
+    // TODO(b/116112787) Make buffer usage a parameter.
     const uint32_t usage = GRALLOC_USAGE_SW_READ_OFTEN | GRALLOC_USAGE_SW_WRITE_OFTEN |
             GRALLOC_USAGE_HW_RENDER | GRALLOC_USAGE_HW_TEXTURE;
     *outBuffer = new GraphicBuffer(renderArea.getReqWidth(), renderArea.getReqHeight(),
-                                   HAL_PIXEL_FORMAT_RGBA_8888, 1, usage, "screenshot");
+                                   static_cast<android_pixel_format>(reqPixelFormat), 1, usage,
+                                   "screenshot");
 
     // This mutex protects syncFd and captureResult for communication of the return values from the
     // main thread back to this Binder thread
@@ -5509,8 +5517,7 @@
     const auto sourceCrop = renderArea.getSourceCrop();
     const auto rotation = renderArea.getRotationFlags();
 
-    // assume ColorMode::SRGB / RenderIntent::COLORIMETRIC
-    engine.setOutputDataSpace(Dataspace::SRGB);
+    engine.setOutputDataSpace(renderArea.getReqDataSpace());
     engine.setDisplayMaxLuminance(DisplayDevice::sDefaultMaxLumiance);
 
     // make sure to clear all GL error flags
diff --git a/services/surfaceflinger/SurfaceFlinger.h b/services/surfaceflinger/SurfaceFlinger.h
index f535c4e..1425818 100644
--- a/services/surfaceflinger/SurfaceFlinger.h
+++ b/services/surfaceflinger/SurfaceFlinger.h
@@ -433,10 +433,14 @@
     virtual sp<IDisplayEventConnection> createDisplayEventConnection(
             ISurfaceComposer::VsyncSource vsyncSource = eVsyncSourceApp);
     virtual status_t captureScreen(const sp<IBinder>& displayToken, sp<GraphicBuffer>* outBuffer,
-                                   Rect sourceCrop, uint32_t reqWidth, uint32_t reqHeight,
-                                   bool useIdentityTransform, ISurfaceComposer::Rotation rotation);
+                                   const ui::Dataspace reqDataspace,
+                                   const ui::PixelFormat reqPixelFormat, Rect sourceCrop,
+                                   uint32_t reqWidth, uint32_t reqHeight, bool useIdentityTransform,
+                                   ISurfaceComposer::Rotation rotation);
     virtual status_t captureLayers(const sp<IBinder>& parentHandle, sp<GraphicBuffer>* outBuffer,
-                                   const Rect& sourceCrop, float frameScale, bool childrenOnly);
+                                   const ui::Dataspace reqDataspace,
+                                   const ui::PixelFormat reqPixelFormat, const Rect& sourceCrop,
+                                   float frameScale, bool childrenOnly);
     virtual status_t getDisplayStats(const sp<IBinder>& displayToken, DisplayStatInfo* stats);
     virtual status_t getDisplayConfigs(const sp<IBinder>& displayToken,
                                        Vector<DisplayInfo>* configs);
@@ -586,7 +590,7 @@
     void renderScreenImplLocked(const RenderArea& renderArea, TraverseLayersFunction traverseLayers,
                                 bool useIdentityTransform);
     status_t captureScreenCommon(RenderArea& renderArea, TraverseLayersFunction traverseLayers,
-                                 sp<GraphicBuffer>* outBuffer,
+                                 sp<GraphicBuffer>* outBuffer, const ui::PixelFormat reqPixelFormat,
                                  bool useIdentityTransform);
     status_t captureScreenImplLocked(const RenderArea& renderArea,
                                      TraverseLayersFunction traverseLayers,
diff --git a/services/surfaceflinger/tests/Credentials_test.cpp b/services/surfaceflinger/tests/Credentials_test.cpp
index 9e762f0..a73ec6c 100644
--- a/services/surfaceflinger/tests/Credentials_test.cpp
+++ b/services/surfaceflinger/tests/Credentials_test.cpp
@@ -267,8 +267,9 @@
     sp<IBinder> display(SurfaceComposerClient::getBuiltInDisplay(ISurfaceComposer::eDisplayIdMain));
     std::function<status_t()> condition = [=]() {
         sp<GraphicBuffer> outBuffer;
-        return ScreenshotClient::capture(display, Rect(), 0 /*reqWidth*/, 0 /*reqHeight*/, false,
-                                         ROTATION, &outBuffer);
+        return ScreenshotClient::capture(display, ui::Dataspace::V0_SRGB,
+                                         ui::PixelFormat::RGBA_8888, Rect(), 0 /*reqWidth*/,
+                                         0 /*reqHeight*/, false, ROTATION, &outBuffer);
     };
     ASSERT_NO_FATAL_FAILURE(checkWithPrivileges<status_t>(condition, NO_ERROR, PERMISSION_DENIED));
 }
@@ -278,8 +279,9 @@
     sp<GraphicBuffer> outBuffer;
     std::function<status_t()> condition = [=]() {
         sp<GraphicBuffer> outBuffer;
-        return ScreenshotClient::captureLayers(mBGSurfaceControl->getHandle(), Rect(), FRAME_SCALE,
-                                               &outBuffer);
+        return ScreenshotClient::captureLayers(mBGSurfaceControl->getHandle(),
+                                               ui::Dataspace::V0_SRGB, ui::PixelFormat::RGBA_8888,
+                                               Rect(), FRAME_SCALE, &outBuffer);
     };
     ASSERT_NO_FATAL_FAILURE(checkWithPrivileges<status_t>(condition, NO_ERROR, PERMISSION_DENIED));
 }
diff --git a/services/surfaceflinger/tests/unittests/CompositionTest.cpp b/services/surfaceflinger/tests/unittests/CompositionTest.cpp
index 4e403b7..f430679 100644
--- a/services/surfaceflinger/tests/unittests/CompositionTest.cpp
+++ b/services/surfaceflinger/tests/unittests/CompositionTest.cpp
@@ -189,7 +189,8 @@
     constexpr bool forSystem = true;
 
     DisplayRenderArea renderArea(mDisplay, sourceCrop, DEFAULT_DISPLAY_WIDTH,
-                                 DEFAULT_DISPLAY_HEIGHT, ui::Transform::ROT_0);
+                                 DEFAULT_DISPLAY_HEIGHT, ui::Dataspace::V0_SRGB,
+                                 ui::Transform::ROT_0);
 
     auto traverseLayers = [this](const LayerVector::Visitor& visitor) {
         return mFlinger.traverseLayersInDisplay(mDisplay, visitor);
@@ -287,7 +288,7 @@
         EXPECT_CALL(*test->mRenderEngine, flush()).WillOnce(Return(ByMove(base::unique_fd())));
         EXPECT_CALL(*test->mRenderEngine, finish()).WillOnce(Return(true));
 
-        EXPECT_CALL(*test->mRenderEngine, setOutputDataSpace(ui::Dataspace::SRGB)).Times(1);
+        EXPECT_CALL(*test->mRenderEngine, setOutputDataSpace(_)).Times(1);
         EXPECT_CALL(*test->mRenderEngine, setDisplayMaxLuminance(DEFAULT_DISPLAY_MAX_LUMINANCE))
                 .Times(1);
         // This expectation retires on saturation as setViewportAndProjection is