Merge "Initial work profile screenshot policy handling" into tm-qpr-dev am: 8dfa3f09b8

Original change: https://googleplex-android-review.googlesource.com/c/platform/frameworks/base/+/19432252

Change-Id: I803ae48d7614a2462d2e9cd9477759c2eb405338
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

packages/SystemUI/AndroidManifest.xml

@@ -394,6 +394,11 @@
                   android:label="@string/screenshot_scroll_label"
                   android:finishOnTaskLaunch="true" />
 
+        <service android:name=".screenshot.ScreenshotProxyService"
+                 android:permission="com.android.systemui.permission.SELF"
+                 android:exported="false" />
+
+
         <service android:name=".screenrecord.RecordingService" />
 
         <receiver android:name=".SysuiRestartReceiver"

packages/SystemUI/src/com/android/systemui/flags/Flags.java

@@ -234,6 +234,7 @@ public class Flags {
     // 1300 - screenshots
 
     public static final UnreleasedFlag SCREENSHOT_REQUEST_PROCESSOR = new UnreleasedFlag(1300);
+    public static final UnreleasedFlag SCREENSHOT_WORK_PROFILE_POLICY = new UnreleasedFlag(1301);
 
     // Pay no attention to the reflection behind the curtain.
     // ========================== Curtain ==========================

packages/SystemUI/src/com/android/systemui/screenshot/IScreenshotProxy.aidl

@@ -0,0 +1,24 @@
+/**
+ * Copyright (c) 2009, The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.screenshot;
+
+/** Interface implemented by ScreenshotProxyService */
+interface IScreenshotProxy {
+
+    /** Is the notification shade currently exanded? */
+    boolean isNotificationShadeExpanded();
+}

packages/SystemUI/src/com/android/systemui/screenshot/ImageCapture.kt

@@ -22,5 +22,5 @@ interface ImageCapture {
 
     fun captureDisplay(displayId: Int, crop: Rect? = null): Bitmap?
 
-    fun captureTask(taskId: Int): Bitmap?
+    suspend fun captureTask(taskId: Int): Bitmap?
 }

packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt

@@ -27,13 +27,19 @@ import android.view.SurfaceControl
 import android.view.SurfaceControl.DisplayCaptureArgs
 import android.view.SurfaceControl.ScreenshotHardwareBuffer
 import androidx.annotation.VisibleForTesting
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Background
 import javax.inject.Inject
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.withContext
 
 private const val TAG = "ImageCaptureImpl"
 
+@SysUISingleton
 open class ImageCaptureImpl @Inject constructor(
     private val displayManager: DisplayManager,
-    private val atmService: IActivityTaskManager
+    private val atmService: IActivityTaskManager,
+    @Background private val bgContext: CoroutineDispatcher
 ) : ImageCapture {
 
     override fun captureDisplay(displayId: Int, crop: Rect?): Bitmap? {
@@ -46,8 +52,8 @@ open class ImageCaptureImpl @Inject constructor(
         return buffer?.asBitmap()
     }
 
-    override fun captureTask(taskId: Int): Bitmap? {
+    override suspend fun captureTask(taskId: Int): Bitmap? {
-        val snapshot = atmService.takeTaskSnapshot(taskId)
+        val snapshot = withContext(bgContext) { atmService.takeTaskSnapshot(taskId) } ?: return null
         return Bitmap.wrapHardwareBuffer(snapshot.hardwareBuffer, snapshot.colorSpace)
     }
 
@@ -67,12 +73,17 @@ open class ImageCaptureImpl @Inject constructor(
     }
 
     @VisibleForTesting
-    open fun captureDisplay(displayToken: IBinder, width: Int, height: Int, crop: Rect): ScreenshotHardwareBuffer? {
+    open fun captureDisplay(
-        val captureArgs = DisplayCaptureArgs.Builder(displayToken)
+        displayToken: IBinder,
+        width: Int,
+        height: Int,
+        crop: Rect
+    ): ScreenshotHardwareBuffer? {
+        val captureArgs =
+            DisplayCaptureArgs.Builder(displayToken)
                 .setSize(width, height)
                 .setSourceCrop(crop)
                 .build()
         return SurfaceControl.captureDisplay(captureArgs)
     }
-
 }

packages/SystemUI/src/com/android/systemui/screenshot/RequestProcessor.kt

@@ -16,51 +16,84 @@
 
 package com.android.systemui.screenshot
 
-import android.net.Uri
+import android.graphics.Insets
-import android.util.Log
-import android.view.WindowManager.TAKE_SCREENSHOT_FULLSCREEN
 import android.view.WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE
-import android.view.WindowManager.TAKE_SCREENSHOT_SELECTED_REGION
 import com.android.internal.util.ScreenshotHelper.HardwareBitmapBundler
 import com.android.internal.util.ScreenshotHelper.ScreenshotRequest
 import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.screenshot.TakeScreenshotService.RequestCallback
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.flags.FeatureFlags
+import com.android.systemui.flags.Flags.SCREENSHOT_WORK_PROFILE_POLICY
 import java.util.function.Consumer
 import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.launch
 
 /**
  * Processes a screenshot request sent from {@link ScreenshotHelper}.
  */
 @SysUISingleton
-internal class RequestProcessor @Inject constructor(
+class RequestProcessor @Inject constructor(
-    private val controller: ScreenshotController,
+    private val capture: ImageCapture,
+    private val policy: ScreenshotPolicy,
+    private val flags: FeatureFlags,
+    /** For the Java Async version, to invoke the callback. */
+    @Application private val mainScope: CoroutineScope
 ) {
-    fun processRequest(
+    /**
-        request: ScreenshotRequest,
+     * Inspects the incoming request, returning a potentially modified request depending on policy.
-        onSavedListener: Consumer<Uri>,
+     *
-        callback: RequestCallback
+     * @param request the request to process
+     */
+    suspend fun process(request: ScreenshotRequest): ScreenshotRequest {
+        var result = request
+
+        // Apply work profile screenshots policy:
+        //
+        // If the focused app belongs to a work profile, transforms a full screen
+        // (or partial) screenshot request to a task snapshot (provided image) screenshot.
+
+        // Whenever displayContentInfo is fetched, the topComponent is also populated
+        // regardless of the managed profile status.
+
+        if (request.type != TAKE_SCREENSHOT_PROVIDED_IMAGE &&
+            flags.isEnabled(SCREENSHOT_WORK_PROFILE_POLICY)
         ) {
 
-        if (request.type == TAKE_SCREENSHOT_PROVIDED_IMAGE) {
+            val info = policy.findPrimaryContent(policy.getDefaultDisplayId())
-            val image = HardwareBitmapBundler.bundleToHardwareBitmap(request.bitmapBundle)
 
-            controller.handleImageAsScreenshot(
+            result = if (policy.isManagedProfile(info.userId)) {
-                image, request.boundsInScreen, request.insets,
+                val image = capture.captureTask(info.taskId)
-                request.taskId, request.userId, request.topComponent, onSavedListener, callback
+                    ?: error("Task snapshot returned a null Bitmap!")
+
+                // Provide the task snapshot as the screenshot
+                ScreenshotRequest(
+                    TAKE_SCREENSHOT_PROVIDED_IMAGE, request.source,
+                    HardwareBitmapBundler.hardwareBitmapToBundle(image),
+                    info.bounds, Insets.NONE, info.taskId, info.userId, info.component
                 )
-            return
+            } else {
-        }
+                // Create a new request of the same type which includes the top component
-
+                ScreenshotRequest(request.source, request.type, info.component)
-        when (request.type) {
-            TAKE_SCREENSHOT_FULLSCREEN ->
-                controller.takeScreenshotFullscreen(null, onSavedListener, callback)
-            TAKE_SCREENSHOT_SELECTED_REGION ->
-                controller.takeScreenshotPartial(null, onSavedListener, callback)
-            else -> Log.w(TAG, "Invalid screenshot option: ${request.type}")
             }
         }
 
-    companion object {
+        return result
-        const val TAG: String = "RequestProcessor"
+    }
+
+    /**
+     * Note: This is for compatibility with existing Java. Prefer the suspending function when
+     * calling from a Coroutine context.
+     *
+     * @param request the request to process
+     * @param callback the callback to provide the processed request, invoked from the main thread
+     */
+    fun processAsync(request: ScreenshotRequest, callback: Consumer<ScreenshotRequest>) {
+        mainScope.launch {
+            val result = process(request)
+            callback.accept(result)
+        }
     }
 }
+
+private const val TAG = "RequestProcessor"

packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotPolicy.kt

@@ -0,0 +1,50 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.screenshot
+
+import android.annotation.UserIdInt
+import android.content.ComponentName
+import android.graphics.Rect
+import android.view.Display
+
+/**
+ * Provides policy decision-making information to screenshot request handling.
+ */
+interface ScreenshotPolicy {
+
+    /** @return true if the user is a managed profile (a.k.a. work profile) */
+    suspend fun isManagedProfile(@UserIdInt userId: Int): Boolean
+
+    /**
+     * Requests information about the owner of display content which occupies a majority of the
+     * screenshot and/or has most recently been interacted with at the time the screenshot was
+     * requested.
+     *
+     * @param displayId the id of the display to inspect
+     * @return content info for the primary content on the display
+     */
+    suspend fun findPrimaryContent(displayId: Int): DisplayContentInfo
+
+    data class DisplayContentInfo(
+        val component: ComponentName,
+        val bounds: Rect,
+        @UserIdInt val userId: Int,
+        val taskId: Int,
+    )
+
+    fun getDefaultDisplayId(): Int = Display.DEFAULT_DISPLAY
+}

packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotPolicyImpl.kt

@@ -0,0 +1,178 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.screenshot
+
+import android.annotation.UserIdInt
+import android.app.ActivityTaskManager
+import android.app.ActivityTaskManager.RootTaskInfo
+import android.app.IActivityTaskManager
+import android.app.WindowConfiguration
+import android.app.WindowConfiguration.activityTypeToString
+import android.app.WindowConfiguration.windowingModeToString
+import android.content.ComponentName
+import android.content.Context
+import android.content.Intent
+import android.graphics.Rect
+import android.os.Process
+import android.os.RemoteException
+import android.os.UserManager
+import android.util.Log
+import android.view.Display.DEFAULT_DISPLAY
+import com.android.internal.infra.ServiceConnector
+import com.android.systemui.SystemUIService
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.screenshot.ScreenshotPolicy.DisplayContentInfo
+import java.util.Arrays
+import javax.inject.Inject
+import kotlin.coroutines.resume
+import kotlin.coroutines.suspendCoroutine
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.withContext
+
+@SysUISingleton
+internal class ScreenshotPolicyImpl @Inject constructor(
+    context: Context,
+    private val userMgr: UserManager,
+    private val atmService: IActivityTaskManager,
+    @Background val bgDispatcher: CoroutineDispatcher,
+) : ScreenshotPolicy {
+
+    private val systemUiContent =
+        DisplayContentInfo(
+            ComponentName(context, SystemUIService::class.java),
+            Rect(),
+            ActivityTaskManager.INVALID_TASK_ID,
+            Process.myUserHandle().identifier,
+        )
+
+    private val proxyConnector: ServiceConnector<IScreenshotProxy> =
+        ServiceConnector.Impl(
+            context,
+            Intent(context, ScreenshotProxyService::class.java),
+            Context.BIND_AUTO_CREATE or Context.BIND_WAIVE_PRIORITY or Context.BIND_NOT_VISIBLE,
+            context.userId,
+            IScreenshotProxy.Stub::asInterface
+        )
+
+    override fun getDefaultDisplayId(): Int {
+        return DEFAULT_DISPLAY
+    }
+
+    override suspend fun isManagedProfile(@UserIdInt userId: Int): Boolean {
+        return withContext(bgDispatcher) { userMgr.isManagedProfile(userId) }
+    }
+
+    private fun nonPipVisibleTask(info: RootTaskInfo): Boolean {
+        return info.windowingMode != WindowConfiguration.WINDOWING_MODE_PINNED &&
+            info.isVisible &&
+            info.isRunning &&
+            info.numActivities > 0 &&
+            info.topActivity != null &&
+            info.childTaskIds.isNotEmpty()
+    }
+
+    /**
+     * Uses RootTaskInfo from ActivityTaskManager to guess at the primary focused task within a
+     * display. If no task is visible or the top task is covered by a system window, the info
+     * reported will reference a SystemUI component instead.
+     */
+    override suspend fun findPrimaryContent(displayId: Int): DisplayContentInfo {
+        // Determine if the notification shade is expanded. If so, task windows are not
+        // visible behind it, so the screenshot should instead be associated with SystemUI.
+        if (isNotificationShadeExpanded()) {
+            return systemUiContent
+        }
+
+        val taskInfoList = getAllRootTaskInfosOnDisplay(displayId)
+        if (DEBUG) {
+            debugLogRootTaskInfos(taskInfoList)
+        }
+
+        // If no visible task is located, then report SystemUI as the foreground content
+        val target = taskInfoList.firstOrNull(::nonPipVisibleTask) ?: return systemUiContent
+
+        val topActivity: ComponentName = target.topActivity ?: error("should not be null")
+        val topChildTask = target.childTaskIds.size - 1
+        val childTaskId = target.childTaskIds[topChildTask]
+        val childTaskUserId = target.childTaskUserIds[topChildTask]
+        val childTaskBounds = target.childTaskBounds[topChildTask]
+
+        return DisplayContentInfo(topActivity, childTaskBounds, childTaskId, childTaskUserId)
+    }
+
+    private fun debugLogRootTaskInfos(taskInfoList: List<RootTaskInfo>) {
+        for (info in taskInfoList) {
+            Log.d(
+                TAG,
+                "[root task info] " +
+                    "taskId=${info.taskId} " +
+                    "parentTaskId=${info.parentTaskId} " +
+                    "position=${info.position} " +
+                    "positionInParent=${info.positionInParent} " +
+                    "isVisible=${info.isVisible()} " +
+                    "visible=${info.visible} " +
+                    "isFocused=${info.isFocused} " +
+                    "isSleeping=${info.isSleeping} " +
+                    "isRunning=${info.isRunning} " +
+                    "windowMode=${windowingModeToString(info.windowingMode)} " +
+                    "activityType=${activityTypeToString(info.activityType)} " +
+                    "topActivity=${info.topActivity} " +
+                    "topActivityInfo=${info.topActivityInfo} " +
+                    "numActivities=${info.numActivities} " +
+                    "childTaskIds=${Arrays.toString(info.childTaskIds)} " +
+                    "childUserIds=${Arrays.toString(info.childTaskUserIds)} " +
+                    "childTaskBounds=${Arrays.toString(info.childTaskBounds)} " +
+                    "childTaskNames=${Arrays.toString(info.childTaskNames)}"
+            )
+
+            for (j in 0 until info.childTaskIds.size) {
+                Log.d(TAG, "    *** [$j] ******")
+                Log.d(TAG, "        ***  childTaskIds[$j]: ${info.childTaskIds[j]}")
+                Log.d(TAG, "        ***  childTaskUserIds[$j]: ${info.childTaskUserIds[j]}")
+                Log.d(TAG, "        ***  childTaskBounds[$j]: ${info.childTaskBounds[j]}")
+                Log.d(TAG, "        ***  childTaskNames[$j]: ${info.childTaskNames[j]}")
+            }
+        }
+    }
+
+    private suspend fun getAllRootTaskInfosOnDisplay(displayId: Int): List<RootTaskInfo> =
+        withContext(bgDispatcher) {
+            try {
+                atmService.getAllRootTaskInfosOnDisplay(displayId)
+            } catch (e: RemoteException) {
+                Log.e(TAG, "getAllRootTaskInfosOnDisplay", e)
+                listOf()
+            }
+        }
+
+    private suspend fun isNotificationShadeExpanded(): Boolean = suspendCoroutine { k ->
+        proxyConnector
+            .postForResult { it.isNotificationShadeExpanded }
+            .whenComplete { expanded, error ->
+                if (error != null) {
+                    Log.e(TAG, "isNotificationShadeExpanded", error)
+                }
+                k.resume(expanded ?: false)
+            }
+    }
+
+    companion object {
+        const val TAG: String = "ScreenshotPolicyImpl"
+        const val DEBUG: Boolean = false
+    }
+}

packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotProxyService.kt

@@ -0,0 +1,51 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.systemui.screenshot
+
+import android.app.Service
+import android.content.Intent
+import android.os.IBinder
+import android.util.Log
+import com.android.systemui.statusbar.phone.panelstate.PanelExpansionStateManager
+import javax.inject.Inject
+
+/**
+ * Provides state from the main SystemUI process on behalf of the Screenshot process.
+ */
+internal class ScreenshotProxyService @Inject constructor(
+    private val mExpansionMgr: PanelExpansionStateManager
+) : Service() {
+
+    private val mBinder: IBinder = object : IScreenshotProxy.Stub() {
+        /**
+         * @return true when the notification shade is partially or fully expanded.
+         */
+        override fun isNotificationShadeExpanded(): Boolean {
+            val expanded = !mExpansionMgr.isClosed()
+            Log.d(TAG, "isNotificationShadeExpanded(): $expanded")
+            return expanded
+        }
+    }
+
+    override fun onBind(intent: Intent): IBinder? {
+        Log.d(TAG, "onBind: $intent")
+        return mBinder
+    }
+
+    companion object {
+        const val TAG = "ScreenshotProxyService"
+    }
+}

packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java

@@ -22,6 +22,7 @@ import static android.content.Intent.ACTION_CLOSE_SYSTEM_DIALOGS;
 import static com.android.internal.util.ScreenshotHelper.SCREENSHOT_MSG_PROCESS_COMPLETE;
 import static com.android.internal.util.ScreenshotHelper.SCREENSHOT_MSG_URI;
 import static com.android.systemui.flags.Flags.SCREENSHOT_REQUEST_PROCESSOR;
+import static com.android.systemui.flags.Flags.SCREENSHOT_WORK_PROFILE_POLICY;
 import static com.android.systemui.screenshot.LogConfig.DEBUG_CALLBACK;
 import static com.android.systemui.screenshot.LogConfig.DEBUG_DISMISS;
 import static com.android.systemui.screenshot.LogConfig.DEBUG_SERVICE;
@@ -97,7 +98,7 @@ public class TakeScreenshotService extends Service {
     };
 
     /** Informs about coarse grained state of the Controller. */
-    interface RequestCallback {
+    public interface RequestCallback {
         /** Respond to the current request indicating the screenshot request failed. */
         void reportError();
 
@@ -124,6 +125,7 @@ public class TakeScreenshotService extends Service {
         mBgExecutor = bgExecutor;
         mFeatureFlags = featureFlags;
         mFeatureFlags.addListener(SCREENSHOT_REQUEST_PROCESSOR, FlagEvent::requestNoRestart);
+        mFeatureFlags.addListener(SCREENSHOT_WORK_PROFILE_POLICY, FlagEvent::requestNoRestart);
         mProcessor = processor;
     }
 
@@ -229,49 +231,57 @@ public class TakeScreenshotService extends Service {
 
         if (mFeatureFlags.isEnabled(SCREENSHOT_REQUEST_PROCESSOR)) {
             Log.d(TAG, "handleMessage: Using request processor");
-            mProcessor.processRequest(screenshotRequest, uriConsumer, requestCallback);
+            mProcessor.processAsync(screenshotRequest,
+                    (request) -> dispatchToController(request, uriConsumer, requestCallback));
             return true;
         }
 
-        switch (screenshotRequest.getType()) {
+        dispatchToController(screenshotRequest, uriConsumer, requestCallback);
+        return true;
+    }
+
+    private void dispatchToController(ScreenshotHelper.ScreenshotRequest request,
+            Consumer<Uri> uriConsumer, RequestCallback callback) {
+
+        ComponentName topComponent = request.getTopComponent();
+
+        switch (request.getType()) {
             case WindowManager.TAKE_SCREENSHOT_FULLSCREEN:
                 if (DEBUG_SERVICE) {
                     Log.d(TAG, "handleMessage: TAKE_SCREENSHOT_FULLSCREEN");
                 }
-                mScreenshot.takeScreenshotFullscreen(topComponent, uriConsumer, requestCallback);
+                mScreenshot.takeScreenshotFullscreen(topComponent, uriConsumer, callback);
                 break;
             case WindowManager.TAKE_SCREENSHOT_SELECTED_REGION:
                 if (DEBUG_SERVICE) {
                     Log.d(TAG, "handleMessage: TAKE_SCREENSHOT_SELECTED_REGION");
                 }
-                mScreenshot.takeScreenshotPartial(topComponent, uriConsumer, requestCallback);
+                mScreenshot.takeScreenshotPartial(topComponent, uriConsumer, callback);
                 break;
             case WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE:
                 if (DEBUG_SERVICE) {
                     Log.d(TAG, "handleMessage: TAKE_SCREENSHOT_PROVIDED_IMAGE");
                 }
                 Bitmap screenshot = ScreenshotHelper.HardwareBitmapBundler.bundleToHardwareBitmap(
-                        screenshotRequest.getBitmapBundle());
+                        request.getBitmapBundle());
-                Rect screenBounds = screenshotRequest.getBoundsInScreen();
+                Rect screenBounds = request.getBoundsInScreen();
-                Insets insets = screenshotRequest.getInsets();
+                Insets insets = request.getInsets();
-                int taskId = screenshotRequest.getTaskId();
+                int taskId = request.getTaskId();
-                int userId = screenshotRequest.getUserId();
+                int userId = request.getUserId();
 
                 if (screenshot == null) {
                     Log.e(TAG, "Got null bitmap from screenshot message");
                     mNotificationsController.notifyScreenshotError(
                             R.string.screenshot_failed_to_capture_text);
-                    requestCallback.reportError();
+                    callback.reportError();
                 } else {
                     mScreenshot.handleImageAsScreenshot(screenshot, screenBounds, insets,
-                            taskId, userId, topComponent, uriConsumer, requestCallback);
+                            taskId, userId, topComponent, uriConsumer, callback);
                 }
                 break;
             default:
-                Log.w(TAG, "Invalid screenshot option: " + msg.what);
+                Log.w(TAG, "Invalid screenshot option: " + request.getType());
-                return false;
         }
-        return true;
     }
 
     private static void sendComplete(Messenger target) {

packages/SystemUI/src/com/android/systemui/screenshot/dagger/ScreenshotModule.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2019 The Android Open Source Project
+ * Copyright (C) 2022 The Android Open Source Project
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -20,6 +20,9 @@ import android.app.Service;
 
 import com.android.systemui.screenshot.ImageCapture;
 import com.android.systemui.screenshot.ImageCaptureImpl;
+import com.android.systemui.screenshot.ScreenshotPolicy;
+import com.android.systemui.screenshot.ScreenshotPolicyImpl;
+import com.android.systemui.screenshot.ScreenshotProxyService;
 import com.android.systemui.screenshot.TakeScreenshotService;
 
 import dagger.Binds;
@@ -33,12 +36,20 @@ import dagger.multibindings.IntoMap;
 @Module
 public abstract class ScreenshotModule {
 
-    /** */
     @Binds
     @IntoMap
     @ClassKey(TakeScreenshotService.class)
-    public abstract Service bindTakeScreenshotService(TakeScreenshotService service);
+    abstract Service bindTakeScreenshotService(TakeScreenshotService service);
 
     @Binds
-    public abstract ImageCapture bindImageCapture(ImageCaptureImpl capture);
+    @IntoMap
+    @ClassKey(ScreenshotProxyService.class)
+    abstract Service bindScreenshotProxyService(ScreenshotProxyService service);
+
+    @Binds
+    abstract ScreenshotPolicy bindScreenshotPolicyImpl(ScreenshotPolicyImpl impl);
+
+    @Binds
+    abstract ImageCapture bindImageCaptureImpl(ImageCaptureImpl capture);
+
 }

packages/SystemUI/tests/src/com/android/systemui/screenshot/FakeImageCapture.kt

@@ -0,0 +1,39 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.systemui.screenshot
+
+import android.graphics.Bitmap
+import android.graphics.Rect
+
+internal class FakeImageCapture : ImageCapture {
+
+    var requestedDisplayId: Int? = null
+    var requestedDisplayCrop: Rect? = null
+    var requestedTaskId: Int? = null
+
+    var image: Bitmap? = null
+
+    override fun captureDisplay(displayId: Int, crop: Rect?): Bitmap? {
+        requestedDisplayId = displayId
+        requestedDisplayCrop = crop
+        return image
+    }
+
+    override suspend fun captureTask(taskId: Int): Bitmap? {
+        requestedTaskId = taskId
+        return image
+    }
+}

packages/SystemUI/tests/src/com/android/systemui/screenshot/FakeScreenshotPolicy.kt

@@ -0,0 +1,40 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.screenshot
+
+import com.android.systemui.screenshot.ScreenshotPolicy.DisplayContentInfo
+
+internal class FakeScreenshotPolicy : ScreenshotPolicy {
+
+    private val userTypes = mutableMapOf<Int, Boolean>()
+    private val contentInfo = mutableMapOf<Int, DisplayContentInfo?>()
+
+    fun setManagedProfile(userId: Int, managedUser: Boolean) {
+        userTypes[userId] = managedUser
+    }
+    override suspend fun isManagedProfile(userId: Int): Boolean {
+        return userTypes[userId] ?: error("No managedProfile value set for userId $userId")
+    }
+
+    fun setDisplayContentInfo(userId: Int, contentInfo: DisplayContentInfo) {
+        this.contentInfo[userId] = contentInfo
+    }
+
+    override suspend fun findPrimaryContent(displayId: Int): DisplayContentInfo {
+        return contentInfo[displayId] ?: error("No DisplayContentInfo set for displayId $displayId")
+    }
+}

packages/SystemUI/tests/src/com/android/systemui/screenshot/ImageCaptureImplTest.kt

@@ -27,6 +27,8 @@ import android.view.SurfaceControl.ScreenshotHardwareBuffer
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.util.mockito.mock
 import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.Dispatchers
 import org.junit.Test
 import org.junit.runner.RunWith
 
@@ -37,7 +39,10 @@ import org.junit.runner.RunWith
 class ImageCaptureImplTest : SysuiTestCase() {
     private val displayManager = mock<DisplayManager>()
     private val atmService = mock<IActivityTaskManager>()
-    private val capture = TestableImageCaptureImpl(displayManager, atmService)
+    private val capture = TestableImageCaptureImpl(
+        displayManager,
+        atmService,
+        Dispatchers.Unconfined)
 
     @Test
     fun captureDisplayWithCrop() {
@@ -59,9 +64,10 @@ class ImageCaptureImplTest : SysuiTestCase() {
 
     class TestableImageCaptureImpl(
         displayManager: DisplayManager,
-        atmService: IActivityTaskManager
+        atmService: IActivityTaskManager,
+        bgDispatcher: CoroutineDispatcher
     ) :
-        ImageCaptureImpl(displayManager, atmService) {
+        ImageCaptureImpl(displayManager, atmService, bgDispatcher) {
 
         var token: IBinder? = null
         var width: Int? = null

packages/SystemUI/tests/src/com/android/systemui/screenshot/RequestProcessorTest.kt

@@ -22,86 +22,253 @@ import android.graphics.ColorSpace
 import android.graphics.Insets
 import android.graphics.Rect
 import android.hardware.HardwareBuffer
-import android.net.Uri
+import android.os.Bundle
 import android.view.WindowManager.ScreenshotSource.SCREENSHOT_KEY_CHORD
 import android.view.WindowManager.ScreenshotSource.SCREENSHOT_OTHER
 import android.view.WindowManager.TAKE_SCREENSHOT_FULLSCREEN
-import android.view.WindowManager.TAKE_SCREENSHOT_SELECTED_REGION
 import android.view.WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE
-
+import android.view.WindowManager.TAKE_SCREENSHOT_SELECTED_REGION
 import com.android.internal.util.ScreenshotHelper.HardwareBitmapBundler
+import com.android.internal.util.ScreenshotHelper.HardwareBitmapBundler.bundleToHardwareBitmap
 import com.android.internal.util.ScreenshotHelper.ScreenshotRequest
-import com.android.systemui.screenshot.TakeScreenshotService.RequestCallback
+import com.android.systemui.flags.FakeFeatureFlags
-import com.android.systemui.util.mockito.argumentCaptor
+import com.android.systemui.flags.Flags
-import com.android.systemui.util.mockito.mock
+import com.android.systemui.screenshot.ScreenshotPolicy.DisplayContentInfo
 import com.google.common.truth.Truth.assertThat
-import java.util.function.Consumer
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.runBlocking
 import org.junit.Test
-import org.mockito.Mockito.eq
+
-import org.mockito.Mockito.verify
+private const val USER_ID = 1
-import org.mockito.Mockito.isNull
+private const val TASK_ID = 1
 
 class RequestProcessorTest {
-    private val controller = mock<ScreenshotController>()
+    private val imageCapture = FakeImageCapture()
-    private val bitmapCaptor = argumentCaptor<Bitmap>()
+    private val component = ComponentName("android.test", "android.test.Component")
+    private val bounds = Rect(25, 25, 75, 75)
 
+    private val scope = CoroutineScope(Dispatchers.Unconfined)
+    private val dispatcher = Dispatchers.Unconfined
+    private val policy = FakeScreenshotPolicy()
+    private val flags = FakeFeatureFlags()
+
+    /** Tests the Java-compatible function wrapper, ensures callback is invoked. */
     @Test
-    fun testFullScreenshot() {
+    fun testProcessAsync() {
+        flags.set(Flags.SCREENSHOT_WORK_PROFILE_POLICY, false)
+
         val request = ScreenshotRequest(TAKE_SCREENSHOT_FULLSCREEN, SCREENSHOT_KEY_CHORD)
-        val onSavedListener = mock<Consumer<Uri>>()
+        val processor = RequestProcessor(imageCapture, policy, flags, scope)
-        val callback = mock<RequestCallback>()
-        val processor = RequestProcessor(controller)
 
-        processor.processRequest(request, onSavedListener, callback)
+        var result: ScreenshotRequest? = null
+        var callbackCount = 0
+        val callback: (ScreenshotRequest) -> Unit = { processedRequest: ScreenshotRequest ->
+            result = processedRequest
+            callbackCount++
+        }
 
-        verify(controller).takeScreenshotFullscreen(/* topComponent */ isNull(),
+        // runs synchronously, using Unconfined Dispatcher
-            eq(onSavedListener), eq(callback))
+        processor.processAsync(request, callback)
+
+        // Callback invoked once returning the same request (no changes)
+        assertThat(callbackCount).isEqualTo(1)
+        assertThat(result).isEqualTo(request)
     }
 
     @Test
-    fun testSelectedRegionScreenshot() {
+    fun testFullScreenshot_workProfilePolicyDisabled() = runBlocking {
+        flags.set(Flags.SCREENSHOT_WORK_PROFILE_POLICY, false)
+
+        val request = ScreenshotRequest(TAKE_SCREENSHOT_FULLSCREEN, SCREENSHOT_KEY_CHORD)
+        val processor = RequestProcessor(imageCapture, policy, flags, scope)
+
+        val processedRequest = processor.process(request)
+
+        // No changes
+        assertThat(processedRequest).isEqualTo(request)
+    }
+
+    @Test
+    fun testFullScreenshot() = runBlocking {
+        flags.set(Flags.SCREENSHOT_WORK_PROFILE_POLICY, true)
+
+        // Indicate that the primary content belongs to a normal user
+        policy.setManagedProfile(USER_ID, false)
+        policy.setDisplayContentInfo(
+            policy.getDefaultDisplayId(),
+            DisplayContentInfo(component, bounds, USER_ID, TASK_ID))
+
+        val request = ScreenshotRequest(TAKE_SCREENSHOT_FULLSCREEN, SCREENSHOT_KEY_CHORD)
+        val processor = RequestProcessor(imageCapture, policy, flags, scope)
+
+        val processedRequest = processor.process(request)
+
+        // Request has topComponent added, but otherwise unchanged.
+        assertThat(processedRequest.type).isEqualTo(TAKE_SCREENSHOT_FULLSCREEN)
+        assertThat(processedRequest.topComponent).isEqualTo(component)
+    }
+
+    @Test
+    fun testFullScreenshot_managedProfile() = runBlocking {
+        flags.set(Flags.SCREENSHOT_WORK_PROFILE_POLICY, true)
+
+        // Provide a fake task bitmap when asked
+        val bitmap = makeHardwareBitmap(100, 100)
+        imageCapture.image = bitmap
+
+        // Indicate that the primary content belongs to a manged profile
+        policy.setManagedProfile(USER_ID, true)
+        policy.setDisplayContentInfo(policy.getDefaultDisplayId(),
+            DisplayContentInfo(component, bounds, USER_ID, TASK_ID))
+
+        val request = ScreenshotRequest(TAKE_SCREENSHOT_FULLSCREEN, SCREENSHOT_KEY_CHORD)
+        val processor = RequestProcessor(imageCapture, policy, flags, scope)
+
+        val processedRequest = processor.process(request)
+
+        // Expect a task snapshot is taken, overriding the full screen mode
+        assertThat(processedRequest.type).isEqualTo(TAKE_SCREENSHOT_PROVIDED_IMAGE)
+        assertThat(bitmap.equalsHardwareBitmapBundle(processedRequest.bitmapBundle)).isTrue()
+        assertThat(processedRequest.boundsInScreen).isEqualTo(bounds)
+        assertThat(processedRequest.insets).isEqualTo(Insets.NONE)
+        assertThat(processedRequest.taskId).isEqualTo(TASK_ID)
+        assertThat(imageCapture.requestedTaskId).isEqualTo(TASK_ID)
+        assertThat(processedRequest.userId).isEqualTo(USER_ID)
+        assertThat(processedRequest.topComponent).isEqualTo(component)
+    }
+
+    @Test
+    fun testSelectedRegionScreenshot_workProfilePolicyDisabled() = runBlocking {
+        flags.set(Flags.SCREENSHOT_WORK_PROFILE_POLICY, false)
+
         val request = ScreenshotRequest(TAKE_SCREENSHOT_SELECTED_REGION, SCREENSHOT_KEY_CHORD)
-        val onSavedListener = mock<Consumer<Uri>>()
+        val processor = RequestProcessor(imageCapture, policy, flags, scope)
-        val callback = mock<RequestCallback>()
-        val processor = RequestProcessor(controller)
 
-        processor.processRequest(request, onSavedListener, callback)
+        val processedRequest = processor.process(request)
 
-        verify(controller).takeScreenshotPartial(/* topComponent */ isNull(),
+        // No changes
-            eq(onSavedListener), eq(callback))
+        assertThat(processedRequest).isEqualTo(request)
      }
 
     @Test
-    fun testProvidedImageScreenshot() {
+    fun testSelectedRegionScreenshot() = runBlocking {
-        val taskId = 1111
+        flags.set(Flags.SCREENSHOT_WORK_PROFILE_POLICY, true)
-        val userId = 2222
-        val bounds = Rect(50, 50, 150, 150)
-        val topComponent = ComponentName("test", "test")
-        val processor = RequestProcessor(controller)
 
-        val buffer = HardwareBuffer.create(100, 100, HardwareBuffer.RGBA_8888, 1,
+        val request = ScreenshotRequest(TAKE_SCREENSHOT_SELECTED_REGION, SCREENSHOT_KEY_CHORD)
-            HardwareBuffer.USAGE_GPU_SAMPLED_IMAGE)
+        val processor = RequestProcessor(imageCapture, policy, flags, scope)
-        val bitmap = Bitmap.wrapHardwareBuffer(buffer, ColorSpace.get(ColorSpace.Named.SRGB))!!
+
+        policy.setManagedProfile(USER_ID, false)
+        policy.setDisplayContentInfo(policy.getDefaultDisplayId(),
+            DisplayContentInfo(component, bounds, USER_ID, TASK_ID))
+
+        val processedRequest = processor.process(request)
+
+        // Request has topComponent added, but otherwise unchanged.
+        assertThat(processedRequest.type).isEqualTo(TAKE_SCREENSHOT_FULLSCREEN)
+        assertThat(processedRequest.topComponent).isEqualTo(component)
+    }
+
+    @Test
+    fun testSelectedRegionScreenshot_managedProfile() = runBlocking {
+        flags.set(Flags.SCREENSHOT_WORK_PROFILE_POLICY, true)
+
+        // Provide a fake task bitmap when asked
+        val bitmap = makeHardwareBitmap(100, 100)
+        imageCapture.image = bitmap
+
+        val request = ScreenshotRequest(TAKE_SCREENSHOT_SELECTED_REGION, SCREENSHOT_KEY_CHORD)
+        val processor = RequestProcessor(imageCapture, policy, flags, scope)
+
+        // Indicate that the primary content belongs to a manged profile
+        policy.setManagedProfile(USER_ID, true)
+        policy.setDisplayContentInfo(policy.getDefaultDisplayId(),
+            DisplayContentInfo(component, bounds, USER_ID, TASK_ID))
+
+        val processedRequest = processor.process(request)
+
+        // Expect a task snapshot is taken, overriding the selected region mode
+        assertThat(processedRequest.type).isEqualTo(TAKE_SCREENSHOT_PROVIDED_IMAGE)
+        assertThat(bitmap.equalsHardwareBitmapBundle(processedRequest.bitmapBundle)).isTrue()
+        assertThat(processedRequest.boundsInScreen).isEqualTo(bounds)
+        assertThat(processedRequest.insets).isEqualTo(Insets.NONE)
+        assertThat(processedRequest.taskId).isEqualTo(TASK_ID)
+        assertThat(imageCapture.requestedTaskId).isEqualTo(TASK_ID)
+        assertThat(processedRequest.userId).isEqualTo(USER_ID)
+        assertThat(processedRequest.topComponent).isEqualTo(component)
+    }
+
+    @Test
+    fun testProvidedImageScreenshot_workProfilePolicyDisabled() = runBlocking {
+        flags.set(Flags.SCREENSHOT_WORK_PROFILE_POLICY, false)
+
+        val bounds = Rect(50, 50, 150, 150)
+        val processor = RequestProcessor(imageCapture, policy, flags, scope)
+
+        val bitmap = makeHardwareBitmap(100, 100)
         val bitmapBundle = HardwareBitmapBundler.hardwareBitmapToBundle(bitmap)
 
         val request = ScreenshotRequest(TAKE_SCREENSHOT_PROVIDED_IMAGE, SCREENSHOT_OTHER,
-            bitmapBundle, bounds, Insets.NONE, taskId, userId, topComponent)
+            bitmapBundle, bounds, Insets.NONE, TASK_ID, USER_ID, component)
 
-        val onSavedListener = mock<Consumer<Uri>>()
+        val processedRequest = processor.process(request)
-        val callback = mock<RequestCallback>()
 
-        processor.processRequest(request, onSavedListener, callback)
+        // No changes
-
+        assertThat(processedRequest).isEqualTo(request)
-        verify(controller).handleImageAsScreenshot(
-            bitmapCaptor.capture(), eq(bounds), eq(Insets.NONE), eq(taskId), eq(userId),
-            eq(topComponent), eq(onSavedListener), eq(callback)
-        )
-
-        assertThat(bitmapCaptor.value.equalsHardwareBitmap(bitmap)).isTrue()
     }
 
-    private fun Bitmap.equalsHardwareBitmap(bitmap: Bitmap): Boolean {
+    @Test
-        return bitmap.hardwareBuffer == this.hardwareBuffer &&
+    fun testProvidedImageScreenshot() = runBlocking {
-                bitmap.colorSpace == this.colorSpace
+        flags.set(Flags.SCREENSHOT_WORK_PROFILE_POLICY, true)
+
+        val bounds = Rect(50, 50, 150, 150)
+        val processor = RequestProcessor(imageCapture, policy, flags, scope)
+
+        policy.setManagedProfile(USER_ID, false)
+
+        val bitmap = makeHardwareBitmap(100, 100)
+        val bitmapBundle = HardwareBitmapBundler.hardwareBitmapToBundle(bitmap)
+
+        val request = ScreenshotRequest(TAKE_SCREENSHOT_PROVIDED_IMAGE, SCREENSHOT_OTHER,
+            bitmapBundle, bounds, Insets.NONE, TASK_ID, USER_ID, component)
+
+        val processedRequest = processor.process(request)
+
+        // No changes
+        assertThat(processedRequest).isEqualTo(request)
+    }
+
+    @Test
+    fun testProvidedImageScreenshot_managedProfile() = runBlocking {
+        flags.set(Flags.SCREENSHOT_WORK_PROFILE_POLICY, true)
+
+        val bounds = Rect(50, 50, 150, 150)
+        val processor = RequestProcessor(imageCapture, policy, flags, scope)
+
+        // Indicate that the screenshot belongs to a manged profile
+        policy.setManagedProfile(USER_ID, true)
+
+        val bitmap = makeHardwareBitmap(100, 100)
+        val bitmapBundle = HardwareBitmapBundler.hardwareBitmapToBundle(bitmap)
+
+        val request = ScreenshotRequest(TAKE_SCREENSHOT_PROVIDED_IMAGE, SCREENSHOT_OTHER,
+            bitmapBundle, bounds, Insets.NONE, TASK_ID, USER_ID, component)
+
+        val processedRequest = processor.process(request)
+
+        // Work profile, but already a task snapshot, so no changes
+        assertThat(processedRequest).isEqualTo(request)
+    }
+
+    private fun makeHardwareBitmap(width: Int, height: Int): Bitmap {
+        val buffer = HardwareBuffer.create(width, height, HardwareBuffer.RGBA_8888, 1,
+            HardwareBuffer.USAGE_GPU_SAMPLED_IMAGE)
+        return Bitmap.wrapHardwareBuffer(buffer, ColorSpace.get(ColorSpace.Named.SRGB))!!
+    }
+
+    private fun Bitmap.equalsHardwareBitmapBundle(bundle: Bundle): Boolean {
+        val provided = bundleToHardwareBitmap(bundle)
+        return provided.hardwareBuffer == this.hardwareBuffer &&
+                provided.colorSpace == this.colorSpace
     }
 }