Allow screen capture for a specified uid.

Currently, screenshots can either be done on the entire display or for a
layer hierarchy. Screenshots can only be taken by a system process or a
process with the correct permission. This change introduces a way for an
app to screenshot itself without introducing security issues

The following checks are done:

1. If the request is from graphics or has READ_FRAME_BUFFER permission
flag, the request is allowed. If a uid was set to non -1, only layers
from the specified uid will be captured. If uid in args is -1,
everything will be in the screenshot, so no layers will be skipped

2. If the request doesn't have the permissions, the calling uid is
compared to the uid in the request. If they match, we allow the
screenshot code to proceed and only screenshot the layers with the
matching uid.

3. Otherwise, we return with PERMISSION_DENIED

Test: ScreenCaptureTest
Bug: 155825630

Change-Id: I129b5a7f489383bf6d2f0ee333a416babc1444a2
diff --git a/libs/gui/LayerState.cpp b/libs/gui/LayerState.cpp
index b215756..d07d435 100644
--- a/libs/gui/LayerState.cpp
+++ b/libs/gui/LayerState.cpp
@@ -509,7 +509,8 @@
     status_t status = output.writeInt32(static_cast<int32_t>(pixelFormat)) ?:
         output.write(sourceCrop) ?:
         output.writeFloat(frameScale) ?:
-        output.writeBool(captureSecureLayers);
+        output.writeBool(captureSecureLayers) ?:
+        output.writeInt32(uid);
     return status;
 }
 
@@ -518,7 +519,8 @@
     status_t status = input.readInt32(&format) ?:
         input.read(sourceCrop) ?:
         input.readFloat(&frameScale) ?:
-        input.readBool(&captureSecureLayers);
+        input.readBool(&captureSecureLayers) ?:
+        input.readInt32(&uid);
 
     pixelFormat = static_cast<ui::PixelFormat>(format);
     return status;
diff --git a/libs/gui/include/gui/LayerState.h b/libs/gui/include/gui/LayerState.h
index 5e1066b..387ddea 100644
--- a/libs/gui/include/gui/LayerState.h
+++ b/libs/gui/include/gui/LayerState.h
@@ -314,12 +314,14 @@
 bool ValidateFrameRate(float frameRate, int8_t compatibility, const char* functionName);
 
 struct CaptureArgs {
+    const static int32_t UNSET_UID = -1;
     virtual ~CaptureArgs() = default;
 
     ui::PixelFormat pixelFormat{ui::PixelFormat::RGBA_8888};
     Rect sourceCrop;
     float frameScale{1};
     bool captureSecureLayers{false};
+    int32_t uid{UNSET_UID};
 
     virtual status_t write(Parcel& output) const;
     virtual status_t read(const Parcel& input);
diff --git a/services/surfaceflinger/Layer.h b/services/surfaceflinger/Layer.h
index 5fbdc34..8d8ab6d 100644
--- a/services/surfaceflinger/Layer.h
+++ b/services/surfaceflinger/Layer.h
@@ -436,9 +436,7 @@
     // Deprecated, please use compositionengine::Output::belongsInOutput()
     // instead.
     // TODO(lpique): Move the remaining callers (screencap) to the new function.
-    bool belongsToDisplay(uint32_t layerStack, bool isPrimaryDisplay) const {
-        return getLayerStack() == layerStack && (!mPrimaryDisplayOnly || isPrimaryDisplay);
-    }
+    bool belongsToDisplay(uint32_t layerStack) const { return getLayerStack() == layerStack; }
 
     FloatRect getBounds(const Region& activeTransparentRegion) const;
     FloatRect getBounds() const;
diff --git a/services/surfaceflinger/RegionSamplingThread.cpp b/services/surfaceflinger/RegionSamplingThread.cpp
index 17daadb..4e841e9 100644
--- a/services/surfaceflinger/RegionSamplingThread.cpp
+++ b/services/surfaceflinger/RegionSamplingThread.cpp
@@ -429,7 +429,7 @@
                   bounds.top, bounds.right, bounds.bottom);
             visitor(layer);
         };
-        mFlinger.traverseLayersInLayerStack(layerStack, filterVisitor);
+        mFlinger.traverseLayersInLayerStack(layerStack, CaptureArgs::UNSET_UID, filterVisitor);
     };
 
     sp<GraphicBuffer> buffer = nullptr;
diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index 9d9a4cf..6fae56a 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -4958,11 +4958,13 @@
         // special permissions.
         case SET_FRAME_RATE:
         case GET_DISPLAY_BRIGHTNESS_SUPPORT:
+        // captureLayers and captureDisplay will handle the permission check in the function
+        case CAPTURE_LAYERS:
+        case CAPTURE_DISPLAY:
         case SET_DISPLAY_BRIGHTNESS: {
             return OK;
         }
-        case CAPTURE_LAYERS:
-        case CAPTURE_DISPLAY:
+
         case ADD_REGION_SAMPLING_LISTENER:
         case REMOVE_REGION_SAMPLING_LISTENER: {
             // codes that require permission check
@@ -5434,10 +5436,33 @@
     }
 }
 
+static status_t validateScreenshotPermissions(const CaptureArgs& captureArgs) {
+    IPCThreadState* ipc = IPCThreadState::self();
+    const int pid = ipc->getCallingPid();
+    const int uid = ipc->getCallingUid();
+    if (uid == AID_GRAPHICS || PermissionCache::checkPermission(sReadFramebuffer, pid, uid)) {
+        return OK;
+    }
+
+    // If the caller doesn't have the correct permissions but is only attempting to screenshot
+    // itself, we allow it to continue.
+    if (captureArgs.uid == uid) {
+        return OK;
+    }
+
+    ALOGE("Permission Denial: can't take screenshot pid=%d, uid=%d", pid, uid);
+    return PERMISSION_DENIED;
+}
+
 status_t SurfaceFlinger::captureDisplay(const DisplayCaptureArgs& args,
                                         ScreenCaptureResults& captureResults) {
     ATRACE_CALL();
 
+    status_t validate = validateScreenshotPermissions(args);
+    if (validate != OK) {
+        return validate;
+    }
+
     if (!args.displayToken) return BAD_VALUE;
 
     auto renderAreaRotation = ui::Transform::toRotationFlags(args.rotation);
@@ -5472,8 +5497,8 @@
                                          renderAreaRotation, args.captureSecureLayers);
     });
 
-    auto traverseLayers = [this, layerStack](const LayerVector::Visitor& visitor) {
-        traverseLayersInLayerStack(layerStack, visitor);
+    auto traverseLayers = [this, args, layerStack](const LayerVector::Visitor& visitor) {
+        traverseLayersInLayerStack(layerStack, args.uid, visitor);
     };
     return captureScreenCommon(std::move(renderAreaFuture), traverseLayers, reqSize,
                                args.pixelFormat, args.useIdentityTransform, captureResults);
@@ -5567,7 +5592,7 @@
     });
 
     auto traverseLayers = [this, layerStack](const LayerVector::Visitor& visitor) {
-        traverseLayersInLayerStack(layerStack, visitor);
+        traverseLayersInLayerStack(layerStack, CaptureArgs::UNSET_UID, visitor);
     };
 
     return captureScreenCommon(std::move(renderAreaFuture), traverseLayers, size,
@@ -5579,6 +5604,11 @@
                                        ScreenCaptureResults& captureResults) {
     ATRACE_CALL();
 
+    status_t validate = validateScreenshotPermissions(args);
+    if (validate != OK) {
+        return validate;
+    }
+
     ui::Size reqSize;
     sp<Layer> parent;
     Rect crop(args.sourceCrop);
@@ -5652,19 +5682,19 @@
     }
 
     bool childrenOnly = args.childrenOnly;
-
     RenderAreaFuture renderAreaFuture = promise::defer([=]() -> std::unique_ptr<RenderArea> {
         return std::make_unique<LayerRenderArea>(*this, parent, crop, reqSize, dataspace,
                                                  childrenOnly, displayViewport,
                                                  captureSecureLayers);
     });
 
-    auto traverseLayers = [parent, childrenOnly,
-                           &excludeLayers](const LayerVector::Visitor& visitor) {
+    auto traverseLayers = [parent, args, &excludeLayers](const LayerVector::Visitor& visitor) {
         parent->traverseChildrenInZOrder(LayerVector::StateSet::Drawing, [&](Layer* layer) {
             if (!layer->isVisible()) {
                 return;
-            } else if (childrenOnly && layer == parent.get()) {
+            } else if (args.childrenOnly && layer == parent.get()) {
+                return;
+            } else if (args.uid != CaptureArgs::UNSET_UID && args.uid != layer->getOwnerUid()) {
                 return;
             }
 
@@ -5890,22 +5920,25 @@
     layersSortedByZ.traverseInReverseZOrder(stateSet, visitor);
 }
 
-void SurfaceFlinger::traverseLayersInLayerStack(ui::LayerStack layerStack,
+void SurfaceFlinger::traverseLayersInLayerStack(ui::LayerStack layerStack, const int32_t uid,
                                                 const LayerVector::Visitor& visitor) {
     // We loop through the first level of layers without traversing,
     // as we need to determine which layers belong to the requested display.
     for (const auto& layer : mDrawingState.layersSortedByZ) {
-        if (!layer->belongsToDisplay(layerStack, false)) {
+        if (!layer->belongsToDisplay(layerStack)) {
             continue;
         }
         // relative layers are traversed in Layer::traverseInZOrder
         layer->traverseInZOrder(LayerVector::StateSet::Drawing, [&](Layer* layer) {
-            if (!layer->belongsToDisplay(layerStack, false)) {
+            if (layer->getPrimaryDisplayOnly()) {
                 return;
             }
             if (!layer->isVisible()) {
                 return;
             }
+            if (uid != CaptureArgs::UNSET_UID && layer->getOwnerUid() != uid) {
+                return;
+            }
             visitor(layer);
         });
     }
diff --git a/services/surfaceflinger/SurfaceFlinger.h b/services/surfaceflinger/SurfaceFlinger.h
index 6c00931..a1b9b14 100644
--- a/services/surfaceflinger/SurfaceFlinger.h
+++ b/services/surfaceflinger/SurfaceFlinger.h
@@ -731,7 +731,9 @@
     sp<DisplayDevice> getDisplayByIdOrLayerStack(uint64_t displayOrLayerStack) REQUIRES(mStateLock);
     sp<DisplayDevice> getDisplayByLayerStack(uint64_t layerStack) REQUIRES(mStateLock);
 
-    void traverseLayersInLayerStack(ui::LayerStack, const LayerVector::Visitor&);
+    // If the uid provided is not UNSET_UID, the traverse will skip any layers that don't have a
+    // matching ownerUid
+    void traverseLayersInLayerStack(ui::LayerStack, const int32_t uid, const LayerVector::Visitor&);
 
     sp<StartPropertySetThread> mStartPropertySetThread;
 
diff --git a/services/surfaceflinger/tests/LayerTransaction_test.cpp b/services/surfaceflinger/tests/LayerTransaction_test.cpp
index 0ef4150..8d715e1 100644
--- a/services/surfaceflinger/tests/LayerTransaction_test.cpp
+++ b/services/surfaceflinger/tests/LayerTransaction_test.cpp
@@ -18,7 +18,6 @@
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wconversion"
 
-#include <private/android_filesystem_config.h>
 #include <thread>
 #include "LayerTransactionTest.h"
 
@@ -26,40 +25,6 @@
 
 using android::hardware::graphics::common::V1_1::BufferUsage;
 
-TEST_F(LayerTransactionTest, SetFlagsSecureEUidSystem) {
-    sp<SurfaceControl> layer;
-    ASSERT_NO_FATAL_FAILURE(layer = createLayer("test", 32, 32));
-    ASSERT_NO_FATAL_FAILURE(fillBufferQueueLayerColor(layer, Color::RED, 32, 32));
-
-    sp<ISurfaceComposer> composer = ComposerService::getComposerService();
-    Transaction()
-            .setFlags(layer, layer_state_t::eLayerSecure, layer_state_t::eLayerSecure)
-            .apply(true);
-    ASSERT_EQ(PERMISSION_DENIED, composer->captureDisplay(mCaptureArgs, mCaptureResults));
-
-    UIDFaker f(AID_SYSTEM);
-
-    // By default the system can capture screenshots with secure layers but they
-    // will be blacked out
-    ASSERT_EQ(NO_ERROR, composer->captureDisplay(mCaptureArgs, mCaptureResults));
-
-    {
-        SCOPED_TRACE("as system");
-        auto shot = screenshot();
-        shot->expectColor(Rect(0, 0, 32, 32), Color::BLACK);
-    }
-
-    // Here we pass captureSecureLayers = true and since we are AID_SYSTEM we should be able
-    // to receive them...we are expected to take care with the results.
-    DisplayCaptureArgs args;
-    args.displayToken = mDisplay;
-    args.captureSecureLayers = true;
-    ASSERT_EQ(NO_ERROR, composer->captureDisplay(args, mCaptureResults));
-    ASSERT_EQ(true, mCaptureResults.capturedSecureLayers);
-    ScreenCapture sc(mCaptureResults.buffer);
-    sc.expectColor(Rect(0, 0, 32, 32), Color::RED);
-}
-
 TEST_F(LayerTransactionTest, SetTransformToDisplayInverse_BufferState) {
     sp<SurfaceControl> layer;
     ASSERT_NO_FATAL_FAILURE(
diff --git a/services/surfaceflinger/tests/ScreenCapture_test.cpp b/services/surfaceflinger/tests/ScreenCapture_test.cpp
index 149e4d7..690f758 100644
--- a/services/surfaceflinger/tests/ScreenCapture_test.cpp
+++ b/services/surfaceflinger/tests/ScreenCapture_test.cpp
@@ -74,6 +74,42 @@
     std::unique_ptr<ScreenCapture> mCapture;
 };
 
+TEST_F(ScreenCaptureTest, SetFlagsSecureEUidSystem) {
+    sp<SurfaceControl> layer;
+    ASSERT_NO_FATAL_FAILURE(
+            layer = createLayer("test", 32, 32,
+                                ISurfaceComposerClient::eSecure |
+                                        ISurfaceComposerClient::eFXSurfaceBufferQueue));
+    ASSERT_NO_FATAL_FAILURE(fillBufferQueueLayerColor(layer, Color::RED, 32, 32));
+
+    Transaction().show(layer).setLayer(layer, INT32_MAX).apply(true);
+
+    sp<ISurfaceComposer> composer = ComposerService::getComposerService();
+    ASSERT_EQ(PERMISSION_DENIED, composer->captureDisplay(mCaptureArgs, mCaptureResults));
+
+    UIDFaker f(AID_SYSTEM);
+
+    // By default the system can capture screenshots with secure layers but they
+    // will be blacked out
+    ASSERT_EQ(NO_ERROR, composer->captureDisplay(mCaptureArgs, mCaptureResults));
+
+    {
+        SCOPED_TRACE("as system");
+        auto shot = screenshot();
+        shot->expectColor(Rect(0, 0, 32, 32), Color::BLACK);
+    }
+
+    // Here we pass captureSecureLayers = true and since we are AID_SYSTEM we should be able
+    // to receive them...we are expected to take care with the results.
+    DisplayCaptureArgs args;
+    args.displayToken = mDisplay;
+    args.captureSecureLayers = true;
+    ASSERT_EQ(NO_ERROR, composer->captureDisplay(args, mCaptureResults));
+    ASSERT_TRUE(mCaptureResults.capturedSecureLayers);
+    ScreenCapture sc(mCaptureResults.buffer);
+    sc.expectColor(Rect(0, 0, 32, 32), Color::RED);
+}
+
 TEST_F(ScreenCaptureTest, CaptureSingleLayer) {
     LayerCaptureArgs captureArgs;
     captureArgs.layerHandle = mBGSurfaceControl->getHandle();
@@ -494,6 +530,128 @@
     mCapture->expectColor(Rect(30, 30, 60, 60), Color::RED);
 }
 
+TEST_F(ScreenCaptureTest, CaptureDisplayWithUid) {
+    uid_t fakeUid = 12345;
+
+    DisplayCaptureArgs captureArgs;
+    captureArgs.displayToken = mDisplay;
+
+    sp<SurfaceControl> layer;
+    ASSERT_NO_FATAL_FAILURE(layer = createLayer("test layer", 32, 32,
+                                                ISurfaceComposerClient::eFXSurfaceBufferQueue,
+                                                mBGSurfaceControl.get()));
+    ASSERT_NO_FATAL_FAILURE(fillBufferQueueLayerColor(layer, Color::RED, 32, 32));
+
+    Transaction().show(layer).setLayer(layer, INT32_MAX).apply();
+
+    // Make sure red layer with the background layer is screenshot.
+    ScreenCapture::captureDisplay(&mCapture, captureArgs);
+    mCapture->expectColor(Rect(0, 0, 32, 32), Color::RED);
+    mCapture->expectBorder(Rect(0, 0, 32, 32), {63, 63, 195, 255});
+
+    // From non system uid, can't request screenshot without a specified uid.
+    UIDFaker f(fakeUid);
+    sp<ISurfaceComposer> composer = ComposerService::getComposerService();
+    ASSERT_EQ(PERMISSION_DENIED, composer->captureDisplay(captureArgs, mCaptureResults));
+
+    // Make screenshot request with current uid set. No layers were created with the current
+    // uid so screenshot will be black.
+    captureArgs.uid = fakeUid;
+    ScreenCapture::captureDisplay(&mCapture, captureArgs);
+    mCapture->expectColor(Rect(0, 0, 32, 32), Color::BLACK);
+    mCapture->expectBorder(Rect(0, 0, 32, 32), Color::BLACK);
+
+    sp<SurfaceControl> layerWithFakeUid;
+    // Create a new layer with the current uid
+    ASSERT_NO_FATAL_FAILURE(layerWithFakeUid =
+                                    createLayer("new test layer", 32, 32,
+                                                ISurfaceComposerClient::eFXSurfaceBufferQueue,
+                                                mBGSurfaceControl.get()));
+    ASSERT_NO_FATAL_FAILURE(fillBufferQueueLayerColor(layerWithFakeUid, Color::GREEN, 32, 32));
+    Transaction()
+            .show(layerWithFakeUid)
+            .setLayer(layerWithFakeUid, INT32_MAX)
+            .setPosition(layerWithFakeUid, 128, 128)
+            .apply();
+
+    // Screenshot from the fakeUid caller with the uid requested allows the layer
+    // with that uid to be screenshotted. Everything else is black
+    ScreenCapture::captureDisplay(&mCapture, captureArgs);
+    mCapture->expectColor(Rect(128, 128, 160, 160), Color::GREEN);
+    mCapture->expectBorder(Rect(128, 128, 160, 160), Color::BLACK);
+}
+
+TEST_F(ScreenCaptureTest, CaptureLayerWithUid) {
+    uid_t fakeUid = 12345;
+
+    sp<SurfaceControl> layer;
+    ASSERT_NO_FATAL_FAILURE(layer = createLayer("test layer", 32, 32,
+                                                ISurfaceComposerClient::eFXSurfaceBufferQueue,
+                                                mBGSurfaceControl.get()));
+    ASSERT_NO_FATAL_FAILURE(fillBufferQueueLayerColor(layer, Color::RED, 32, 32));
+
+    Transaction().show(layer).setLayer(layer, INT32_MAX).apply();
+
+    LayerCaptureArgs captureArgs;
+    captureArgs.layerHandle = mBGSurfaceControl->getHandle();
+    captureArgs.childrenOnly = false;
+
+    // Make sure red layer with the background layer is screenshot.
+    ScreenCapture::captureLayers(&mCapture, captureArgs);
+    mCapture->expectColor(Rect(0, 0, 32, 32), Color::RED);
+    mCapture->expectBorder(Rect(0, 0, 32, 32), {63, 63, 195, 255});
+
+    // From non system uid, can't request screenshot without a specified uid.
+    std::unique_ptr<UIDFaker> uidFaker = std::make_unique<UIDFaker>(fakeUid);
+
+    sp<ISurfaceComposer> composer = ComposerService::getComposerService();
+    ASSERT_EQ(PERMISSION_DENIED, composer->captureLayers(captureArgs, mCaptureResults));
+
+    // Make screenshot request with current uid set. No layers were created with the current
+    // uid so screenshot will be black.
+    captureArgs.uid = fakeUid;
+    ScreenCapture::captureLayers(&mCapture, captureArgs);
+    mCapture->expectColor(Rect(0, 0, 32, 32), Color::TRANSPARENT);
+    mCapture->expectBorder(Rect(0, 0, 32, 32), Color::TRANSPARENT);
+
+    sp<SurfaceControl> layerWithFakeUid;
+    // Create a new layer with the current uid
+    ASSERT_NO_FATAL_FAILURE(layerWithFakeUid =
+                                    createLayer("new test layer", 32, 32,
+                                                ISurfaceComposerClient::eFXSurfaceBufferQueue,
+                                                mBGSurfaceControl.get()));
+    ASSERT_NO_FATAL_FAILURE(fillBufferQueueLayerColor(layerWithFakeUid, Color::GREEN, 32, 32));
+    Transaction()
+            .show(layerWithFakeUid)
+            .setLayer(layerWithFakeUid, INT32_MAX)
+            .setPosition(layerWithFakeUid, 128, 128)
+            // reparent a layer that was created with a different uid to the new layer.
+            .reparent(layer, layerWithFakeUid->getHandle())
+            .apply();
+
+    // Screenshot from the fakeUid caller with the uid requested allows the layer
+    // with that uid to be screenshotted. The child layer is skipped since it was created
+    // from a different uid.
+    ScreenCapture::captureLayers(&mCapture, captureArgs);
+    mCapture->expectColor(Rect(128, 128, 160, 160), Color::GREEN);
+    mCapture->expectBorder(Rect(128, 128, 160, 160), Color::TRANSPARENT);
+
+    // Clear fake calling uid so it's back to system.
+    uidFaker = nullptr;
+    // Screenshot from the test caller with the uid requested allows the layer
+    // with that uid to be screenshotted. The child layer is skipped since it was created
+    // from a different uid.
+    ScreenCapture::captureLayers(&mCapture, captureArgs);
+    mCapture->expectColor(Rect(128, 128, 160, 160), Color::GREEN);
+    mCapture->expectBorder(Rect(128, 128, 160, 160), Color::TRANSPARENT);
+
+    // Screenshot from the fakeUid caller with no uid requested allows everything to be screenshot.
+    captureArgs.uid = -1;
+    ScreenCapture::captureLayers(&mCapture, captureArgs);
+    mCapture->expectColor(Rect(128, 128, 160, 160), Color::RED);
+    mCapture->expectBorder(Rect(128, 128, 160, 160), {63, 63, 195, 255});
+}
+
 // In the following tests we verify successful skipping of a parent layer,
 // so we use the same verification logic and only change how we mutate
 // the parent layer to verify that various properties are ignored.
diff --git a/services/surfaceflinger/tests/unittests/CompositionTest.cpp b/services/surfaceflinger/tests/unittests/CompositionTest.cpp
index 48c4e18..dc8b31e 100644
--- a/services/surfaceflinger/tests/unittests/CompositionTest.cpp
+++ b/services/surfaceflinger/tests/unittests/CompositionTest.cpp
@@ -241,7 +241,8 @@
                                                 ui::Dataspace::V0_SRGB, ui::Transform::ROT_0);
 
     auto traverseLayers = [this](const LayerVector::Visitor& visitor) {
-        return mFlinger.traverseLayersInLayerStack(mDisplay->getLayerStack(), visitor);
+        return mFlinger.traverseLayersInLayerStack(mDisplay->getLayerStack(),
+                                                   CaptureArgs::UNSET_UID, visitor);
     };
 
     // TODO: Eliminate expensive/real allocation if possible.
diff --git a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
index 3941d42..8972907 100644
--- a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
+++ b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
@@ -343,9 +343,9 @@
                                                 regionSampling, captureResults);
     }
 
-    auto traverseLayersInLayerStack(ui::LayerStack layerStack,
+    auto traverseLayersInLayerStack(ui::LayerStack layerStack, int32_t uid,
                                     const LayerVector::Visitor& visitor) {
-        return mFlinger->SurfaceFlinger::traverseLayersInLayerStack(layerStack, visitor);
+        return mFlinger->SurfaceFlinger::traverseLayersInLayerStack(layerStack, uid, visitor);
     }
 
     auto getDisplayNativePrimaries(const sp<IBinder>& displayToken,
diff --git a/services/surfaceflinger/tests/utils/ScreenshotUtils.h b/services/surfaceflinger/tests/utils/ScreenshotUtils.h
index 56628f8..081d18b 100644
--- a/services/surfaceflinger/tests/utils/ScreenshotUtils.h
+++ b/services/surfaceflinger/tests/utils/ScreenshotUtils.h
@@ -31,16 +31,20 @@
         captureScreen(sc, SurfaceComposerClient::getInternalDisplayToken());
     }
 
-    static void captureScreen(std::unique_ptr<ScreenCapture>* sc, sp<IBinder> displayToken) {
+    static void captureDisplay(std::unique_ptr<ScreenCapture>* sc,
+                               const DisplayCaptureArgs& captureArgs) {
         const auto sf = ComposerService::getComposerService();
         SurfaceComposerClient::Transaction().apply(true);
 
+        ScreenCaptureResults captureResults;
+        ASSERT_EQ(NO_ERROR, sf->captureDisplay(captureArgs, captureResults));
+        *sc = std::make_unique<ScreenCapture>(captureResults.buffer);
+    }
+
+    static void captureScreen(std::unique_ptr<ScreenCapture>* sc, sp<IBinder> displayToken) {
         DisplayCaptureArgs args;
         args.displayToken = displayToken;
-
-        ScreenCaptureResults captureResults;
-        ASSERT_EQ(NO_ERROR, sf->captureDisplay(args, captureResults));
-        *sc = std::make_unique<ScreenCapture>(captureResults.buffer);
+        captureDisplay(sc, args);
     }
 
     static void captureLayers(std::unique_ptr<ScreenCapture>* sc,