[Partial Screensharing] Add enterprise policies handling flag

Adds a feature flag to enable/disable enterprise
policies handling in partial screen sharing feature.

Also adds a logic to forbid cross-profile sharing
using recent tasks when this flag is disabled.

Cross-profile sharing using the app selector
is already forbidden.

Bug: 233348916
Test: com.android.systemui.mediaprojection.appselector.MediaProjectionAppSelectorControllerTest
Test: manual check that recent tasks don't have apps from
  a profile different from the host app profile
Change-Id: I7247cbc9e45565e157ea70de77f5f7e6563723ad
diff --git a/packages/SystemUI/src/com/android/systemui/flags/Flags.kt b/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
index c9a2a0b..08160a3 100644
--- a/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
+++ b/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
@@ -406,6 +406,17 @@
     val WM_DESKTOP_WINDOWING_2 =
         sysPropBooleanFlag(1112, "persist.wm.debug.desktop_mode_2", default = false)
 
+    // TODO(b/254513207): Tracking Bug to delete
+    @Keep
+    @JvmField
+    val WM_ENABLE_PARTIAL_SCREEN_SHARING_ENTERPRISE_POLICIES =
+        unreleasedFlag(
+            1113,
+            name = "screen_record_enterprise_policies",
+            namespace = DeviceConfig.NAMESPACE_WINDOW_MANAGER,
+            teamfood = false
+        )
+
     // 1200 - predictive back
     @Keep
     @JvmField
diff --git a/packages/SystemUI/src/com/android/systemui/media/MediaProjectionAppSelectorActivity.kt b/packages/SystemUI/src/com/android/systemui/media/MediaProjectionAppSelectorActivity.kt
index ceb4845..a692ad7 100644
--- a/packages/SystemUI/src/com/android/systemui/media/MediaProjectionAppSelectorActivity.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/MediaProjectionAppSelectorActivity.kt
@@ -18,6 +18,7 @@
 import android.app.ActivityOptions
 import android.content.Intent
 import android.content.res.Configuration
+import android.content.res.Resources
 import android.media.projection.IMediaProjection
 import android.media.projection.MediaProjectionManager.EXTRA_MEDIA_PROJECTION
 import android.os.Binder
@@ -27,6 +28,7 @@
 import android.os.UserHandle
 import android.view.ViewGroup
 import com.android.internal.annotations.VisibleForTesting
+import com.android.internal.app.AbstractMultiProfilePagerAdapter.MyUserIdProvider
 import com.android.internal.app.ChooserActivity
 import com.android.internal.app.ResolverListController
 import com.android.internal.app.chooser.NotSelectableTargetInfo
@@ -59,16 +61,12 @@
     private lateinit var configurationController: ConfigurationController
     private lateinit var controller: MediaProjectionAppSelectorController
     private lateinit var recentsViewController: MediaProjectionRecentsViewController
+    private lateinit var component: MediaProjectionAppSelectorComponent
 
     override fun getLayoutResource() = R.layout.media_projection_app_selector
 
     public override fun onCreate(bundle: Bundle?) {
-        val component =
-            componentFactory.create(
-                activity = this,
-                view = this,
-                resultHandler = this
-            )
+        component = componentFactory.create(activity = this, view = this, resultHandler = this)
 
         // Create a separate configuration controller for this activity as the configuration
         // might be different from the global one
@@ -76,11 +74,12 @@
         controller = component.controller
         recentsViewController = component.recentsViewController
 
-        val queryIntent = Intent(Intent.ACTION_MAIN).apply { addCategory(Intent.CATEGORY_LAUNCHER) }
-        intent.putExtra(Intent.EXTRA_INTENT, queryIntent)
+        intent.configureChooserIntent(
+            resources,
+            component.hostUserHandle,
+            component.personalProfileUserHandle
+        )
 
-        val title = getString(R.string.media_projection_permission_app_selector_title)
-        intent.putExtra(Intent.EXTRA_TITLE, title)
         super.onCreate(bundle)
         controller.init()
     }
@@ -183,6 +182,13 @@
 
     override fun shouldShowContentPreview() = true
 
+    override fun shouldShowContentPreviewWhenEmpty(): Boolean = true
+
+    override fun createMyUserIdProvider(): MyUserIdProvider =
+        object : MyUserIdProvider() {
+            override fun getMyUserId(): Int = component.hostUserHandle.identifier
+        }
+
     override fun createContentPreviewView(parent: ViewGroup): ViewGroup =
         recentsViewController.createView(parent)
 
@@ -193,6 +199,34 @@
          * instance through activity result.
          */
         const val EXTRA_CAPTURE_REGION_RESULT_RECEIVER = "capture_region_result_receiver"
+
+        /** UID of the app that originally launched the media projection flow (host app user) */
+        const val EXTRA_HOST_APP_USER_HANDLE = "launched_from_user_handle"
         const val KEY_CAPTURE_TARGET = "capture_region"
+
+        /** Set up intent for the [ChooserActivity] */
+        private fun Intent.configureChooserIntent(
+            resources: Resources,
+            hostUserHandle: UserHandle,
+            personalProfileUserHandle: UserHandle
+        ) {
+            // Specify the query intent to show icons for all apps on the chooser screen
+            val queryIntent =
+                Intent(Intent.ACTION_MAIN).apply { addCategory(Intent.CATEGORY_LAUNCHER) }
+            putExtra(Intent.EXTRA_INTENT, queryIntent)
+
+            // Update the title of the chooser
+            val title = resources.getString(R.string.media_projection_permission_app_selector_title)
+            putExtra(Intent.EXTRA_TITLE, title)
+
+            // Select host app's profile tab by default
+            val selectedProfile =
+                if (hostUserHandle == personalProfileUserHandle) {
+                    PROFILE_PERSONAL
+                } else {
+                    PROFILE_WORK
+                }
+            putExtra(EXTRA_SELECTED_PROFILE, selectedProfile)
+        }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/media/MediaProjectionPermissionActivity.java b/packages/SystemUI/src/com/android/systemui/media/MediaProjectionPermissionActivity.java
index bfa67a8..d830fc4 100644
--- a/packages/SystemUI/src/com/android/systemui/media/MediaProjectionPermissionActivity.java
+++ b/packages/SystemUI/src/com/android/systemui/media/MediaProjectionPermissionActivity.java
@@ -22,6 +22,7 @@
 import static com.android.systemui.screenrecord.ScreenShareOptionKt.SINGLE_APP;
 
 import android.app.Activity;
+import android.app.ActivityManager;
 import android.app.AlertDialog;
 import android.content.DialogInterface;
 import android.content.Intent;
@@ -35,6 +36,7 @@
 import android.os.IBinder;
 import android.os.RemoteException;
 import android.os.ServiceManager;
+import android.os.UserHandle;
 import android.text.BidiFormatter;
 import android.text.SpannableString;
 import android.text.TextPaint;
@@ -208,8 +210,14 @@
                 final Intent intent = new Intent(this, MediaProjectionAppSelectorActivity.class);
                 intent.putExtra(MediaProjectionManager.EXTRA_MEDIA_PROJECTION,
                         projection.asBinder());
+                intent.putExtra(MediaProjectionAppSelectorActivity.EXTRA_HOST_APP_USER_HANDLE,
+                        UserHandle.getUserHandleForUid(getLaunchedFromUid()));
                 intent.setFlags(Intent.FLAG_ACTIVITY_FORWARD_RESULT);
-                startActivity(intent);
+
+                // Start activity from the current foreground user to avoid creating a separate
+                // SystemUI process without access to recent tasks because it won't have
+                // WM Shell running inside.
+                startActivityAsUser(intent, UserHandle.of(ActivityManager.getCurrentUser()));
             }
         } catch (RemoteException e) {
             Log.e(TAG, "Error granting projection permission", e);
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorComponent.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorComponent.kt
index 6c41caa..1d86343 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorComponent.kt
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorComponent.kt
@@ -19,9 +19,11 @@
 import android.app.Activity
 import android.content.ComponentName
 import android.content.Context
+import android.os.UserHandle
 import com.android.launcher3.icons.IconFactory
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.media.MediaProjectionAppSelectorActivity
+import com.android.systemui.media.MediaProjectionAppSelectorActivity.Companion.EXTRA_HOST_APP_USER_HANDLE
 import com.android.systemui.mediaprojection.appselector.data.ActivityTaskManagerThumbnailLoader
 import com.android.systemui.mediaprojection.appselector.data.AppIconLoader
 import com.android.systemui.mediaprojection.appselector.data.IconLoaderLibAppIconLoader
@@ -30,6 +32,8 @@
 import com.android.systemui.mediaprojection.appselector.data.ShellRecentTaskListProvider
 import com.android.systemui.mediaprojection.appselector.view.MediaProjectionRecentsViewController
 import com.android.systemui.mediaprojection.appselector.view.TaskPreviewSizeProvider
+import com.android.systemui.settings.UserTracker
+import com.android.systemui.shared.system.ActivityManagerWrapper
 import com.android.systemui.statusbar.phone.ConfigurationControllerImpl
 import com.android.systemui.statusbar.policy.ConfigurationController
 import dagger.Binds
@@ -39,6 +43,7 @@
 import dagger.Subcomponent
 import dagger.multibindings.ClassKey
 import dagger.multibindings.IntoMap
+import java.lang.IllegalArgumentException
 import javax.inject.Qualifier
 import javax.inject.Scope
 import kotlinx.coroutines.CoroutineScope
@@ -46,6 +51,12 @@
 
 @Qualifier @Retention(AnnotationRetention.BINARY) annotation class MediaProjectionAppSelector
 
+@Qualifier @Retention(AnnotationRetention.BINARY) annotation class HostUserHandle
+
+@Qualifier @Retention(AnnotationRetention.BINARY) annotation class PersonalProfile
+
+@Qualifier @Retention(AnnotationRetention.BINARY) annotation class WorkProfile
+
 @Retention(AnnotationRetention.RUNTIME) @Scope annotation class MediaProjectionAppSelectorScope
 
 @Module(subcomponents = [MediaProjectionAppSelectorComponent::class])
@@ -83,7 +94,7 @@
         @MediaProjectionAppSelector
         @MediaProjectionAppSelectorScope
         fun provideAppSelectorComponentName(context: Context): ComponentName =
-                ComponentName(context, MediaProjectionAppSelectorActivity::class.java)
+            ComponentName(context, MediaProjectionAppSelectorActivity::class.java)
 
         @Provides
         @MediaProjectionAppSelector
@@ -93,9 +104,32 @@
         ): ConfigurationController = ConfigurationControllerImpl(activity)
 
         @Provides
-        fun bindIconFactory(
-            context: Context
-        ): IconFactory = IconFactory.obtain(context)
+        @PersonalProfile
+        @MediaProjectionAppSelectorScope
+        fun personalUserHandle(activityManagerWrapper: ActivityManagerWrapper): UserHandle {
+            // Current foreground user is the 'personal' profile
+            return UserHandle.of(activityManagerWrapper.currentUserId)
+        }
+
+        @Provides
+        @WorkProfile
+        @MediaProjectionAppSelectorScope
+        fun workProfileUserHandle(userTracker: UserTracker): UserHandle? =
+            userTracker.userProfiles.find { it.isManagedProfile }?.userHandle
+
+        @Provides
+        @HostUserHandle
+        @MediaProjectionAppSelectorScope
+        fun hostUserHandle(activity: MediaProjectionAppSelectorActivity): UserHandle {
+            val extras =
+                activity.intent.extras
+                    ?: error("MediaProjectionAppSelectorActivity should be launched with extras")
+            return extras.getParcelable(EXTRA_HOST_APP_USER_HANDLE)
+                ?: error("MediaProjectionAppSelectorActivity should be provided with " +
+                        "$EXTRA_HOST_APP_USER_HANDLE extra")
+        }
+
+        @Provides fun bindIconFactory(context: Context): IconFactory = IconFactory.obtain(context)
 
         @Provides
         @MediaProjectionAppSelector
@@ -124,6 +158,8 @@
 
     val controller: MediaProjectionAppSelectorController
     val recentsViewController: MediaProjectionRecentsViewController
+    @get:HostUserHandle val hostUserHandle: UserHandle
+    @get:PersonalProfile val personalProfileUserHandle: UserHandle
 
     @MediaProjectionAppSelector val configurationController: ConfigurationController
 }
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorController.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorController.kt
index d744a40b..52c7ca3 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorController.kt
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorController.kt
@@ -17,24 +17,36 @@
 package com.android.systemui.mediaprojection.appselector
 
 import android.content.ComponentName
+import android.os.UserHandle
+import com.android.systemui.flags.FeatureFlags
+import com.android.systemui.flags.Flags
 import com.android.systemui.mediaprojection.appselector.data.RecentTask
 import com.android.systemui.mediaprojection.appselector.data.RecentTaskListProvider
+import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.launch
-import javax.inject.Inject
 
 @MediaProjectionAppSelectorScope
-class MediaProjectionAppSelectorController @Inject constructor(
+class MediaProjectionAppSelectorController
+@Inject
+constructor(
     private val recentTaskListProvider: RecentTaskListProvider,
     private val view: MediaProjectionAppSelectorView,
+    private val flags: FeatureFlags,
+    @HostUserHandle private val hostUserHandle: UserHandle,
     @MediaProjectionAppSelector private val scope: CoroutineScope,
     @MediaProjectionAppSelector private val appSelectorComponentName: ComponentName
 ) {
 
     fun init() {
         scope.launch {
-            val tasks = recentTaskListProvider.loadRecentTasks().sortTasks()
+            val recentTasks = recentTaskListProvider.loadRecentTasks()
+
+            val tasks = recentTasks
+                .filterDevicePolicyRestrictedTasks()
+                .sortedTasks()
+
             view.bind(tasks)
         }
     }
@@ -43,9 +55,20 @@
         scope.cancel()
     }
 
-    private fun List<RecentTask>.sortTasks(): List<RecentTask> =
-        sortedBy {
-            // Show normal tasks first and only then tasks with opened app selector
-            it.topActivityComponent == appSelectorComponentName
+    /**
+     * Removes all recent tasks that are different from the profile of the host app to avoid any
+     * cross-profile sharing
+     */
+    private fun List<RecentTask>.filterDevicePolicyRestrictedTasks(): List<RecentTask> =
+        if (flags.isEnabled(Flags.WM_ENABLE_PARTIAL_SCREEN_SHARING_ENTERPRISE_POLICIES)) {
+            // TODO(b/263950746): filter tasks based on the enterprise policies
+            this
+        } else {
+            filter { UserHandle.of(it.userId) == hostUserHandle }
         }
+
+    private fun List<RecentTask>.sortedTasks(): List<RecentTask> = sortedBy {
+        // Show normal tasks first and only then tasks with opened app selector
+        it.topActivityComponent == appSelectorComponentName
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTask.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTask.kt
index cd994b8..41e2286 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTask.kt
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTask.kt
@@ -17,11 +17,12 @@
 package com.android.systemui.mediaprojection.appselector.data
 
 import android.annotation.ColorInt
+import android.annotation.UserIdInt
 import android.content.ComponentName
 
 data class RecentTask(
     val taskId: Int,
-    val userId: Int,
+    @UserIdInt val userId: Int,
     val topActivityComponent: ComponentName?,
     val baseIntentComponent: ComponentName?,
     @ColorInt val colorBackground: Int?
diff --git a/packages/SystemUI/src/com/android/systemui/screenrecord/ScreenRecordPermissionDialog.kt b/packages/SystemUI/src/com/android/systemui/screenrecord/ScreenRecordPermissionDialog.kt
index 44b18ec..68e3dcd 100644
--- a/packages/SystemUI/src/com/android/systemui/screenrecord/ScreenRecordPermissionDialog.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenrecord/ScreenRecordPermissionDialog.kt
@@ -23,6 +23,7 @@
 import android.os.Handler
 import android.os.Looper
 import android.os.ResultReceiver
+import android.os.UserHandle
 import android.view.View
 import android.view.View.GONE
 import android.view.View.VISIBLE
@@ -77,6 +78,14 @@
                     MediaProjectionAppSelectorActivity.EXTRA_CAPTURE_REGION_RESULT_RECEIVER,
                     CaptureTargetResultReceiver()
                 )
+
+                // Send SystemUI's user handle as the host app user handle because SystemUI
+                // is the 'host app' (the app that receives screen capture data)
+                intent.putExtra(
+                    MediaProjectionAppSelectorActivity.EXTRA_HOST_APP_USER_HANDLE,
+                    UserHandle.of(UserHandle.myUserId())
+                )
+
                 val animationController = dialogLaunchAnimator.createActivityLaunchController(v!!)
                 if (animationController == null) {
                     dismiss()
diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorControllerTest.kt
index 19d2d33..1042ea7 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorControllerTest.kt
@@ -1,12 +1,16 @@
 package com.android.systemui.mediaprojection.appselector
 
 import android.content.ComponentName
+import android.os.UserHandle
 import android.testing.AndroidTestingRunner
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.flags.FeatureFlags
+import com.android.systemui.flags.Flags
 import com.android.systemui.mediaprojection.appselector.data.RecentTask
 import com.android.systemui.mediaprojection.appselector.data.RecentTaskListProvider
 import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.whenever
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import org.junit.Test
@@ -21,11 +25,17 @@
     private val scope = CoroutineScope(Dispatchers.Unconfined)
     private val appSelectorComponentName = ComponentName("com.test", "AppSelector")
 
+    private val hostUserHandle = UserHandle.of(123)
+    private val otherUserHandle = UserHandle.of(456)
+
     private val view: MediaProjectionAppSelectorView = mock()
+    private val featureFlags: FeatureFlags = mock()
 
     private val controller = MediaProjectionAppSelectorController(
         taskListProvider,
         view,
+        featureFlags,
+        hostUserHandle,
         scope,
         appSelectorComponentName
     )
@@ -98,15 +108,72 @@
         )
     }
 
+    @Test
+    fun initRecentTasksWithAppSelectorTasks_enterprisePoliciesDisabled_bindsOnlyTasksWithHostProfile() {
+        givenEnterprisePoliciesFeatureFlag(enabled = false)
+
+        val tasks = listOf(
+            createRecentTask(taskId = 1, userId = hostUserHandle.identifier),
+            createRecentTask(taskId = 2, userId = otherUserHandle.identifier),
+            createRecentTask(taskId = 3, userId = hostUserHandle.identifier),
+            createRecentTask(taskId = 4, userId = otherUserHandle.identifier),
+            createRecentTask(taskId = 5, userId = hostUserHandle.identifier),
+        )
+        taskListProvider.tasks = tasks
+
+        controller.init()
+
+        verify(view).bind(
+            listOf(
+                createRecentTask(taskId = 1, userId = hostUserHandle.identifier),
+                createRecentTask(taskId = 3, userId = hostUserHandle.identifier),
+                createRecentTask(taskId = 5, userId = hostUserHandle.identifier),
+            )
+        )
+    }
+
+    @Test
+    fun initRecentTasksWithAppSelectorTasks_enterprisePoliciesEnabled_bindsAllTasks() {
+        givenEnterprisePoliciesFeatureFlag(enabled = true)
+
+        val tasks = listOf(
+            createRecentTask(taskId = 1, userId = hostUserHandle.identifier),
+            createRecentTask(taskId = 2, userId = otherUserHandle.identifier),
+            createRecentTask(taskId = 3, userId = hostUserHandle.identifier),
+            createRecentTask(taskId = 4, userId = otherUserHandle.identifier),
+            createRecentTask(taskId = 5, userId = hostUserHandle.identifier),
+        )
+        taskListProvider.tasks = tasks
+
+        controller.init()
+
+        // TODO(b/233348916) should filter depending on the policies
+        verify(view).bind(
+            listOf(
+                createRecentTask(taskId = 1, userId = hostUserHandle.identifier),
+                createRecentTask(taskId = 2, userId = otherUserHandle.identifier),
+                createRecentTask(taskId = 3, userId = hostUserHandle.identifier),
+                createRecentTask(taskId = 4, userId = otherUserHandle.identifier),
+                createRecentTask(taskId = 5, userId = hostUserHandle.identifier),
+            )
+        )
+    }
+
+    private fun givenEnterprisePoliciesFeatureFlag(enabled: Boolean) {
+        whenever(featureFlags.isEnabled(Flags.WM_ENABLE_PARTIAL_SCREEN_SHARING_ENTERPRISE_POLICIES))
+            .thenReturn(enabled)
+    }
+
     private fun createRecentTask(
         taskId: Int,
-        topActivityComponent: ComponentName? = null
+        topActivityComponent: ComponentName? = null,
+        userId: Int = hostUserHandle.identifier
     ): RecentTask {
         return RecentTask(
             taskId = taskId,
             topActivityComponent = topActivityComponent,
             baseIntentComponent = ComponentName("com", "Test"),
-            userId = 0,
+            userId = userId,
             colorBackground = 0
         )
     }