1 files changed, 23 insertions, 16 deletions
diff --git a/PermissionController/src/com/android/permissioncontroller/role/utils/RoleUiBehaviorUtils.java b/PermissionController/src/com/android/permissioncontroller/role/utils/RoleUiBehaviorUtils.java
index 40088c7e0..6081695b5 100644
--- a/PermissionController/src/com/android/permissioncontroller/role/utils/RoleUiBehaviorUtils.java
+++ b/PermissionController/src/com/android/permissioncontroller/role/utils/RoleUiBehaviorUtils.java
@@ -26,10 +26,11 @@ import android.util.Log;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
-import androidx.preference.Preference;
 
 import com.android.modules.utils.build.SdkLevel;
-import com.android.permissioncontroller.role.ui.TwoTargetPreference;
+import com.android.permissioncontroller.role.ui.RoleApplicationPreference;
+import com.android.permissioncontroller.role.ui.RolePreference;
+import com.android.permissioncontroller.role.ui.UserRestrictionAwarePreference;
 import com.android.permissioncontroller.role.ui.behavior.RoleUiBehavior;
 import com.android.role.controller.model.Role;
 
@@ -102,13 +103,9 @@ public final class RoleUiBehaviorUtils {
      * @see RoleUiBehavior#preparePreferenceAsUser
      */
     public static void preparePreferenceAsUser(@NonNull Role role,
-            @NonNull TwoTargetPreference preference, @NonNull UserHandle user,
+            @NonNull RolePreference preference, @NonNull UserHandle user,
             @NonNull Context context) {
-        if (SdkLevel.isAtLeastU() && role.isExclusive()) {
-            final UserManager userManager = context.getSystemService(UserManager.class);
-            preference.setEnabled(!userManager.hasUserRestrictionForUser(
-                    UserManager.DISALLOW_CONFIG_DEFAULT_APPS, user));
-        }
+        prepareUserRestrictionAwarePreferenceAsUser(role, preference, user, context);
 
         RoleUiBehavior uiBehavior = getUiBehavior(role);
         if (uiBehavior == null) {
@@ -134,22 +131,32 @@ public final class RoleUiBehaviorUtils {
      * @see RoleUiBehavior#prepareApplicationPreferenceAsUser
      */
     public static void prepareApplicationPreferenceAsUser(@NonNull Role role,
-            @NonNull Preference preference, @NonNull ApplicationInfo applicationInfo,
-            @NonNull UserHandle user, @NonNull Context context) {
-        if (SdkLevel.isAtLeastU() && role.isExclusive()) {
-            final UserManager userManager = context.getSystemService(UserManager.class);
-            preference.setEnabled(!userManager.hasUserRestrictionForUser(
-                    UserManager.DISALLOW_CONFIG_DEFAULT_APPS, user));
-        }
+            @NonNull RoleApplicationPreference preference,
+            @NonNull ApplicationInfo applicationInfo, @NonNull UserHandle user,
+            @NonNull Context context) {
+        prepareUserRestrictionAwarePreferenceAsUser(role, preference, user, context);
 
         RoleUiBehavior uiBehavior = getUiBehavior(role);
         if (uiBehavior == null) {
             return;
         }
-        uiBehavior.prepareApplicationPreferenceAsUser(role, preference, applicationInfo, user,
+        uiBehavior.prepareApplicationPreferenceAsUser(
+                role, preference.asTwoStatePreference(), applicationInfo, user,
                 context);
     }
 
+    private static void prepareUserRestrictionAwarePreferenceAsUser(@NonNull Role role,
+            @NonNull UserRestrictionAwarePreference preference, @NonNull UserHandle user,
+            @NonNull Context context) {
+        if (SdkLevel.isAtLeastU() && role.isExclusive()) {
+            UserManager userManager = context.getSystemService(UserManager.class);
+            boolean hasDisallowConfigDefaultApps = userManager.hasUserRestrictionForUser(
+                    UserManager.DISALLOW_CONFIG_DEFAULT_APPS, user);
+            preference.setUserRestriction(hasDisallowConfigDefaultApps
+                    ? UserManager.DISALLOW_CONFIG_DEFAULT_APPS : null);
+        }
+    }
+
     /**
      * @see RoleUiBehavior#getConfirmationMessage
      */