From 6aee9d4d59c9dcb9cb88700acd9b29c0addbfd79 Mon Sep 17 00:00:00 2001
From: Jonathan Scott <scottjonathan@google.com>
Date: Tue, 13 Dec 2022 13:56:22 +0000
Subject: [PATCH] Correctly identify the COPE user for screenshot policy.

Test: btest "android.devicepolicy.cts.ScreenCaptureDisabledTest#setScreenCaptureDisabled_true_works[IncludeRunOnParentOfOrganizationOwnedProfileOwnerUsingParentInstance]" -sy
Bug: 258561494
Change-Id: I9e940e5ba7faadbcfdb71461400886fa8a11ca69
---
 .../DevicePolicyManagerService.java           | 31 ++++++-------------
 1 file changed, 9 insertions(+), 22 deletions(-)

diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index 7c697058c394f..667f41f0ad4db 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -2588,7 +2588,6 @@ public class DevicePolicyManagerService extends BaseIDevicePolicyManager {
         ensureLocked();
         // Try to find an admin which can use reqPolicy
         final ComponentName poAdminComponent = mOwners.getProfileOwnerComponent(userId);
-        final ComponentName doAdminComponent = mOwners.getDeviceOwnerComponent();
 
         if (poAdminComponent != null) {
             return getProfileOwnerLocked(userId);
@@ -7762,9 +7761,15 @@ public class DevicePolicyManagerService extends BaseIDevicePolicyManager {
     // be disabled device-wide.
     private void pushScreenCapturePolicy(int adminUserId) {
         // Update screen capture device-wide if disabled by the DO or COPE PO on the parent profile.
-        ActiveAdmin admin =
-                getDeviceOwnerOrProfileOwnerOfOrganizationOwnedDeviceParentLocked(
-                        UserHandle.USER_SYSTEM);
+        // TODO(b/261999445): remove
+        ActiveAdmin admin;
+        if (isHeadlessFlagEnabled()) {
+            admin = getDeviceOwnerOrProfileOwnerOfOrganizationOwnedDeviceParentLocked(
+                    mUserManagerInternal.getProfileParentId(adminUserId));
+        } else {
+            admin = getDeviceOwnerOrProfileOwnerOfOrganizationOwnedDeviceParentLocked(
+                    UserHandle.USER_SYSTEM);
+        }
         if (admin != null && admin.disableScreenCapture) {
             setScreenCaptureDisabled(UserHandle.USER_ALL);
         } else {
@@ -8817,10 +8822,6 @@ public class DevicePolicyManagerService extends BaseIDevicePolicyManager {
         return admin;
     }
 
-    /**
-     * @deprecated Use the version which does not take a user id.
-     */
-    @Deprecated
     ActiveAdmin getDeviceOwnerOrProfileOwnerOfOrganizationOwnedDeviceParentLocked(int userId) {
         ensureLocked();
         ActiveAdmin admin = getDeviceOwnerAdminLocked();
@@ -8831,16 +8832,6 @@ public class DevicePolicyManagerService extends BaseIDevicePolicyManager {
         return admin != null ? admin.getParentActiveAdmin() : null;
     }
 
-    ActiveAdmin getDeviceOwnerOrProfileOwnerOfOrganizationOwnedDeviceParentLocked() {
-        ensureLocked();
-        ActiveAdmin admin = getDeviceOwnerAdminLocked();
-        if (admin != null) {
-            return admin;
-        }
-        admin = getProfileOwnerOfOrganizationOwnedDeviceLocked();
-        return admin != null ? admin.getParentActiveAdmin() : null;
-    }
-
     @Override
     public void clearDeviceOwner(String packageName) {
         Objects.requireNonNull(packageName, "packageName is null");
@@ -9374,10 +9365,6 @@ public class DevicePolicyManagerService extends BaseIDevicePolicyManager {
         return admin;
     }
 
-    /**
-     * @deprecated use the version which does not take a user id.
-     */
-    @Deprecated
     @GuardedBy("getLockObject()")
     ActiveAdmin getProfileOwnerOfOrganizationOwnedDeviceLocked(int userHandle) {
         return mInjector.binderWithCleanCallingIdentity(() -> {