Move exclude layer to CaptureArgs.

Test: atest ScreenCaptureTest ScreenCaptureChildOnlyTest
Bug: 267324693
Change-Id: I796dbb3dbc83da37a54ef36a7d05cf0ba7e0c431

4 files changed, 108 insertions, 42 deletions

diff --git a/core/java/android/window/ScreenCapture.java b/core/java/android/window/ScreenCapture.java
index 8a7efb93d961..d8e64d48f456 100644
--- a/core/java/android/window/ScreenCapture.java
+++ b/core/java/android/window/ScreenCapture.java
@@ -272,6 +272,8 @@ public class ScreenCapture {
         public final long mUid;
         public final boolean mGrayscale;
 
+        final SurfaceControl[] mExcludeLayers;
+
         private CaptureArgs(CaptureArgs.Builder<? extends CaptureArgs.Builder<?>> builder) {
             mPixelFormat = builder.mPixelFormat;
             mSourceCrop.set(builder.mSourceCrop);
@@ -281,6 +283,7 @@ public class ScreenCapture {
             mAllowProtected = builder.mAllowProtected;
             mUid = builder.mUid;
             mGrayscale = builder.mGrayscale;
+            mExcludeLayers = builder.mExcludeLayers;
         }
 
         private CaptureArgs(Parcel in) {
@@ -292,6 +295,46 @@ public class ScreenCapture {
             mAllowProtected = in.readBoolean();
             mUid = in.readLong();
             mGrayscale = in.readBoolean();
+
+            int excludeLayersLength = in.readInt();
+            if (excludeLayersLength > 0) {
+                mExcludeLayers = new SurfaceControl[excludeLayersLength];
+                for (int index = 0; index < excludeLayersLength; index++) {
+                    mExcludeLayers[index] = SurfaceControl.CREATOR.createFromParcel(in);
+                }
+            } else {
+                mExcludeLayers = null;
+            }
+        }
+
+        /** Release any layers if set using {@link Builder#setExcludeLayers(SurfaceControl[])}. */
+        public void release() {
+            if (mExcludeLayers.length == 0) {
+                return;
+            }
+
+            for (SurfaceControl surfaceControl : mExcludeLayers) {
+                if (surfaceControl != null) {
+                    surfaceControl.release();
+                }
+            }
+        }
+
+        /**
+         * Returns an array of {@link SurfaceControl#mNativeObject} corresponding to
+         * {@link #mExcludeLayers}. Used only in native code.
+         */
+        private long[] getNativeExcludeLayers() {
+            if (mExcludeLayers == null || mExcludeLayers.length == 0) {
+                return new long[0];
+            }
+
+            long[] nativeExcludeLayers = new long[mExcludeLayers.length];
+            for (int index = 0; index < mExcludeLayers.length; index++) {
+                nativeExcludeLayers[index] = mExcludeLayers[index].mNativeObject;
+            }
+
+            return nativeExcludeLayers;
         }
 
         /**
@@ -308,6 +351,7 @@ public class ScreenCapture {
             private boolean mAllowProtected;
             private long mUid = -1;
             private boolean mGrayscale;
+            private SurfaceControl[] mExcludeLayers;
 
             /**
              * Construct a new {@link CaptureArgs} with the set parameters. The builder remains
@@ -397,6 +441,14 @@ public class ScreenCapture {
             }
 
             /**
+             * An array of {@link SurfaceControl} layer handles to exclude.
+             */
+            public T setExcludeLayers(@Nullable SurfaceControl[] excludeLayers) {
+                mExcludeLayers = excludeLayers;
+                return getThis();
+            }
+
+            /**
              * Each sub class should return itself to allow the builder to chain properly
              */
             T getThis() {
@@ -419,6 +471,15 @@ public class ScreenCapture {
             dest.writeBoolean(mAllowProtected);
             dest.writeLong(mUid);
             dest.writeBoolean(mGrayscale);
+
+            if (mExcludeLayers != null) {
+                dest.writeInt(mExcludeLayers.length);
+                for (SurfaceControl excludeLayer : mExcludeLayers) {
+                    excludeLayer.writeToParcel(dest, flags);
+                }
+            } else {
+                dest.writeInt(0);
+            }
         }
 
         public static final Parcelable.Creator<CaptureArgs> CREATOR =
@@ -529,21 +590,12 @@ public class ScreenCapture {
      */
     public static class LayerCaptureArgs extends CaptureArgs {
         private final long mNativeLayer;
-        private final long[] mNativeExcludeLayers;
         private final boolean mChildrenOnly;
 
         private LayerCaptureArgs(Builder builder) {
             super(builder);
             mChildrenOnly = builder.mChildrenOnly;
             mNativeLayer = builder.mLayer.mNativeObject;
-            if (builder.mExcludeLayers != null) {
-                mNativeExcludeLayers = new long[builder.mExcludeLayers.length];
-                for (int i = 0; i < builder.mExcludeLayers.length; i++) {
-                    mNativeExcludeLayers[i] = builder.mExcludeLayers[i].mNativeObject;
-                }
-            } else {
-                mNativeExcludeLayers = null;
-            }
         }
 
         /**
@@ -551,7 +603,6 @@ public class ScreenCapture {
          */
         public static class Builder extends CaptureArgs.Builder<Builder> {
             private SurfaceControl mLayer;
-            private SurfaceControl[] mExcludeLayers;
             private boolean mChildrenOnly = true;
 
             /**
@@ -575,6 +626,7 @@ public class ScreenCapture {
                 setAllowProtected(args.mAllowProtected);
                 setUid(args.mUid);
                 setGrayscale(args.mGrayscale);
+                setExcludeLayers(args.mExcludeLayers);
             }
 
             public Builder(SurfaceControl layer) {
@@ -590,14 +642,6 @@ public class ScreenCapture {
             }
 
             /**
-             * An array of layer handles to exclude.
-             */
-            public Builder setExcludeLayers(@Nullable SurfaceControl[] excludeLayers) {
-                mExcludeLayers = excludeLayers;
-                return this;
-            }
-
-            /**
              * Whether to include the layer itself in the screenshot or just the children and their
              * descendants.
              */
diff --git a/core/jni/android_window_ScreenCapture.cpp b/core/jni/android_window_ScreenCapture.cpp
index c1929c6535fb..1b67a0da57e1 100644
--- a/core/jni/android_window_ScreenCapture.cpp
+++ b/core/jni/android_window_ScreenCapture.cpp
@@ -45,6 +45,7 @@ static struct {
     jfieldID allowProtected;
     jfieldID uid;
     jfieldID grayscale;
+    jmethodID getNativeExcludeLayers;
 } gCaptureArgsClassInfo;
 
 static struct {
@@ -56,7 +57,6 @@ static struct {
 
 static struct {
     jfieldID layer;
-    jfieldID excludeLayers;
     jfieldID childrenOnly;
 } gLayerCaptureArgsClassInfo;
 
@@ -168,6 +168,23 @@ static void getCaptureArgs(JNIEnv* env, jobject captureArgsObject, CaptureArgs&
     captureArgs.uid = env->GetLongField(captureArgsObject, gCaptureArgsClassInfo.uid);
     captureArgs.grayscale =
             env->GetBooleanField(captureArgsObject, gCaptureArgsClassInfo.grayscale);
+
+    jlongArray excludeObjectArray = reinterpret_cast<jlongArray>(
+            env->CallObjectMethod(captureArgsObject, gCaptureArgsClassInfo.getNativeExcludeLayers));
+    if (excludeObjectArray != nullptr) {
+        ScopedLongArrayRO excludeArray(env, excludeObjectArray);
+        const jsize len = excludeArray.size();
+        captureArgs.excludeHandles.reserve(len);
+
+        for (jsize i = 0; i < len; i++) {
+            auto excludeObject = reinterpret_cast<SurfaceControl*>(excludeArray[i]);
+            if (excludeObject == nullptr) {
+                jniThrowNullPointerException(env, "Exclude layer is null");
+                return;
+            }
+            captureArgs.excludeHandles.emplace(excludeObject->getHandle());
+        }
+    }
 }
 
 static DisplayCaptureArgs displayCaptureArgsFromObject(JNIEnv* env,
@@ -207,6 +224,7 @@ static jint nativeCaptureLayers(JNIEnv* env, jclass clazz, jobject layerCaptureA
                                 jlong screenCaptureListenerObject) {
     LayerCaptureArgs captureArgs;
     getCaptureArgs(env, layerCaptureArgsObject, captureArgs);
+
     SurfaceControl* layer = reinterpret_cast<SurfaceControl*>(
             env->GetLongField(layerCaptureArgsObject, gLayerCaptureArgsClassInfo.layer));
     if (layer == nullptr) {
@@ -217,23 +235,6 @@ static jint nativeCaptureLayers(JNIEnv* env, jclass clazz, jobject layerCaptureA
     captureArgs.childrenOnly =
             env->GetBooleanField(layerCaptureArgsObject, gLayerCaptureArgsClassInfo.childrenOnly);
 
-    jlongArray excludeObjectArray = reinterpret_cast<jlongArray>(
-            env->GetObjectField(layerCaptureArgsObject, gLayerCaptureArgsClassInfo.excludeLayers));
-    if (excludeObjectArray != nullptr) {
-        ScopedLongArrayRO excludeArray(env, excludeObjectArray);
-        const jsize len = excludeArray.size();
-        captureArgs.excludeHandles.reserve(len);
-
-        for (jsize i = 0; i < len; i++) {
-            auto excludeObject = reinterpret_cast<SurfaceControl*>(excludeArray[i]);
-            if (excludeObject == nullptr) {
-                jniThrowNullPointerException(env, "Exclude layer is null");
-                return BAD_VALUE;
-            }
-            captureArgs.excludeHandles.emplace(excludeObject->getHandle());
-        }
-    }
-
     sp<gui::IScreenCaptureListener> captureListener =
             reinterpret_cast<gui::IScreenCaptureListener*>(screenCaptureListenerObject);
     return ScreenshotClient::captureLayers(captureArgs, captureListener);
@@ -318,6 +319,9 @@ int register_android_window_ScreenCapture(JNIEnv* env) {
     gCaptureArgsClassInfo.uid = GetFieldIDOrDie(env, captureArgsClazz, "mUid", "J");
     gCaptureArgsClassInfo.grayscale = GetFieldIDOrDie(env, captureArgsClazz, "mGrayscale", "Z");
 
+    gCaptureArgsClassInfo.getNativeExcludeLayers =
+            GetMethodIDOrDie(env, captureArgsClazz, "getNativeExcludeLayers", "()[J");
+
     jclass displayCaptureArgsClazz =
             FindClassOrDie(env, "android/window/ScreenCapture$DisplayCaptureArgs");
     gDisplayCaptureArgsClassInfo.displayToken =
@@ -333,8 +337,6 @@ int register_android_window_ScreenCapture(JNIEnv* env) {
             FindClassOrDie(env, "android/window/ScreenCapture$LayerCaptureArgs");
     gLayerCaptureArgsClassInfo.layer =
             GetFieldIDOrDie(env, layerCaptureArgsClazz, "mNativeLayer", "J");
-    gLayerCaptureArgsClassInfo.excludeLayers =
-            GetFieldIDOrDie(env, layerCaptureArgsClazz, "mNativeExcludeLayers", "[J");
     gLayerCaptureArgsClassInfo.childrenOnly =
             GetFieldIDOrDie(env, layerCaptureArgsClazz, "mChildrenOnly", "Z");
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
index 3dbb745f0c6c..f74fda668f90 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
@@ -74,8 +74,10 @@ import android.util.Log;
 import android.util.Pair;
 import android.util.SparseArray;
 import android.view.IWindowManager;
+import android.view.SurfaceControl;
 import android.view.View;
 import android.view.ViewGroup;
+import android.view.ViewRootImpl;
 import android.view.WindowInsets;
 import android.view.WindowManager;
 import android.window.ScreenCapture;
@@ -1133,13 +1135,24 @@ public class BubbleController implements ConfigurationChangeListener,
     /**
      * Performs a screenshot that may exclude the bubble layer, if one is present. The screenshot
      * can be access via the supplied {@link ScreenshotSync#get()} asynchronously.
-     *
-     * TODO(b/267324693): Implement the exclude layer functionality in screenshot.
      */
     public void getScreenshotExcludingBubble(int displayId,
             Pair<ScreenCaptureListener, ScreenshotSync> screenCaptureListener) {
         try {
-            mWmService.captureDisplay(displayId, null, screenCaptureListener.first);
+            ScreenCapture.CaptureArgs args = null;
+            if (mStackView != null) {
+                ViewRootImpl viewRoot = mStackView.getViewRootImpl();
+                if (viewRoot != null) {
+                    SurfaceControl bubbleLayer = viewRoot.getSurfaceControl();
+                    if (bubbleLayer != null) {
+                        args = new ScreenCapture.CaptureArgs.Builder<>()
+                                .setExcludeLayers(new SurfaceControl[] {bubbleLayer})
+                                .build();
+                    }
+                }
+            }
+
+            mWmService.captureDisplay(displayId, args, screenCaptureListener.first);
         } catch (RemoteException e) {
             Log.e(TAG, "Failed to capture screenshot");
         }
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index a7a90604f228..6fc81ac9c2c6 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -9314,7 +9314,14 @@ public class WindowManagerService extends IWindowManager.Stub
             throw new SecurityException("Requires READ_FRAME_BUFFER permission");
         }
 
-        ScreenCapture.captureLayers(getCaptureArgs(displayId, captureArgs), listener);
+        ScreenCapture.LayerCaptureArgs layerCaptureArgs = getCaptureArgs(displayId, captureArgs);
+        ScreenCapture.captureLayers(layerCaptureArgs, listener);
+
+        if (Binder.getCallingUid() != SYSTEM_UID) {
+            // Release the SurfaceControl objects only if the caller is not in system server as no
+            // parcelling occurs in this case.
+            layerCaptureArgs.release();
+        }
     }
 
     @VisibleForTesting