Merge "inputflinger_input_dispatcher_fuzzer : Fix memory leak" into main
diff --git a/libs/input/input_flags.aconfig b/libs/input/input_flags.aconfig
index 4e187ed..5bb30db 100644
--- a/libs/input/input_flags.aconfig
+++ b/libs/input/input_flags.aconfig
@@ -255,3 +255,12 @@
   bug: "277261245"
 }
 
+flag {
+  name: "prevent_merging_input_pointer_devices"
+  namespace: "desktop_input"
+  description: "Prevent merging input sub-devices that provide pointer input streams"
+  bug: "389689566"
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
+}
diff --git a/services/gpuservice/Android.bp b/services/gpuservice/Android.bp
index 689221f..01287b0 100644
--- a/services/gpuservice/Android.bp
+++ b/services/gpuservice/Android.bp
@@ -7,6 +7,13 @@
     default_applicable_licenses: ["frameworks_native_license"],
 }
 
+aconfig_declarations {
+    name: "gpuservice_flags",
+    package: "com.android.frameworks.gpuservice.flags",
+    container: "system",
+    srcs: ["gpuservice_flags.aconfig"],
+}
+
 cc_defaults {
     name: "gpuservice_defaults",
     cflags: [
@@ -20,6 +27,11 @@
 }
 
 cc_aconfig_library {
+    name: "gpuservice_multiuser_flags_c_lib",
+    aconfig_declarations: "gpuservice_flags",
+}
+
+cc_aconfig_library {
     name: "gpuservice_flags_c_lib",
     aconfig_declarations: "graphicsenv_flags",
 }
@@ -92,6 +104,9 @@
     srcs: [
         ":libgpuservice_sources",
     ],
+    shared_libs: [
+        "gpuservice_multiuser_flags_c_lib",
+    ],
 }
 
 cc_defaults {
@@ -126,4 +141,7 @@
     static_libs: [
         "libgpuservice",
     ],
+    shared_libs: [
+        "gpuservice_multiuser_flags_c_lib",
+    ],
 }
diff --git a/services/gpuservice/GpuService.cpp b/services/gpuservice/GpuService.cpp
index fadb1fd..f74b4fa 100644
--- a/services/gpuservice/GpuService.cpp
+++ b/services/gpuservice/GpuService.cpp
@@ -24,7 +24,9 @@
 #include <binder/IResultReceiver.h>
 #include <binder/Parcel.h>
 #include <binder/PermissionCache.h>
+#include <com_android_frameworks_gpuservice_flags.h>
 #include <cutils/properties.h>
+#include <cutils/multiuser.h>
 #include <gpumem/GpuMem.h>
 #include <gpuwork/GpuWork.h>
 #include <gpustats/GpuStats.h>
@@ -38,6 +40,8 @@
 #include <thread>
 #include <memory>
 
+namespace gpuservice_flags = com::android::frameworks::gpuservice::flags;
+
 namespace android {
 
 using base::StringAppendF;
@@ -113,11 +117,22 @@
 
     // only system_server with the ACCESS_GPU_SERVICE permission is allowed to set
     // persist.graphics.egl
-    if (uid != AID_SYSTEM ||
-        !PermissionCache::checkPermission(sAccessGpuServicePermission, pid, uid)) {
-        ALOGE("Permission Denial: can't set persist.graphics.egl from setAngleAsSystemDriver() "
+    if (gpuservice_flags::multiuser_permission_check()) {
+        // retrieve the appid of Settings app on multiuser builds
+        const int multiuserappid = multiuser_get_app_id(uid);
+        if (multiuserappid != AID_SYSTEM ||
+            !PermissionCache::checkPermission(sAccessGpuServicePermission, pid, uid)) {
+            ALOGE("Permission Denial: can't set persist.graphics.egl from setAngleAsSystemDriver() "
+                "pid=%d, uid=%d\n, multiuserappid=%d", pid, uid, multiuserappid);
+            return;
+        }
+    } else {
+        if (uid != AID_SYSTEM ||
+            !PermissionCache::checkPermission(sAccessGpuServicePermission, pid, uid)) {
+            ALOGE("Permission Denial: can't set persist.graphics.egl from setAngleAsSystemDriver() "
                 "pid=%d, uid=%d\n", pid, uid);
-        return;
+            return;
+        }
     }
 
     std::lock_guard<std::mutex> lock(mLock);
diff --git a/services/gpuservice/gpuservice_flags.aconfig b/services/gpuservice/gpuservice_flags.aconfig
new file mode 100644
index 0000000..be6a7bb
--- /dev/null
+++ b/services/gpuservice/gpuservice_flags.aconfig
@@ -0,0 +1,12 @@
+package: "com.android.frameworks.gpuservice.flags"
+container: "system"
+
+flag {
+    name: "multiuser_permission_check"
+    namespace: "gpu"
+    description: "Whether to consider headless system user mode/multiuser when checking toggleAngleAsSystemDriver permission."
+    bug: "389867658"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
+}
diff --git a/services/gpuservice/tests/fuzzers/Android.bp b/services/gpuservice/tests/fuzzers/Android.bp
index d4d48c4..7be3253 100644
--- a/services/gpuservice/tests/fuzzers/Android.bp
+++ b/services/gpuservice/tests/fuzzers/Android.bp
@@ -13,6 +13,9 @@
         "libgpuservice",
         "liblog",
     ],
+    shared_libs: [
+        "gpuservice_multiuser_flags_c_lib",
+    ],
     fuzz_config: {
         cc: [
             "paulthomson@google.com",
diff --git a/services/gpuservice/tests/unittests/Android.bp b/services/gpuservice/tests/unittests/Android.bp
index d2184d8..0dac24d 100644
--- a/services/gpuservice/tests/unittests/Android.bp
+++ b/services/gpuservice/tests/unittests/Android.bp
@@ -89,6 +89,7 @@
     ],
     header_libs: ["bpf_headers"],
     shared_libs: [
+        "gpuservice_multiuser_flags_c_lib",
         "libbase",
         "libbinder",
         "libbpf_bcc",
diff --git a/services/inputflinger/reader/InputReader.cpp b/services/inputflinger/reader/InputReader.cpp
index 207806d..7ab000b 100644
--- a/services/inputflinger/reader/InputReader.cpp
+++ b/services/inputflinger/reader/InputReader.cpp
@@ -62,6 +62,28 @@
             identifier1.location == identifier2.location);
 }
 
+/**
+ * Determines if the device classes passed for two devices represent incompatible combinations
+ * that should not be merged into into a single InputDevice.
+ */
+
+bool isCompatibleSubDevice(ftl::Flags<InputDeviceClass> classes1,
+                           ftl::Flags<InputDeviceClass> classes2) {
+    if (!com::android::input::flags::prevent_merging_input_pointer_devices()) {
+        return true;
+    }
+
+    const ftl::Flags<InputDeviceClass> pointerFlags =
+            ftl::Flags<InputDeviceClass>{InputDeviceClass::TOUCH, InputDeviceClass::TOUCH_MT,
+                                         InputDeviceClass::CURSOR, InputDeviceClass::TOUCHPAD};
+
+    // Do not merge devices that both have any type of pointer event.
+    if (classes1.any(pointerFlags) && classes2.any(pointerFlags)) return false;
+
+    // Safe to merge
+    return true;
+}
+
 bool isStylusPointerGestureStart(const NotifyMotionArgs& motionArgs) {
     const auto actionMasked = MotionEvent::getActionMasked(motionArgs.action);
     if (actionMasked != AMOTION_EVENT_ACTION_HOVER_ENTER &&
@@ -271,7 +293,8 @@
     }
 
     InputDeviceIdentifier identifier = mEventHub->getDeviceIdentifier(eventHubId);
-    std::shared_ptr<InputDevice> device = createDeviceLocked(when, eventHubId, identifier);
+    ftl::Flags<InputDeviceClass> classes = mEventHub->getDeviceClasses(eventHubId);
+    std::shared_ptr<InputDevice> device = createDeviceLocked(when, eventHubId, identifier, classes);
 
     mPendingArgs += device->configure(when, mConfig, /*changes=*/{});
     mPendingArgs += device->reset(when);
@@ -354,12 +377,16 @@
 }
 
 std::shared_ptr<InputDevice> InputReader::createDeviceLocked(
-        nsecs_t when, int32_t eventHubId, const InputDeviceIdentifier& identifier) {
-    auto deviceIt = std::find_if(mDevices.begin(), mDevices.end(), [identifier](auto& devicePair) {
-        const InputDeviceIdentifier identifier2 =
-                devicePair.second->getDeviceInfo().getIdentifier();
-        return isSubDevice(identifier, identifier2);
-    });
+        nsecs_t when, int32_t eventHubId, const InputDeviceIdentifier& identifier,
+        ftl::Flags<InputDeviceClass> classes) {
+    auto deviceIt =
+            std::find_if(mDevices.begin(), mDevices.end(), [identifier, classes](auto& devicePair) {
+                const InputDeviceIdentifier identifier2 =
+                        devicePair.second->getDeviceInfo().getIdentifier();
+                const ftl::Flags<InputDeviceClass> classes2 = devicePair.second->getClasses();
+                return isSubDevice(identifier, identifier2) &&
+                        isCompatibleSubDevice(classes, classes2);
+            });
 
     std::shared_ptr<InputDevice> device;
     if (deviceIt != mDevices.end()) {
diff --git a/services/inputflinger/reader/include/InputReader.h b/services/inputflinger/reader/include/InputReader.h
index 931766b..0d6e102 100644
--- a/services/inputflinger/reader/include/InputReader.h
+++ b/services/inputflinger/reader/include/InputReader.h
@@ -25,6 +25,7 @@
 #include <vector>
 
 #include "EventHub.h"
+#include "InputDevice.h"
 #include "InputListener.h"
 #include "InputReaderBase.h"
 #include "InputReaderContext.h"
@@ -127,7 +128,8 @@
 protected:
     // These members are protected so they can be instrumented by test cases.
     virtual std::shared_ptr<InputDevice> createDeviceLocked(nsecs_t when, int32_t deviceId,
-                                                            const InputDeviceIdentifier& identifier)
+                                                            const InputDeviceIdentifier& identifier,
+                                                            ftl::Flags<InputDeviceClass> classes)
             REQUIRES(mLock);
 
     // With each iteration of the loop, InputReader reads and processes one incoming message from
diff --git a/services/inputflinger/tests/InputReader_test.cpp b/services/inputflinger/tests/InputReader_test.cpp
index d9a75a5..43d2378 100644
--- a/services/inputflinger/tests/InputReader_test.cpp
+++ b/services/inputflinger/tests/InputReader_test.cpp
@@ -1405,6 +1405,68 @@
     ASSERT_EQ(mFakeEventHub->fakeReadKernelWakeup(3), false);
 }
 
+TEST_F(InputReaderTest, MergeableInputDevices) {
+    constexpr int32_t eventHubIds[2] = {END_RESERVED_ID, END_RESERVED_ID + 1};
+
+    // By default, all of the default-created eventhub devices will have the same identifier
+    // (implicitly vid 0, pid 0, etc.), which is why we expect them to be merged.
+    ASSERT_NO_FATAL_FAILURE(addDevice(eventHubIds[0], "1st", InputDeviceClass::KEYBOARD, nullptr));
+    ASSERT_NO_FATAL_FAILURE(addDevice(eventHubIds[1], "2nd", InputDeviceClass::JOYSTICK, nullptr));
+
+    // The two devices will be merged to one input device as they have same identifier, and none are
+    // pointer devices.
+    ASSERT_EQ(1U, mFakePolicy->getInputDevices().size());
+}
+
+TEST_F(InputReaderTest, MergeableDevicesWithTouch) {
+    constexpr int32_t eventHubIds[3] = {END_RESERVED_ID, END_RESERVED_ID + 1, END_RESERVED_ID + 2};
+
+    // By default, all of the default-created eventhub devices will have the same identifier
+    // (implicitly vid 0, pid 0, etc.), which is why we expect them to be merged.
+    ASSERT_NO_FATAL_FAILURE(addDevice(eventHubIds[0], "1st", InputDeviceClass::TOUCH_MT, nullptr));
+    ASSERT_NO_FATAL_FAILURE(addDevice(eventHubIds[1], "2nd", InputDeviceClass::KEYBOARD, nullptr));
+    ASSERT_NO_FATAL_FAILURE(addDevice(eventHubIds[2], "3rd", InputDeviceClass::GAMEPAD, nullptr));
+
+    // The three devices will be merged to one input device as they have same identifier, and only
+    // one is a pointer device.
+    ASSERT_EQ(1U, mFakePolicy->getInputDevices().size());
+}
+
+TEST_F(InputReaderTest, UnmergeableTouchDevices) {
+    SCOPED_FLAG_OVERRIDE(prevent_merging_input_pointer_devices, true);
+
+    constexpr int32_t eventHubIds[3] = {END_RESERVED_ID, END_RESERVED_ID + 1, END_RESERVED_ID + 2};
+
+    // By default, all of the default-created eventhub devices will have the same identifier
+    // (implicitly vid 0, pid 0, etc.), which is why they can potentially be merged.
+    ASSERT_NO_FATAL_FAILURE(addDevice(eventHubIds[0], "1st", InputDeviceClass::TOUCH, nullptr));
+    ASSERT_NO_FATAL_FAILURE(addDevice(eventHubIds[1], "2nd", InputDeviceClass::TOUCH_MT, nullptr));
+    ASSERT_NO_FATAL_FAILURE(addDevice(eventHubIds[2], "2nd", InputDeviceClass::CURSOR, nullptr));
+
+    // The three devices will not be merged, as they have same identifier, but are all pointer
+    // devices.
+    ASSERT_EQ(3U, mFakePolicy->getInputDevices().size());
+}
+
+TEST_F(InputReaderTest, MergeableMixedDevices) {
+    SCOPED_FLAG_OVERRIDE(prevent_merging_input_pointer_devices, true);
+
+    constexpr int32_t eventHubIds[4] = {END_RESERVED_ID, END_RESERVED_ID + 1, END_RESERVED_ID + 2,
+                                        END_RESERVED_ID + 3};
+
+    // By default, all of the default-created eventhub devices will have the same identifier
+    // (implicitly vid 0, pid 0, etc.), which is why they can potentially be merged.
+    ASSERT_NO_FATAL_FAILURE(addDevice(eventHubIds[0], "1st", InputDeviceClass::TOUCH, nullptr));
+    ASSERT_NO_FATAL_FAILURE(addDevice(eventHubIds[1], "2nd", InputDeviceClass::TOUCH_MT, nullptr));
+    ASSERT_NO_FATAL_FAILURE(addDevice(eventHubIds[2], "3rd", InputDeviceClass::DPAD, nullptr));
+    ASSERT_NO_FATAL_FAILURE(addDevice(eventHubIds[3], "4th", InputDeviceClass::JOYSTICK, nullptr));
+
+    // Non-touch devices can be merged with one of the touch devices, as they have same identifier,
+    // but the two touch devices will not combine with each other. It is not specified which touch
+    // device the non-touch devices merge with.
+    ASSERT_EQ(2U, mFakePolicy->getInputDevices().size());
+}
+
 // --- InputReaderIntegrationTest ---
 
 // These tests create and interact with the InputReader only through its interface.
diff --git a/services/inputflinger/tests/InstrumentedInputReader.cpp b/services/inputflinger/tests/InstrumentedInputReader.cpp
index 110ca5f..53fc8a1 100644
--- a/services/inputflinger/tests/InstrumentedInputReader.cpp
+++ b/services/inputflinger/tests/InstrumentedInputReader.cpp
@@ -38,13 +38,14 @@
 }
 
 std::shared_ptr<InputDevice> InstrumentedInputReader::createDeviceLocked(
-        nsecs_t when, int32_t eventHubId, const InputDeviceIdentifier& identifier) REQUIRES(mLock) {
+        nsecs_t when, int32_t eventHubId, const InputDeviceIdentifier& identifier,
+        ftl::Flags<InputDeviceClass> classes) REQUIRES(mLock) {
     if (!mNextDevices.empty()) {
         std::shared_ptr<InputDevice> device(std::move(mNextDevices.front()));
         mNextDevices.pop();
         return device;
     }
-    return InputReader::createDeviceLocked(when, eventHubId, identifier);
+    return InputReader::createDeviceLocked(when, eventHubId, identifier, classes);
 }
 
 } // namespace android
diff --git a/services/inputflinger/tests/InstrumentedInputReader.h b/services/inputflinger/tests/InstrumentedInputReader.h
index e9c7bb4..9abf30c 100644
--- a/services/inputflinger/tests/InstrumentedInputReader.h
+++ b/services/inputflinger/tests/InstrumentedInputReader.h
@@ -43,8 +43,9 @@
     using InputReader::loopOnce;
 
 protected:
-    virtual std::shared_ptr<InputDevice> createDeviceLocked(
-            nsecs_t when, int32_t eventHubId, const InputDeviceIdentifier& identifier);
+    virtual std::shared_ptr<InputDevice> createDeviceLocked(nsecs_t when, int32_t eventHubId,
+                                                            const InputDeviceIdentifier& identifier,
+                                                            ftl::Flags<InputDeviceClass> classes);
 
     class FakeInputReaderContext : public ContextImpl {
     public:
diff --git a/services/surfaceflinger/RegionSamplingThread.cpp b/services/surfaceflinger/RegionSamplingThread.cpp
index 1c4a11a..514adac 100644
--- a/services/surfaceflinger/RegionSamplingThread.cpp
+++ b/services/surfaceflinger/RegionSamplingThread.cpp
@@ -356,12 +356,10 @@
     screenshotArgs.seamlessTransition = false;
 
     std::vector<std::pair<Layer*, sp<LayerFE>>> layers;
-    auto displayState =
-            mFlinger.getSnapshotsFromMainThread(screenshotArgs, getLayerSnapshotsFn, layers);
-    FenceResult fenceResult =
-            mFlinger.captureScreenshot(screenshotArgs, buffer, kRegionSampling, kGrayscale,
-                                       kIsProtected, nullptr, displayState, layers)
-                    .get();
+    mFlinger.getSnapshotsFromMainThread(screenshotArgs, getLayerSnapshotsFn, layers);
+    FenceResult fenceResult = mFlinger.captureScreenshot(screenshotArgs, buffer, kRegionSampling,
+                                                         kGrayscale, kIsProtected, nullptr, layers)
+                                      .get();
     if (fenceResult.ok()) {
         fenceResult.value()->waitForever(LOG_TAG);
     }
diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index 1163390..6bc3aad 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -7196,14 +7196,13 @@
 
 namespace {
 
-ui::Dataspace pickBestDataspace(ui::Dataspace requestedDataspace,
-                                const compositionengine::impl::OutputCompositionState& state,
+ui::Dataspace pickBestDataspace(ui::Dataspace requestedDataspace, ui::ColorMode colorMode,
                                 bool capturingHdrLayers, bool hintForSeamlessTransition) {
     if (requestedDataspace != ui::Dataspace::UNKNOWN) {
         return requestedDataspace;
     }
 
-    const auto dataspaceForColorMode = ui::pickDataspaceFor(state.colorMode);
+    const auto dataspaceForColorMode = ui::pickDataspaceFor(colorMode);
 
     // TODO: Enable once HDR screenshots are ready.
     if constexpr (/* DISABLES CODE */ (false)) {
@@ -7515,11 +7514,11 @@
     return protectedLayerFound;
 }
 
-// Getting layer snapshots and display should take place on main thread.
-// Accessing display requires mStateLock, and contention for this lock
-// is reduced when grabbed from the main thread, thus also reducing
-// risk of deadlocks.
-std::optional<SurfaceFlinger::OutputCompositionState> SurfaceFlinger::getSnapshotsFromMainThread(
+// Getting layer snapshots and accessing display state should take place on
+// main thread. Accessing display requires mStateLock, and contention for
+// this lock is reduced when grabbed from the main thread, thus also reducing
+// risk of deadlocks. Returns false if no display is found.
+bool SurfaceFlinger::getSnapshotsFromMainThread(
         ScreenshotArgs& args, GetLayerSnapshotsFunction getLayerSnapshotsFn,
         std::vector<std::pair<Layer*, sp<LayerFE>>>& layers) {
     return mScheduler
@@ -7559,8 +7558,8 @@
     }
 
     std::vector<std::pair<Layer*, sp<LayerFE>>> layers;
-    auto displayState = getSnapshotsFromMainThread(args, getLayerSnapshotsFn, layers);
-    if (!displayState) {
+    bool hasDisplayState = getSnapshotsFromMainThread(args, getLayerSnapshotsFn, layers);
+    if (!hasDisplayState) {
         ALOGD("Display state not found");
         invokeScreenCaptureError(NO_MEMORY, captureListener);
     }
@@ -7637,12 +7636,13 @@
 
     auto futureFence =
             captureScreenshot(args, texture, false /* regionSampling */, grayscale, isProtected,
-                              captureListener, displayState, layers, hdrTexture, gainmapTexture);
+                              captureListener, layers, hdrTexture, gainmapTexture);
     futureFence.get();
 }
 
-std::optional<SurfaceFlinger::OutputCompositionState> SurfaceFlinger::getDisplayStateOnMainThread(
-        ScreenshotArgs& args) {
+// Returns true if display is found and args was populated with display state
+// data. Otherwise, returns false.
+bool SurfaceFlinger::getDisplayStateOnMainThread(ScreenshotArgs& args) {
     sp<const DisplayDevice> display = nullptr;
     {
         Mutex::Autolock lock(mStateLock);
@@ -7677,49 +7677,49 @@
         }
 
         if (display != nullptr) {
-            return std::optional{display->getCompositionDisplay()->getState()};
+            const auto& state = display->getCompositionDisplay()->getState();
+            args.displayBrightnessNits = state.displayBrightnessNits;
+            args.sdrWhitePointNits = state.sdrWhitePointNits;
+            args.renderIntent = state.renderIntent;
+            args.colorMode = state.colorMode;
+            return true;
         }
     }
-    return std::nullopt;
+    return false;
 }
 
 ftl::SharedFuture<FenceResult> SurfaceFlinger::captureScreenshot(
-        const ScreenshotArgs& args, const std::shared_ptr<renderengine::ExternalTexture>& buffer,
+        ScreenshotArgs& args, const std::shared_ptr<renderengine::ExternalTexture>& buffer,
         bool regionSampling, bool grayscale, bool isProtected,
         const sp<IScreenCaptureListener>& captureListener,
-        const std::optional<OutputCompositionState>& displayState,
         const std::vector<std::pair<Layer*, sp<LayerFE>>>& layers,
         const std::shared_ptr<renderengine::ExternalTexture>& hdrBuffer,
         const std::shared_ptr<renderengine::ExternalTexture>& gainmapBuffer) {
     SFTRACE_CALL();
 
     ScreenCaptureResults captureResults;
-
-    float displayBrightnessNits = displayState.value().displayBrightnessNits;
-    float sdrWhitePointNits = displayState.value().sdrWhitePointNits;
-
     ftl::SharedFuture<FenceResult> renderFuture;
 
+    float hdrSdrRatio = args.displayBrightnessNits / args.sdrWhitePointNits;
+
     if (hdrBuffer && gainmapBuffer) {
         ftl::SharedFuture<FenceResult> hdrRenderFuture =
                 renderScreenImpl(args, hdrBuffer, regionSampling, grayscale, isProtected,
-                                 captureResults, displayState, layers);
+                                 captureResults, layers);
         captureResults.buffer = buffer->getBuffer();
         captureResults.optionalGainMap = gainmapBuffer->getBuffer();
 
         renderFuture =
                 ftl::Future(std::move(hdrRenderFuture))
-                        .then([&, displayBrightnessNits, sdrWhitePointNits,
-                               dataspace = captureResults.capturedDataspace, buffer, hdrBuffer,
-                               gainmapBuffer](FenceResult fenceResult) -> FenceResult {
+                        .then([&, hdrSdrRatio, dataspace = captureResults.capturedDataspace, buffer,
+                               hdrBuffer, gainmapBuffer](FenceResult fenceResult) -> FenceResult {
                             if (!fenceResult.ok()) {
                                 return fenceResult;
                             }
 
                             return getRenderEngine()
                                     .tonemapAndDrawGainmap(hdrBuffer, fenceResult.value()->get(),
-                                                           displayBrightnessNits /
-                                                                   sdrWhitePointNits,
+                                                           hdrSdrRatio,
                                                            static_cast<ui::Dataspace>(dataspace),
                                                            buffer, gainmapBuffer)
                                     .get();
@@ -7727,17 +7727,16 @@
                         .share();
     } else {
         renderFuture = renderScreenImpl(args, buffer, regionSampling, grayscale, isProtected,
-                                        captureResults, displayState, layers);
+                                        captureResults, layers);
     }
 
     if (captureListener) {
         // Defer blocking on renderFuture back to the Binder thread.
         return ftl::Future(std::move(renderFuture))
                 .then([captureListener, captureResults = std::move(captureResults),
-                       displayBrightnessNits,
-                       sdrWhitePointNits](FenceResult fenceResult) mutable -> FenceResult {
+                       hdrSdrRatio](FenceResult fenceResult) mutable -> FenceResult {
                     captureResults.fenceResult = std::move(fenceResult);
-                    captureResults.hdrSdrRatio = displayBrightnessNits / sdrWhitePointNits;
+                    captureResults.hdrSdrRatio = hdrSdrRatio;
                     captureListener->onScreenCaptureCompleted(captureResults);
                     return base::unexpected(NO_ERROR);
                 })
@@ -7747,9 +7746,8 @@
 }
 
 ftl::SharedFuture<FenceResult> SurfaceFlinger::renderScreenImpl(
-        const ScreenshotArgs& args, const std::shared_ptr<renderengine::ExternalTexture>& buffer,
+        ScreenshotArgs& args, const std::shared_ptr<renderengine::ExternalTexture>& buffer,
         bool regionSampling, bool grayscale, bool isProtected, ScreenCaptureResults& captureResults,
-        const std::optional<OutputCompositionState>& displayState,
         const std::vector<std::pair<Layer*, sp<LayerFE>>>& layers) {
     SFTRACE_CALL();
 
@@ -7763,49 +7761,37 @@
                 layerFE->mSnapshot->geomLayerTransform.inverse();
     }
 
-    auto capturedBuffer = buffer;
-
-    auto renderIntent = RenderIntent::TONE_MAP_COLORIMETRIC;
-    auto sdrWhitePointNits = DisplayDevice::sDefaultMaxLumiance;
-    auto displayBrightnessNits = DisplayDevice::sDefaultMaxLumiance;
-
-    captureResults.capturedDataspace = args.dataspace;
-
     const bool enableLocalTonemapping =
             FlagManager::getInstance().local_tonemap_screenshots() && !args.seamlessTransition;
 
-    if (displayState) {
-        const auto& state = displayState.value();
-        captureResults.capturedDataspace =
-                pickBestDataspace(args.dataspace, state, captureResults.capturedHdrLayers,
-                                  args.seamlessTransition);
-        sdrWhitePointNits = state.sdrWhitePointNits;
+    captureResults.capturedDataspace =
+            pickBestDataspace(args.dataspace, args.colorMode, captureResults.capturedHdrLayers,
+                              args.seamlessTransition);
 
-        if (!captureResults.capturedHdrLayers) {
-            displayBrightnessNits = sdrWhitePointNits;
-        } else {
-            displayBrightnessNits = state.displayBrightnessNits;
-            if (!enableLocalTonemapping) {
-                // Only clamp the display brightness if this is not a seamless transition.
-                // Otherwise for seamless transitions it's important to match the current
-                // display state as the buffer will be shown under these same conditions, and we
-                // want to avoid any flickers
-                if (sdrWhitePointNits > 1.0f && !args.seamlessTransition) {
-                    // Restrict the amount of HDR "headroom" in the screenshot to avoid
-                    // over-dimming the SDR portion. 2.0 chosen by experimentation
-                    constexpr float kMaxScreenshotHeadroom = 2.0f;
-                    displayBrightnessNits = std::min(sdrWhitePointNits * kMaxScreenshotHeadroom,
-                                                     displayBrightnessNits);
-                }
-            }
-        }
-
-        // Screenshots leaving the device should be colorimetric
-        if (args.dataspace == ui::Dataspace::UNKNOWN && args.seamlessTransition) {
-            renderIntent = state.renderIntent;
-        }
+    // Only clamp the display brightness if this is not a seamless transition.
+    // Otherwise for seamless transitions it's important to match the current
+    // display state as the buffer will be shown under these same conditions, and we
+    // want to avoid any flickers.
+    if (captureResults.capturedHdrLayers && !enableLocalTonemapping &&
+        args.sdrWhitePointNits > 1.0f && !args.seamlessTransition) {
+        // Restrict the amount of HDR "headroom" in the screenshot to avoid
+        // over-dimming the SDR portion. 2.0 chosen by experimentation
+        constexpr float kMaxScreenshotHeadroom = 2.0f;
+        // TODO: Aim to update displayBrightnessNits earlier in screenshot
+        // path so ScreenshotArgs can be passed as const
+        args.displayBrightnessNits = std::min(args.sdrWhitePointNits * kMaxScreenshotHeadroom,
+                                              args.displayBrightnessNits);
+    } else {
+        args.displayBrightnessNits = args.sdrWhitePointNits;
     }
 
+    auto renderIntent = RenderIntent::TONE_MAP_COLORIMETRIC;
+    // Screenshots leaving the device should be colorimetric
+    if (args.dataspace == ui::Dataspace::UNKNOWN && args.seamlessTransition) {
+        renderIntent = args.renderIntent;
+    }
+
+    auto capturedBuffer = buffer;
     captureResults.buffer = capturedBuffer->getBuffer();
 
     ui::LayerStack layerStack{ui::DEFAULT_LAYER_STACK};
@@ -7815,8 +7801,7 @@
     }
 
     auto present = [this, buffer = capturedBuffer, dataspace = captureResults.capturedDataspace,
-                    sdrWhitePointNits, displayBrightnessNits, grayscale, isProtected, layers,
-                    layerStack, regionSampling, args, renderIntent,
+                    grayscale, isProtected, layers, layerStack, regionSampling, args, renderIntent,
                     enableLocalTonemapping]() -> FenceResult {
         std::unique_ptr<compositionengine::CompositionEngine> compositionEngine =
                 mFactory.createCompositionEngine();
@@ -7842,9 +7827,10 @@
         if (enableLocalTonemapping) {
             // Boost the whole scene so that SDR white is at 1.0 while still communicating the hdr
             // sdr ratio via display brightness / sdrWhite nits.
-            targetBrightness = sdrWhitePointNits / displayBrightnessNits;
+            targetBrightness = args.sdrWhitePointNits / args.displayBrightnessNits;
         } else if (dataspace == ui::Dataspace::BT2020_HLG) {
-            const float maxBrightnessNits = displayBrightnessNits / sdrWhitePointNits * 203;
+            const float maxBrightnessNits =
+                    args.displayBrightnessNits / args.sdrWhitePointNits * 203;
             // With a low dimming ratio, don't fit the entire curve. Otherwise mixed content
             // will appear way too bright.
             if (maxBrightnessNits < 1000.f) {
@@ -7870,8 +7856,8 @@
                                         .buffer = std::move(buffer),
                                         .displayId = args.displayId,
                                         .reqBufferSize = args.reqSize,
-                                        .sdrWhitePointNits = sdrWhitePointNits,
-                                        .displayBrightnessNits = displayBrightnessNits,
+                                        .sdrWhitePointNits = args.sdrWhitePointNits,
+                                        .displayBrightnessNits = args.displayBrightnessNits,
                                         .targetBrightness = targetBrightness,
                                         .layerAlpha = layerAlpha,
                                         .regionSampling = regionSampling,
diff --git a/services/surfaceflinger/SurfaceFlinger.h b/services/surfaceflinger/SurfaceFlinger.h
index 935a2da..3f454ba 100644
--- a/services/surfaceflinger/SurfaceFlinger.h
+++ b/services/surfaceflinger/SurfaceFlinger.h
@@ -896,32 +896,41 @@
         // If true, the render result may be used for system animations
         // that must preserve the exact colors of the display
         bool seamlessTransition{false};
+
+        // Current display brightness of the output composition state
+        float displayBrightnessNits{-1.f};
+
+        // SDR white point of the output composition state
+        float sdrWhitePointNits{-1.f};
+
+        // Current active color mode of the output composition state
+        ui::ColorMode colorMode{ui::ColorMode::NATIVE};
+
+        // Current active render intent of the output composition state
+        ui::RenderIntent renderIntent{ui::RenderIntent::COLORIMETRIC};
     };
 
-    std::optional<OutputCompositionState> getSnapshotsFromMainThread(
-            ScreenshotArgs& args, GetLayerSnapshotsFunction getLayerSnapshotsFn,
-            std::vector<std::pair<Layer*, sp<LayerFE>>>& layers);
+    bool getSnapshotsFromMainThread(ScreenshotArgs& args,
+                                    GetLayerSnapshotsFunction getLayerSnapshotsFn,
+                                    std::vector<std::pair<Layer*, sp<LayerFE>>>& layers);
 
     void captureScreenCommon(ScreenshotArgs& args, GetLayerSnapshotsFunction, ui::Size bufferSize,
                              ui::PixelFormat, bool allowProtected, bool grayscale,
                              const sp<IScreenCaptureListener>&);
 
-    std::optional<OutputCompositionState> getDisplayStateOnMainThread(ScreenshotArgs& args)
-            REQUIRES(kMainThreadContext);
+    bool getDisplayStateOnMainThread(ScreenshotArgs& args) REQUIRES(kMainThreadContext);
 
     ftl::SharedFuture<FenceResult> captureScreenshot(
-            const ScreenshotArgs& args,
-            const std::shared_ptr<renderengine::ExternalTexture>& buffer, bool regionSampling,
-            bool grayscale, bool isProtected, const sp<IScreenCaptureListener>& captureListener,
-            const std::optional<OutputCompositionState>& displayState,
+            ScreenshotArgs& args, const std::shared_ptr<renderengine::ExternalTexture>& buffer,
+            bool regionSampling, bool grayscale, bool isProtected,
+            const sp<IScreenCaptureListener>& captureListener,
             const std::vector<std::pair<Layer*, sp<LayerFE>>>& layers,
             const std::shared_ptr<renderengine::ExternalTexture>& hdrBuffer = nullptr,
             const std::shared_ptr<renderengine::ExternalTexture>& gainmapBuffer = nullptr);
 
     ftl::SharedFuture<FenceResult> renderScreenImpl(
-            const ScreenshotArgs& args, const std::shared_ptr<renderengine::ExternalTexture>&,
+            ScreenshotArgs& args, const std::shared_ptr<renderengine::ExternalTexture>&,
             bool regionSampling, bool grayscale, bool isProtected, ScreenCaptureResults&,
-            const std::optional<OutputCompositionState>& displayState,
             const std::vector<std::pair<Layer*, sp<LayerFE>>>& layers);
 
     void readPersistentProperties();
diff --git a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
index 9a2e254..41f2b6e 100644
--- a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
+++ b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
@@ -469,7 +469,7 @@
         ftl::FakeGuard guard(kMainThreadContext);
 
         ScreenCaptureResults captureResults;
-        auto displayState = std::optional{display->getCompositionDisplay()->getState()};
+        const auto& state = display->getCompositionDisplay()->getState();
         auto layers = getLayerSnapshotsFn();
 
         SurfaceFlinger::ScreenshotArgs screenshotArgs;
@@ -480,10 +480,14 @@
         screenshotArgs.dataspace = dataspace;
         screenshotArgs.isSecure = isSecure;
         screenshotArgs.seamlessTransition = seamlessTransition;
+        screenshotArgs.displayBrightnessNits = state.displayBrightnessNits;
+        screenshotArgs.sdrWhitePointNits = state.sdrWhitePointNits;
+        screenshotArgs.renderIntent = state.renderIntent;
+        screenshotArgs.colorMode = state.colorMode;
 
         return mFlinger->renderScreenImpl(screenshotArgs, buffer, regionSampling,
                                           false /* grayscale */, false /* isProtected */,
-                                          captureResults, displayState, layers);
+                                          captureResults, layers);
     }
 
     auto getLayerSnapshotsForScreenshotsFn(ui::LayerStack layerStack, uint32_t uid) {