Merge "Revert "Revert "Allow screenshot to carry more information."""

6 files changed, 138 insertions, 41 deletions

diff --git a/core/java/android/view/SurfaceControl.java b/core/java/android/view/SurfaceControl.java
index 0b91d2126254..0043d3258772 100644
--- a/core/java/android/view/SurfaceControl.java
+++ b/core/java/android/view/SurfaceControl.java
@@ -88,10 +88,10 @@ public final class SurfaceControl implements Parcelable {
     private static native void nativeDestroy(long nativeObject);
     private static native void nativeDisconnect(long nativeObject);
 
-    private static native GraphicBuffer nativeScreenshot(IBinder displayToken,
+    private static native ScreenshotGraphicBuffer nativeScreenshot(IBinder displayToken,
             Rect sourceCrop, int width, int height, boolean useIdentityTransform, int rotation,
             boolean captureSecureLayers);
-    private static native GraphicBuffer nativeCaptureLayers(IBinder layerHandleToken,
+    private static native ScreenshotGraphicBuffer nativeCaptureLayers(IBinder layerHandleToken,
             Rect sourceCrop, float frameScale);
 
     private static native long nativeCreateTransaction();
@@ -431,6 +431,46 @@ public final class SurfaceControl implements Parcelable {
     public static final int METADATA_TASK_ID = 3;
 
     /**
+     * A wrapper around GraphicBuffer that contains extra information about how to
+     * interpret the screenshot GraphicBuffer.
+     * @hide
+     */
+    public static class ScreenshotGraphicBuffer {
+        private final GraphicBuffer mGraphicBuffer;
+        private final ColorSpace mColorSpace;
+
+        public ScreenshotGraphicBuffer(GraphicBuffer graphicBuffer, ColorSpace colorSpace) {
+            mGraphicBuffer = graphicBuffer;
+            mColorSpace = colorSpace;
+        }
+
+       /**
+        * Create ScreenshotGraphicBuffer from existing native GraphicBuffer object.
+        * @param width The width in pixels of the buffer
+        * @param height The height in pixels of the buffer
+        * @param format The format of each pixel as specified in {@link PixelFormat}
+        * @param usage Hint indicating how the buffer will be used
+        * @param unwrappedNativeObject The native object of GraphicBuffer
+        * @param namedColorSpace Integer value of a named color space {@link ColorSpace.Named}
+        */
+        private static ScreenshotGraphicBuffer createFromNative(int width, int height, int format,
+                int usage, long unwrappedNativeObject, int namedColorSpace) {
+            GraphicBuffer graphicBuffer = GraphicBuffer.createFromExisting(width, height, format,
+                    usage, unwrappedNativeObject);
+            ColorSpace colorSpace = ColorSpace.get(ColorSpace.Named.values()[namedColorSpace]);
+            return new ScreenshotGraphicBuffer(graphicBuffer, colorSpace);
+        }
+
+        public ColorSpace getColorSpace() {
+            return mColorSpace;
+        }
+
+        public GraphicBuffer getGraphicBuffer() {
+            return mGraphicBuffer;
+        }
+    }
+
+    /**
      * Builder class for {@link SurfaceControl} objects.
      *
      * By default the surface will be hidden, and have "unset" bounds, meaning it can
@@ -1823,10 +1863,10 @@ public final class SurfaceControl implements Parcelable {
             throw new IllegalArgumentException("consumer must not be null");
         }
 
-        final GraphicBuffer buffer = screenshotToBuffer(display, sourceCrop, width, height,
-                useIdentityTransform, rotation);
+        final ScreenshotGraphicBuffer buffer = screenshotToBuffer(display, sourceCrop, width,
+                height, useIdentityTransform, rotation);
         try {
-            consumer.attachAndQueueBuffer(buffer);
+            consumer.attachAndQueueBuffer(buffer.getGraphicBuffer());
         } catch (RuntimeException e) {
             Log.w(TAG, "Failed to take screenshot - " + e.getMessage());
         }
@@ -1869,17 +1909,16 @@ public final class SurfaceControl implements Parcelable {
         }
 
         SurfaceControl.rotateCropForSF(sourceCrop, rotation);
-        final GraphicBuffer buffer = screenshotToBuffer(displayToken, sourceCrop, width, height,
-                useIdentityTransform, rotation);
+        final ScreenshotGraphicBuffer buffer = screenshotToBuffer(displayToken, sourceCrop, width,
+                height, useIdentityTransform, rotation);
 
         if (buffer == null) {
             Log.w(TAG, "Failed to take screenshot");
             return null;
         }
-        // TODO(b/116112787) Now that hardware bitmap creation can take color space, we
-        // should continue to fix screenshot.
-        return Bitmap.wrapHardwareBuffer(HardwareBuffer.createFromGraphicBuffer(buffer),
-                ColorSpace.get(ColorSpace.Named.SRGB));
+        return Bitmap.wrapHardwareBuffer(
+                HardwareBuffer.createFromGraphicBuffer(buffer.getGraphicBuffer()),
+                buffer.getColorSpace());
     }
 
     /**
@@ -1905,8 +1944,8 @@ public final class SurfaceControl implements Parcelable {
      * @return Returns a GraphicBuffer that contains the captured content.
      * @hide
      */
-    public static GraphicBuffer screenshotToBuffer(IBinder display, Rect sourceCrop, int width,
-            int height, boolean useIdentityTransform, int rotation) {
+    public static ScreenshotGraphicBuffer screenshotToBuffer(IBinder display, Rect sourceCrop,
+            int width, int height, boolean useIdentityTransform, int rotation) {
         if (display == null) {
             throw new IllegalArgumentException("displayToken must not be null");
         }
@@ -1925,7 +1964,7 @@ public final class SurfaceControl implements Parcelable {
      *
      * @hide
      */
-    public static GraphicBuffer screenshotToBufferWithSecureLayersUnsafe(IBinder display,
+    public static ScreenshotGraphicBuffer screenshotToBufferWithSecureLayersUnsafe(IBinder display,
             Rect sourceCrop, int width, int height, boolean useIdentityTransform,
             int rotation) {
         if (display == null) {
@@ -1959,7 +1998,7 @@ public final class SurfaceControl implements Parcelable {
      * @return Returns a GraphicBuffer that contains the layer capture.
      * @hide
      */
-    public static GraphicBuffer captureLayers(IBinder layerHandleToken, Rect sourceCrop,
+    public static ScreenshotGraphicBuffer captureLayers(IBinder layerHandleToken, Rect sourceCrop,
             float frameScale) {
         return nativeCaptureLayers(layerHandleToken, sourceCrop, frameScale);
     }
diff --git a/core/jni/android_view_SurfaceControl.cpp b/core/jni/android_view_SurfaceControl.cpp
index 4a6c72bda283..94f96babaf90 100644
--- a/core/jni/android_view_SurfaceControl.cpp
+++ b/core/jni/android_view_SurfaceControl.cpp
@@ -126,6 +126,41 @@ static struct {
     jfieldID white;
 } gDisplayPrimariesClassInfo;
 
+static struct {
+    jclass clazz;
+    jmethodID builder;
+} gScreenshotGraphicBufferClassInfo;
+
+class JNamedColorSpace {
+public:
+    // ColorSpace.Named.SRGB.ordinal() = 0;
+    static constexpr jint SRGB = 0;
+
+    // ColorSpace.Named.DISPLAY_P3.ordinal() = 6;
+    static constexpr jint DISPLAY_P3 = 6;
+};
+
+constexpr jint fromDataspaceToNamedColorSpaceValue(const ui::Dataspace dataspace) {
+    switch (dataspace) {
+        case ui::Dataspace::DISPLAY_P3:
+            return JNamedColorSpace::DISPLAY_P3;
+        default:
+            return JNamedColorSpace::SRGB;
+    }
+}
+
+constexpr ui::Dataspace pickDataspaceFromColorMode(const ui::ColorMode colorMode) {
+    switch (colorMode) {
+        case ui::ColorMode::DISPLAY_P3:
+        case ui::ColorMode::BT2100_PQ:
+        case ui::ColorMode::BT2100_HLG:
+        case ui::ColorMode::DISPLAY_BT2020:
+            return ui::Dataspace::DISPLAY_P3;
+        default:
+            return ui::Dataspace::V0_SRGB;
+    }
+}
+
 // ----------------------------------------------------------------------------
 
 static jlong nativeCreateTransaction(JNIEnv* env, jclass clazz) {
@@ -210,9 +245,12 @@ static jobject nativeScreenshot(JNIEnv* env, jclass clazz,
     if (displayToken == NULL) {
         return NULL;
     }
+    const ui::ColorMode colorMode = SurfaceComposerClient::getActiveColorMode(displayToken);
+    const ui::Dataspace dataspace = pickDataspaceFromColorMode(colorMode);
+
     Rect sourceCrop = rectFromObj(env, sourceCropObj);
     sp<GraphicBuffer> buffer;
-    status_t res = ScreenshotClient::capture(displayToken, ui::Dataspace::V0_SRGB,
+    status_t res = ScreenshotClient::capture(displayToken, dataspace,
             ui::PixelFormat::RGBA_8888,
             sourceCrop, width, height,
             useIdentityTransform, rotation, captureSecureLayers, &buffer);
@@ -220,13 +258,15 @@ static jobject nativeScreenshot(JNIEnv* env, jclass clazz,
         return NULL;
     }
 
-    return env->CallStaticObjectMethod(gGraphicBufferClassInfo.clazz,
-            gGraphicBufferClassInfo.builder,
+    const jint namedColorSpace = fromDataspaceToNamedColorSpaceValue(dataspace);
+    return env->CallStaticObjectMethod(gScreenshotGraphicBufferClassInfo.clazz,
+            gScreenshotGraphicBufferClassInfo.builder,
             buffer->getWidth(),
             buffer->getHeight(),
             buffer->getPixelFormat(),
             (jint)buffer->getUsage(),
-            (jlong)buffer.get());
+            (jlong)buffer.get(),
+            namedColorSpace);
 }
 
 static jobject nativeCaptureLayers(JNIEnv* env, jclass clazz, jobject layerHandleToken,
@@ -243,20 +283,23 @@ static jobject nativeCaptureLayers(JNIEnv* env, jclass clazz, jobject layerHandl
     }
 
     sp<GraphicBuffer> buffer;
-    status_t res = ScreenshotClient::captureChildLayers(layerHandle, ui::Dataspace::V0_SRGB,
+    const ui::Dataspace dataspace = ui::Dataspace::V0_SRGB;
+    status_t res = ScreenshotClient::captureChildLayers(layerHandle, dataspace,
                                                         ui::PixelFormat::RGBA_8888, sourceCrop,
                                                         frameScale, &buffer);
     if (res != NO_ERROR) {
         return NULL;
     }
 
-    return env->CallStaticObjectMethod(gGraphicBufferClassInfo.clazz,
-                                       gGraphicBufferClassInfo.builder,
+    const jint namedColorSpace = fromDataspaceToNamedColorSpaceValue(dataspace);
+    return env->CallStaticObjectMethod(gScreenshotGraphicBufferClassInfo.clazz,
+                                       gScreenshotGraphicBufferClassInfo.builder,
                                        buffer->getWidth(),
                                        buffer->getHeight(),
                                        buffer->getPixelFormat(),
                                        (jint)buffer->getUsage(),
-                                       (jlong)buffer.get());
+                                       (jlong)buffer.get(),
+                                       namedColorSpace);
 }
 
 static void nativeApplyTransaction(JNIEnv* env, jclass clazz, jlong transactionObj, jboolean sync) {
@@ -1306,9 +1349,13 @@ static const JNINativeMethod sSurfaceControlMethods[] = {
             (void*)nativeSetOverrideScalingMode },
     {"nativeGetHandle", "(J)Landroid/os/IBinder;",
             (void*)nativeGetHandle },
-    {"nativeScreenshot", "(Landroid/os/IBinder;Landroid/graphics/Rect;IIZIZ)Landroid/graphics/GraphicBuffer;",
+    {"nativeScreenshot",
+            "(Landroid/os/IBinder;Landroid/graphics/Rect;IIZIZ)"
+            "Landroid/view/SurfaceControl$ScreenshotGraphicBuffer;",
             (void*)nativeScreenshot },
-    {"nativeCaptureLayers", "(Landroid/os/IBinder;Landroid/graphics/Rect;F)Landroid/graphics/GraphicBuffer;",
+    {"nativeCaptureLayers",
+            "(Landroid/os/IBinder;Landroid/graphics/Rect;F)"
+            "Landroid/view/SurfaceControl$ScreenshotGraphicBuffer;",
             (void*)nativeCaptureLayers },
     {"nativeSetInputWindowInfo", "(JJLandroid/view/InputWindowHandle;)V",
             (void*)nativeSetInputWindowInfo },
@@ -1386,6 +1433,14 @@ int register_android_view_SurfaceControl(JNIEnv* env)
     gGraphicBufferClassInfo.builder = GetStaticMethodIDOrDie(env, graphicsBufferClazz,
             "createFromExisting", "(IIIIJ)Landroid/graphics/GraphicBuffer;");
 
+    jclass screenshotGraphicsBufferClazz = FindClassOrDie(env,
+            "android/view/SurfaceControl$ScreenshotGraphicBuffer");
+    gScreenshotGraphicBufferClassInfo.clazz =
+            MakeGlobalRefOrDie(env, screenshotGraphicsBufferClazz);
+    gScreenshotGraphicBufferClassInfo.builder = GetStaticMethodIDOrDie(env,
+            screenshotGraphicsBufferClazz,
+            "createFromNative", "(IIIIJI)Landroid/view/SurfaceControl$ScreenshotGraphicBuffer;");
+
     jclass displayedContentSampleClazz = FindClassOrDie(env,
             "android/hardware/display/DisplayedContentSample");
     gDisplayedContentSampleClassInfo.clazz = MakeGlobalRefOrDie(env, displayedContentSampleClazz);
diff --git a/services/core/java/com/android/server/display/DisplayManagerService.java b/services/core/java/com/android/server/display/DisplayManagerService.java
index 32f34b8378ce..3010324488b8 100644
--- a/services/core/java/com/android/server/display/DisplayManagerService.java
+++ b/services/core/java/com/android/server/display/DisplayManagerService.java
@@ -37,7 +37,6 @@ import android.content.pm.ParceledListSlice;
 import android.content.res.Resources;
 import android.content.res.TypedArray;
 import android.graphics.ColorSpace;
-import android.graphics.GraphicBuffer;
 import android.graphics.Point;
 import android.graphics.Rect;
 import android.hardware.SensorManager;
@@ -1275,11 +1274,12 @@ public final class DisplayManagerService extends SystemService {
         if (token == null) {
             return false;
         }
-        final GraphicBuffer gb = SurfaceControl.screenshotToBufferWithSecureLayersUnsafe(
-                token, new Rect(), 0 /* width */, 0 /* height */, false /* useIdentityTransform */,
-                0 /* rotation */);
+        final SurfaceControl.ScreenshotGraphicBuffer gb =
+                SurfaceControl.screenshotToBufferWithSecureLayersUnsafe(
+                        token, new Rect(), 0 /* width */, 0 /* height */,
+                        false /* useIdentityTransform */, 0 /* rotation */);
         try {
-            outSurface.attachAndQueueBuffer(gb);
+            outSurface.attachAndQueueBuffer(gb.getGraphicBuffer());
         } catch (RuntimeException e) {
             Slog.w(TAG, "Failed to take screenshot - " + e.getMessage());
         }
diff --git a/services/core/java/com/android/server/wm/TaskScreenshotAnimatable.java b/services/core/java/com/android/server/wm/TaskScreenshotAnimatable.java
index e0d85e876476..4379b7cdace6 100644
--- a/services/core/java/com/android/server/wm/TaskScreenshotAnimatable.java
+++ b/services/core/java/com/android/server/wm/TaskScreenshotAnimatable.java
@@ -41,7 +41,7 @@ class TaskScreenshotAnimatable implements SurfaceAnimator.Animatable {
         return new TaskScreenshotAnimatable(task, getBufferFromTask(task));
     }
 
-    private static GraphicBuffer getBufferFromTask(Task task) {
+    private static SurfaceControl.ScreenshotGraphicBuffer getBufferFromTask(Task task) {
         if (task == null) {
             return null;
         }
@@ -51,7 +51,10 @@ class TaskScreenshotAnimatable implements SurfaceAnimator.Animatable {
                 task.getSurfaceControl().getHandle(), tmpRect, 1f);
     }
 
-    private TaskScreenshotAnimatable(Task task, GraphicBuffer buffer) {
+    private TaskScreenshotAnimatable(Task task,
+            SurfaceControl.ScreenshotGraphicBuffer screenshotBuffer) {
+        GraphicBuffer buffer = screenshotBuffer == null
+                ? null : screenshotBuffer.getGraphicBuffer();
         mTask = task;
         mWidth = (buffer != null) ? buffer.getWidth() : 1;
         mHeight = (buffer != null) ? buffer.getHeight() : 1;
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotController.java b/services/core/java/com/android/server/wm/TaskSnapshotController.java
index beb3d82c05fb..6fe8b43db139 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotController.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotController.java
@@ -279,8 +279,11 @@ class TaskSnapshotController {
             Slog.w(TAG_WM, "Failed to take screenshot. No main window for " + task);
             return null;
         }
-        final GraphicBuffer buffer = SurfaceControl.captureLayers(
-                task.getSurfaceControl().getHandle(), mTmpRect, scaleFraction);
+        final SurfaceControl.ScreenshotGraphicBuffer screenshotBuffer =
+                SurfaceControl.captureLayers(
+                        task.getSurfaceControl().getHandle(), mTmpRect, scaleFraction);
+        final GraphicBuffer buffer = screenshotBuffer != null ? screenshotBuffer.getGraphicBuffer()
+                : null;
         if (buffer == null || buffer.getWidth() <= 1 || buffer.getHeight() <= 1) {
             if (DEBUG_SCREENSHOT) {
                 Slog.w(TAG_WM, "Failed to take screenshot for " + task);
diff --git a/services/core/java/com/android/server/wm/WallpaperController.java b/services/core/java/com/android/server/wm/WallpaperController.java
index dddc6b755db0..166a33d7a560 100644
--- a/services/core/java/com/android/server/wm/WallpaperController.java
+++ b/services/core/java/com/android/server/wm/WallpaperController.java
@@ -34,8 +34,6 @@ import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
 import static com.android.server.wm.WindowManagerService.H.WALLPAPER_DRAW_PENDING_TIMEOUT;
 
 import android.graphics.Bitmap;
-import android.graphics.ColorSpace;
-import android.graphics.GraphicBuffer;
 import android.graphics.Rect;
 import android.hardware.HardwareBuffer;
 import android.os.Bundle;
@@ -738,17 +736,16 @@ class WallpaperController {
         final Rect bounds = wallpaperWindowState.getBounds();
         bounds.offsetTo(0, 0);
 
-        GraphicBuffer wallpaperBuffer = SurfaceControl.captureLayers(
+        SurfaceControl.ScreenshotGraphicBuffer wallpaperBuffer = SurfaceControl.captureLayers(
                 wallpaperWindowState.getSurfaceControl().getHandle(), bounds, 1 /* frameScale */);
 
         if (wallpaperBuffer == null) {
             Slog.w(TAG_WM, "Failed to screenshot wallpaper");
             return null;
         }
-        // TODO(b/116112787) Now that hardware bitmap creation can take color space, we
-        // should continue to fix screenshot.
-        return Bitmap.wrapHardwareBuffer(HardwareBuffer.createFromGraphicBuffer(wallpaperBuffer),
-                                         ColorSpace.get(ColorSpace.Named.SRGB));
+        return Bitmap.wrapHardwareBuffer(
+                HardwareBuffer.createFromGraphicBuffer(wallpaperBuffer.getGraphicBuffer()),
+                wallpaperBuffer.getColorSpace());
     }
 
     private WindowState getTopVisibleWallpaper() {