From 079504b35b2f1376bf6597f4229e4c3e3dad3417 Mon Sep 17 00:00:00 2001
From: Elis Elliott <eliselliott@google.com>
Date: Fri, 17 Feb 2023 16:55:35 +0000
Subject: [PATCH] Use new setUserRestriction API in policy engine.

Test: btest a.d.c.TelephonyTest -c yes
Test: btest a.d.c.SafeBootTest -c yes
Test: btest a.d.c.InstallUnknownSourcesTest -c yes
Test: btest a.d.c.TimeTest -c yes
Test: btest a.d.c.ModifyUsersTest -c yes
Test: btest a.d.c.DebuggingTest -c yes
Change-Id: I267e310af2ec8656f6b02daf9ad50a6b2f0ec141
---
 .../server/devicepolicy/DevicePolicyManagerService.java      | 5 +++--
 .../android/server/devicepolicy/PolicyEnforcerCallbacks.java | 5 ++---
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index 821a5f667b7b7..1f891ea4bc7fc 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -12823,11 +12823,12 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
         checkCanExecuteOrThrowUnsafe(DevicePolicyManager.OPERATION_SET_USER_RESTRICTION);
 
         if (useDevicePolicyEngine(caller, /* delegateScope= */ null)) {
+            int affectedUserId = parent ? getProfileParentId(userId) : userId;
             EnforcingAdmin admin = enforcePermissionForUserRestriction(
                     who,
                     key,
                     caller.getPackageName(),
-                    userId);
+                    affectedUserId);
             PolicyDefinition<Boolean> policyDefinition =
                     PolicyDefinition.getPolicyDefinitionForUserRestriction(key);
             if (enabledFromThisOwner) {
@@ -12835,7 +12836,7 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
                         policyDefinition,
                         admin,
                         new BooleanPolicyValue(true),
-                        parent ? getProfileParentId(userId) : userId);
+                        affectedUserId);
             } else {
                 // Remove any local and global policy that was set by the admin
                 if (!policyDefinition.isLocalOnlyPolicy()) {
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/PolicyEnforcerCallbacks.java b/services/devicepolicy/java/com/android/server/devicepolicy/PolicyEnforcerCallbacks.java
index daa8a26e00d1c..fd91249fac627 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/PolicyEnforcerCallbacks.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/PolicyEnforcerCallbacks.java
@@ -215,10 +215,9 @@ final class PolicyEnforcerCallbacks {
             }
             UserRestrictionPolicyKey parsedKey =
                     (UserRestrictionPolicyKey) policyKey;
-            // TODO: call into new UserManager API when merged
             UserManagerInternal userManager = LocalServices.getService(UserManagerInternal.class);
-//            userManager.setUserRestriction(
-//                    userId, parsedKey.getRestriction(), enabled != null && enabled);
+            userManager.setUserRestriction(
+                    userId, parsedKey.getRestriction(), enabled != null && enabled);
             return true;
         }));
     }