DMS: Avoid SurfaceControl query for display token

This CL looks up the IBinder token for operations on physical displays
instead of calling SurfaceControl.getBuiltInDisplay, which does not
expect logical display IDs.

Bug: 116025192
Test: DisplayManagerServiceTest
Change-Id: I8c3af1f71638608f84ae86a8d5ca51b4158fea02
diff --git a/services/core/java/com/android/server/display/DisplayManagerService.java b/services/core/java/com/android/server/display/DisplayManagerService.java
index b89768a..c9df86e 100644
--- a/services/core/java/com/android/server/display/DisplayManagerService.java
+++ b/services/core/java/com/android/server/display/DisplayManagerService.java
@@ -1241,27 +1241,59 @@
         }
     }
 
+    @Nullable
+    private IBinder getDisplayToken(int displayId) {
+        synchronized (mSyncRoot) {
+            final LogicalDisplay display = mLogicalDisplays.get(displayId);
+            if (display != null) {
+                final DisplayDevice device = display.getPrimaryDisplayDeviceLocked();
+                if (device != null) {
+                    return device.getDisplayTokenLocked();
+                }
+            }
+        }
+
+        return null;
+    }
+
+    private boolean screenshotInternal(int displayId, Surface outSurface) {
+        final IBinder token = getDisplayToken(displayId);
+        if (token == null) {
+            return false;
+        }
+        SurfaceControl.screenshot(token, outSurface);
+        return true;
+    }
+
     @VisibleForTesting
     DisplayedContentSamplingAttributes getDisplayedContentSamplingAttributesInternal(
             int displayId) {
-        IBinder displayToken = SurfaceControl.getBuiltInDisplay(displayId);
-        return SurfaceControl.getDisplayedContentSamplingAttributes(displayToken);
+        final IBinder token = getDisplayToken(displayId);
+        if (token == null) {
+            return null;
+        }
+        return SurfaceControl.getDisplayedContentSamplingAttributes(token);
     }
 
     @VisibleForTesting
     boolean setDisplayedContentSamplingEnabledInternal(
             int displayId, boolean enable, int componentMask, int maxFrames) {
-        IBinder displayToken = SurfaceControl.getBuiltInDisplay(displayId);
+        final IBinder token = getDisplayToken(displayId);
+        if (token == null) {
+            return false;
+        }
         return SurfaceControl.setDisplayedContentSamplingEnabled(
-                displayToken, enable, componentMask, maxFrames);
+                token, enable, componentMask, maxFrames);
     }
 
     @VisibleForTesting
     DisplayedContentSample getDisplayedContentSampleInternal(int displayId,
             long maxFrames, long timestamp) {
-        IBinder displayToken = SurfaceControl.getBuiltInDisplay(displayId);
-        return SurfaceControl.getDisplayedContentSample(
-            displayToken, maxFrames, timestamp);
+        final IBinder token = getDisplayToken(displayId);
+        if (token == null) {
+            return null;
+        }
+        return SurfaceControl.getDisplayedContentSample(token, maxFrames, timestamp);
     }
 
     private void clearViewportsLocked() {
@@ -2257,20 +2289,7 @@
 
         @Override
         public boolean screenshot(int displayId, Surface outSurface) {
-            synchronized (mSyncRoot) {
-                final LogicalDisplay display = mLogicalDisplays.get(displayId);
-                if (display != null) {
-                    final DisplayDevice device = display.getPrimaryDisplayDeviceLocked();
-                    if (device != null) {
-                        final IBinder token = device.getDisplayTokenLocked();
-                        if (token != null) {
-                            SurfaceControl.screenshot(token, outSurface);
-                            return true;
-                        }
-                    }
-                }
-            }
-            return false;
+            return screenshotInternal(displayId, outSurface);
         }
 
         @Override