SurfaceFlinger: Add exclusion list for captureLayers.

Among other use cases, WM needs to be able to omit the IME from Task Snapshots,
even while it is on-screen.

Bug: 126614127
Test: Transaction_test.cpp#CaptureLayerExclude,CaptureLayerExcludeTree
Change-Id: I055d99106c9ce2ed90d64eca06961d88cbd5e2d4
diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index 5676c59..fccd910 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -5405,10 +5405,11 @@
                                useIdentityTransform);
 }
 
-status_t SurfaceFlinger::captureLayers(const sp<IBinder>& layerHandleBinder,
-                                       sp<GraphicBuffer>* outBuffer, const Dataspace reqDataspace,
-                                       const ui::PixelFormat reqPixelFormat, const Rect& sourceCrop,
-                                       float frameScale, bool childrenOnly) {
+status_t SurfaceFlinger::captureLayers(
+        const sp<IBinder>& layerHandleBinder, sp<GraphicBuffer>* outBuffer,
+        const Dataspace reqDataspace, const ui::PixelFormat reqPixelFormat, const Rect& sourceCrop,
+        const std::unordered_set<sp<IBinder>, ISurfaceComposer::SpHash<IBinder>>& excludeHandles,
+        float frameScale, bool childrenOnly) {
     ATRACE_CALL();
 
     class LayerRenderArea : public RenderArea {
@@ -5530,15 +5531,36 @@
         reqHeight = 1;
     }
 
-    LayerRenderArea renderArea(this, parent, crop, reqWidth, reqHeight, reqDataspace, childrenOnly);
+    std::unordered_set<sp<Layer>, ISurfaceComposer::SpHash<Layer>> excludeLayers;
+    for (const auto& handle : excludeHandles) {
+        BBinder* local = handle->localBinder();
+        if (local != nullptr) {
+            auto layerHandle = reinterpret_cast<Layer::Handle*>(local);
+            excludeLayers.emplace(layerHandle->owner.promote());
+        } else {
+            ALOGW("Invalid layer handle passed as excludeLayer to captureLayers");
+            return NAME_NOT_FOUND;
+        }
+    }
 
-    auto traverseLayers = [parent, childrenOnly](const LayerVector::Visitor& visitor) {
+    LayerRenderArea renderArea(this, parent, crop, reqWidth, reqHeight, reqDataspace, childrenOnly);
+    auto traverseLayers = [parent, childrenOnly,
+                           &excludeLayers](const LayerVector::Visitor& visitor) {
         parent->traverseChildrenInZOrder(LayerVector::StateSet::Drawing, [&](Layer* layer) {
             if (!layer->isVisible()) {
                 return;
             } else if (childrenOnly && layer == parent.get()) {
                 return;
             }
+
+            sp<Layer> p = layer;
+            while (p != nullptr) {
+                if (excludeLayers.count(p) != 0) {
+                    return;
+                }
+                p = p->getParent();
+            }
+
             visitor(layer);
         });
     };
diff --git a/services/surfaceflinger/SurfaceFlinger.h b/services/surfaceflinger/SurfaceFlinger.h
index 5cd0f21..7e8e836 100644
--- a/services/surfaceflinger/SurfaceFlinger.h
+++ b/services/surfaceflinger/SurfaceFlinger.h
@@ -399,9 +399,12 @@
                            Rect sourceCrop, uint32_t reqWidth, uint32_t reqHeight,
                            bool useIdentityTransform, ISurfaceComposer::Rotation rotation,
                            bool captureSecureLayers) override;
-    status_t captureLayers(const sp<IBinder>& parentHandle, sp<GraphicBuffer>* outBuffer,
-                           const ui::Dataspace reqDataspace, const ui::PixelFormat reqPixelFormat,
-                           const Rect& sourceCrop, float frameScale, bool childrenOnly) override;
+    status_t captureLayers(
+            const sp<IBinder>& parentHandle, sp<GraphicBuffer>* outBuffer,
+            const ui::Dataspace reqDataspace, const ui::PixelFormat reqPixelFormat,
+            const Rect& sourceCrop,
+            const std::unordered_set<sp<IBinder>, ISurfaceComposer::SpHash<IBinder>>& exclude,
+            float frameScale, bool childrenOnly) override;
     status_t getDisplayStats(const sp<IBinder>& displayToken, DisplayStatInfo* stats) override;
     status_t getDisplayConfigs(const sp<IBinder>& displayToken,
                                Vector<DisplayInfo>* configs) override {
diff --git a/services/surfaceflinger/tests/Transaction_test.cpp b/services/surfaceflinger/tests/Transaction_test.cpp
index d62afa5..ba854e3 100644
--- a/services/surfaceflinger/tests/Transaction_test.cpp
+++ b/services/surfaceflinger/tests/Transaction_test.cpp
@@ -231,6 +231,20 @@
         *sc = std::make_unique<ScreenCapture>(outBuffer);
     }
 
+    static void captureChildLayersExcluding(
+            std::unique_ptr<ScreenCapture>* sc, sp<IBinder>& parentHandle,
+            std::unordered_set<sp<IBinder>, ISurfaceComposer::SpHash<IBinder>> excludeLayers) {
+        sp<ISurfaceComposer> sf(ComposerService::getComposerService());
+        SurfaceComposerClient::Transaction().apply(true);
+
+        sp<GraphicBuffer> outBuffer;
+        ASSERT_EQ(NO_ERROR,
+                  sf->captureLayers(parentHandle, &outBuffer, ui::Dataspace::V0_SRGB,
+                                    ui::PixelFormat::RGBA_8888, Rect::EMPTY_RECT, excludeLayers,
+                                    1.0f, true));
+        *sc = std::make_unique<ScreenCapture>(outBuffer);
+    }
+
     void expectColor(const Rect& rect, const Color& color, uint8_t tolerance = 0) {
         ASSERT_EQ(HAL_PIXEL_FORMAT_RGBA_8888, mOutBuffer->getPixelFormat());
         expectBufferColor(mOutBuffer, mPixels, rect, color, tolerance);
@@ -5238,6 +5252,57 @@
     mCapture->expectChildColor(0, 0);
 }
 
+TEST_F(ScreenCaptureTest, CaptureLayerExclude) {
+    auto fgHandle = mFGSurfaceControl->getHandle();
+
+    sp<SurfaceControl> child = createSurface(mClient, "Child surface", 10, 10,
+                                             PIXEL_FORMAT_RGBA_8888, 0, mFGSurfaceControl.get());
+    fillSurfaceRGBA8(child, 200, 200, 200);
+    sp<SurfaceControl> child2 = createSurface(mClient, "Child surface", 10, 10,
+                                              PIXEL_FORMAT_RGBA_8888, 0, mFGSurfaceControl.get());
+    fillSurfaceRGBA8(child2, 200, 0, 200);
+
+    SurfaceComposerClient::Transaction()
+            .show(child)
+            .show(child2)
+            .setLayer(child, 1)
+            .setLayer(child2, 2)
+            .apply(true);
+
+    // Child2 would be visible but its excluded, so we should see child1 color instead.
+    ScreenCapture::captureChildLayersExcluding(&mCapture, fgHandle, {child2->getHandle()});
+    mCapture->checkPixel(10, 10, 0, 0, 0);
+    mCapture->checkPixel(0, 0, 200, 200, 200);
+}
+
+// Like the last test but verifies that children are also exclude.
+TEST_F(ScreenCaptureTest, CaptureLayerExcludeTree) {
+    auto fgHandle = mFGSurfaceControl->getHandle();
+
+    sp<SurfaceControl> child = createSurface(mClient, "Child surface", 10, 10,
+                                             PIXEL_FORMAT_RGBA_8888, 0, mFGSurfaceControl.get());
+    fillSurfaceRGBA8(child, 200, 200, 200);
+    sp<SurfaceControl> child2 = createSurface(mClient, "Child surface", 10, 10,
+                                              PIXEL_FORMAT_RGBA_8888, 0, mFGSurfaceControl.get());
+    fillSurfaceRGBA8(child2, 200, 0, 200);
+    sp<SurfaceControl> child3 = createSurface(mClient, "Child surface", 10, 10,
+                                              PIXEL_FORMAT_RGBA_8888, 0, child2.get());
+    fillSurfaceRGBA8(child2, 200, 0, 200);
+
+    SurfaceComposerClient::Transaction()
+            .show(child)
+            .show(child2)
+            .show(child3)
+            .setLayer(child, 1)
+            .setLayer(child2, 2)
+            .apply(true);
+
+    // Child2 would be visible but its excluded, so we should see child1 color instead.
+    ScreenCapture::captureChildLayersExcluding(&mCapture, fgHandle, {child2->getHandle()});
+    mCapture->checkPixel(10, 10, 0, 0, 0);
+    mCapture->checkPixel(0, 0, 200, 200, 200);
+}
+
 TEST_F(ScreenCaptureTest, CaptureTransparent) {
     sp<SurfaceControl> child = createSurface(mClient, "Child surface", 10, 10,
                                              PIXEL_FORMAT_RGBA_8888, 0, mFGSurfaceControl.get());