diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index 548aa32..eb3ae26 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -5492,37 +5492,24 @@
             }
             case 1035: {
                 const int modeId = data.readInt32();
-                mDebugDisplayModeSetByBackdoor = false;
 
-                const auto displayId = [&]() -> std::optional<PhysicalDisplayId> {
-                    uint64_t inputDisplayId = 0;
-                    if (data.readUint64(&inputDisplayId) == NO_ERROR) {
-                        const auto token = getPhysicalDisplayToken(
-                                static_cast<PhysicalDisplayId>(inputDisplayId));
-                        if (!token) {
-                            ALOGE("No display with id: %" PRIu64, inputDisplayId);
-                            return std::nullopt;
-                        }
-
-                        return std::make_optional<PhysicalDisplayId>(inputDisplayId);
+                const auto display = [&]() -> sp<IBinder> {
+                    uint64_t value;
+                    if (data.readUint64(&value) != NO_ERROR) {
+                        return getInternalDisplayToken();
                     }
 
-                    return getInternalDisplayId();
+                    if (const auto id = DisplayId::fromValue<PhysicalDisplayId>(value)) {
+                        return getPhysicalDisplayToken(*id);
+                    }
+
+                    ALOGE("Invalid physical display ID");
+                    return nullptr;
                 }();
 
-                if (!displayId) {
-                    ALOGE("No display found");
-                    return NO_ERROR;
-                }
-
-                status_t result = setActiveMode(getPhysicalDisplayToken(*displayId), modeId);
-                if (result != NO_ERROR) {
-                    return result;
-                }
-
-                mDebugDisplayModeSetByBackdoor = true;
-
-                return NO_ERROR;
+                const status_t result = setActiveMode(display, modeId);
+                mDebugDisplayModeSetByBackdoor = result == NO_ERROR;
+                return result;
             }
             case 1036: {
                 if (data.readInt32() > 0) {
@@ -5759,7 +5746,7 @@
                                captureListener);
 }
 
-status_t SurfaceFlinger::captureDisplay(uint64_t displayIdOrLayerStack,
+status_t SurfaceFlinger::captureDisplay(DisplayId displayId,
                                         const sp<IScreenCaptureListener>& captureListener) {
     ui::LayerStack layerStack;
     wp<const DisplayDevice> displayWeak;
@@ -5767,21 +5754,14 @@
     ui::Dataspace dataspace;
     {
         Mutex::Autolock lock(mStateLock);
-        auto display = getDisplayDeviceLocked(PhysicalDisplayId{displayIdOrLayerStack});
 
-        // Fall back to first display whose layer stack matches.
-        if (!display) {
-            const auto layerStack = static_cast<ui::LayerStack>(displayIdOrLayerStack);
-            display = findDisplay(WithLayerStack(layerStack));
-        }
-
+        const auto display = getDisplayDeviceLocked(displayId);
         if (!display) {
             return NAME_NOT_FOUND;
         }
 
-        layerStack = display->getLayerStack();
         displayWeak = display;
-
+        layerStack = display->getLayerStack();
         size = display->getLayerStackSpaceRect().getSize();
 
         dataspace =
@@ -5863,7 +5843,11 @@
             }
         }
 
-        const auto display = findDisplay(WithLayerStack(parent->getLayerStack()));
+        const auto display =
+                findDisplay([layerStack = parent->getLayerStack()](const auto& display) {
+                    return display.getLayerStack() == layerStack;
+                });
+
         if (!display) {
             return NAME_NOT_FOUND;
         }