Add test APIs for policy size tracking

Also changed flag to a bug-fix flag

Bug: 324060156
Test: atest DevicePolicyManagerTest
Change-Id: Ie8ffeb5b50d39c211256a253f80f3d170d01f04a

18 files changed, 192 insertions, 28 deletions

diff --git a/core/api/test-current.txt b/core/api/test-current.txt
index f36aeab2cb17..5138b589b692 100644
--- a/core/api/test-current.txt
+++ b/core/api/test-current.txt
@@ -589,6 +589,7 @@ package android.app.admin {
     method @RequiresPermission(android.Manifest.permission.FORCE_DEVICE_POLICY_MANAGER_LOGS) public long forceNetworkLogs();
     method @RequiresPermission(android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS) public void forceRemoveActiveAdmin(@NonNull android.content.ComponentName, int);
     method @RequiresPermission(android.Manifest.permission.FORCE_DEVICE_POLICY_MANAGER_LOGS) public long forceSecurityLogs();
+    method @FlaggedApi("android.app.admin.flags.device_policy_size_tracking_internal_bug_fix_enabled") @RequiresPermission("android.permission.MANAGE_DEVICE_POLICY_STORAGE_LIMIT") public void forceSetMaxPolicyStorageLimit(int);
     method public void forceUpdateUserSetupComplete(int);
     method @NonNull public java.util.Set<java.lang.String> getDefaultCrossProfilePackages();
     method @Deprecated public int getDeviceOwnerType(@NonNull android.content.ComponentName);
@@ -599,6 +600,7 @@ package android.app.admin {
     method public long getLastSecurityLogRetrievalTime();
     method public java.util.List<java.lang.String> getOwnerInstalledCaCerts(@NonNull android.os.UserHandle);
     method @NonNull @RequiresPermission(android.Manifest.permission.MANAGE_DEVICE_ADMINS) public java.util.Set<java.lang.String> getPolicyExemptApps();
+    method @FlaggedApi("android.app.admin.flags.device_policy_size_tracking_internal_bug_fix_enabled") @RequiresPermission("android.permission.MANAGE_DEVICE_POLICY_STORAGE_LIMIT") public int getPolicySizeForAdmin(@NonNull android.app.admin.EnforcingAdmin);
     method public boolean isCurrentInputMethodSetByOwner();
     method public boolean isFactoryResetProtectionPolicySupported();
     method @RequiresPermission(anyOf={android.Manifest.permission.MANAGE_USERS, android.Manifest.permission.INTERACT_ACROSS_USERS}) public boolean isNewUserDisclaimerAcknowledged();
@@ -667,6 +669,10 @@ package android.app.admin {
     field @NonNull public static final android.app.admin.DpcAuthority DPC_AUTHORITY;
   }
 
+  public final class EnforcingAdmin implements android.os.Parcelable {
+    ctor @FlaggedApi("android.app.admin.flags.device_policy_size_tracking_internal_bug_fix_enabled") public EnforcingAdmin(@NonNull String, @NonNull android.app.admin.Authority, @NonNull android.os.UserHandle, @Nullable android.content.ComponentName);
+  }
+
   public final class FlagUnion extends android.app.admin.ResolutionMechanism<java.lang.Integer> {
     method public int describeContents();
     method public void writeToParcel(@NonNull android.os.Parcel, int);
diff --git a/core/java/android/app/admin/AccountTypePolicyKey.java b/core/java/android/app/admin/AccountTypePolicyKey.java
index 51f313755e59..02e492bb06aa 100644
--- a/core/java/android/app/admin/AccountTypePolicyKey.java
+++ b/core/java/android/app/admin/AccountTypePolicyKey.java
@@ -54,7 +54,7 @@ public final class AccountTypePolicyKey extends PolicyKey {
     @TestApi
     public AccountTypePolicyKey(@NonNull String key, @NonNull String accountType) {
         super(key);
-        if (Flags.devicePolicySizeTrackingInternalEnabled()) {
+        if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
             PolicySizeVerifier.enforceMaxStringLength(accountType, "accountType");
         }
         mAccountType = Objects.requireNonNull((accountType));
diff --git a/core/java/android/app/admin/BundlePolicyValue.java b/core/java/android/app/admin/BundlePolicyValue.java
index cb5e9861141d..c993671f4fc1 100644
--- a/core/java/android/app/admin/BundlePolicyValue.java
+++ b/core/java/android/app/admin/BundlePolicyValue.java
@@ -31,7 +31,7 @@ public final class BundlePolicyValue extends PolicyValue<Bundle> {
 
     public BundlePolicyValue(Bundle value) {
         super(value);
-        if (Flags.devicePolicySizeTrackingInternalEnabled()) {
+        if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
             PolicySizeVerifier.enforceMaxBundleFieldsLength(value);
         }
     }
diff --git a/core/java/android/app/admin/ComponentNamePolicyValue.java b/core/java/android/app/admin/ComponentNamePolicyValue.java
index a957dbf132bb..a7a2f7d27e0d 100644
--- a/core/java/android/app/admin/ComponentNamePolicyValue.java
+++ b/core/java/android/app/admin/ComponentNamePolicyValue.java
@@ -31,7 +31,7 @@ public final class ComponentNamePolicyValue extends PolicyValue<ComponentName> {
 
     public ComponentNamePolicyValue(@NonNull ComponentName value) {
         super(value);
-        if (Flags.devicePolicySizeTrackingInternalEnabled()) {
+        if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
             PolicySizeVerifier.enforceMaxComponentNameLength(value);
         }
     }
diff --git a/core/java/android/app/admin/DevicePolicyManager.java b/core/java/android/app/admin/DevicePolicyManager.java
index ea6f45e8e201..17b3cf62d961 100644
--- a/core/java/android/app/admin/DevicePolicyManager.java
+++ b/core/java/android/app/admin/DevicePolicyManager.java
@@ -54,6 +54,7 @@ import static android.Manifest.permission.REQUEST_PASSWORD_COMPLEXITY;
 import static android.Manifest.permission.SET_TIME;
 import static android.Manifest.permission.SET_TIME_ZONE;
 import static android.app.admin.DeviceAdminInfo.HEADLESS_DEVICE_OWNER_MODE_UNSUPPORTED;
+import static android.app.admin.flags.Flags.FLAG_DEVICE_POLICY_SIZE_TRACKING_INTERNAL_BUG_FIX_ENABLED;
 import static android.app.admin.flags.Flags.FLAG_DEVICE_THEFT_API_ENABLED;
 import static android.app.admin.flags.Flags.FLAG_ESIM_MANAGEMENT_ENABLED;
 import static android.app.admin.flags.Flags.FLAG_DEVICE_POLICY_SIZE_TRACKING_ENABLED;
@@ -17560,6 +17561,48 @@ public class DevicePolicyManager {
     }
 
     /**
+     * Force sets the maximum storage size allowed for policies associated with an admin regardless
+     * of the default value set in the system, unlike {@link #setMaxPolicyStorageLimit} which can
+     * only set it to a value higher than the default value set by the system.Setting a limit of -1
+     * effectively removes any storage restrictions.
+     *
+     * @param storageLimit Maximum storage allowed in bytes. Use -1 to disable limits.
+     *
+     * @hide
+     */
+    @TestApi
+    @RequiresPermission(permission.MANAGE_DEVICE_POLICY_STORAGE_LIMIT)
+    @FlaggedApi(FLAG_DEVICE_POLICY_SIZE_TRACKING_INTERNAL_BUG_FIX_ENABLED)
+    public void forceSetMaxPolicyStorageLimit(int storageLimit) {
+        if (mService != null) {
+            try {
+                mService.forceSetMaxPolicyStorageLimit(mContext.getPackageName(), storageLimit);
+            } catch (RemoteException e) {
+                throw e.rethrowFromSystemServer();
+            }
+        }
+    }
+
+    /**
+     * Retrieves the size of the current policies set by the {@code admin}.
+     *
+     * @hide
+     */
+    @TestApi
+    @RequiresPermission(permission.MANAGE_DEVICE_POLICY_STORAGE_LIMIT)
+    @FlaggedApi(FLAG_DEVICE_POLICY_SIZE_TRACKING_INTERNAL_BUG_FIX_ENABLED)
+    public int getPolicySizeForAdmin(@NonNull EnforcingAdmin admin) {
+        if (mService != null) {
+            try {
+                return mService.getPolicySizeForAdmin(mContext.getPackageName(), admin);
+            } catch (RemoteException e) {
+                throw e.rethrowFromSystemServer();
+            }
+        }
+        return -1;
+    }
+
+    /**
      * @return The headless device owner mode for the current set DO, returns
      * {@link DeviceAdminInfo#HEADLESS_DEVICE_OWNER_MODE_UNSUPPORTED} if no DO is set.
      *
diff --git a/core/java/android/app/admin/EnforcingAdmin.java b/core/java/android/app/admin/EnforcingAdmin.java
index 7c718f6651a2..f70a53f61671 100644
--- a/core/java/android/app/admin/EnforcingAdmin.java
+++ b/core/java/android/app/admin/EnforcingAdmin.java
@@ -16,9 +16,13 @@
 
 package android.app.admin;
 
+import static android.app.admin.flags.Flags.FLAG_DEVICE_POLICY_SIZE_TRACKING_INTERNAL_BUG_FIX_ENABLED;
+
+import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.SystemApi;
+import android.annotation.TestApi;
 import android.content.ComponentName;
 import android.os.Parcel;
 import android.os.Parcelable;
@@ -60,6 +64,8 @@ public final class EnforcingAdmin implements Parcelable {
      *
      * @hide
      */
+    @FlaggedApi(FLAG_DEVICE_POLICY_SIZE_TRACKING_INTERNAL_BUG_FIX_ENABLED)
+    @TestApi
     public EnforcingAdmin(
             @NonNull String packageName, @NonNull Authority authority,
             @NonNull UserHandle userHandle, @Nullable ComponentName componentName) {
@@ -101,6 +107,16 @@ public final class EnforcingAdmin implements Parcelable {
         return mUserHandle;
     }
 
+    /**
+     * Returns the {@link ComponentName} of the admin if applicable.
+     *
+     * @hide
+     */
+    @Nullable
+    public ComponentName getComponentName() {
+        return mComponentName;
+    }
+
     @Override
     public boolean equals(@Nullable Object o) {
         if (this == o) return true;
diff --git a/core/java/android/app/admin/IDevicePolicyManager.aidl b/core/java/android/app/admin/IDevicePolicyManager.aidl
index 2002326d76bd..d1837132e1a4 100644
--- a/core/java/android/app/admin/IDevicePolicyManager.aidl
+++ b/core/java/android/app/admin/IDevicePolicyManager.aidl
@@ -623,8 +623,10 @@ interface IDevicePolicyManager {
 
     int[] getSubscriptionIds(String callerPackageName);
 
-    void setMaxPolicyStorageLimit(String packageName, int storageLimit);
-    int getMaxPolicyStorageLimit(String packageName);
+    void setMaxPolicyStorageLimit(String callerPackageName, int storageLimit);
+    void forceSetMaxPolicyStorageLimit(String callerPackageName, int storageLimit);
+    int getMaxPolicyStorageLimit(String callerPackageName);
+    int getPolicySizeForAdmin(String callerPackageName, in EnforcingAdmin admin);
 
     int getHeadlessDeviceOwnerMode(String callerPackageName);
 }
diff --git a/core/java/android/app/admin/LockTaskPolicy.java b/core/java/android/app/admin/LockTaskPolicy.java
index a36ea0508a95..68b4ad84d81a 100644
--- a/core/java/android/app/admin/LockTaskPolicy.java
+++ b/core/java/android/app/admin/LockTaskPolicy.java
@@ -135,7 +135,7 @@ public final class LockTaskPolicy extends PolicyValue<LockTaskPolicy> {
     }
 
     private void setPackagesInternal(Set<String> packages) {
-        if (Flags.devicePolicySizeTrackingInternalEnabled()) {
+        if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
             for (String p : packages) {
                 PolicySizeVerifier.enforceMaxPackageNameLength(p);
             }
diff --git a/core/java/android/app/admin/PackagePermissionPolicyKey.java b/core/java/android/app/admin/PackagePermissionPolicyKey.java
index 389585f036db..1a04f6c908bc 100644
--- a/core/java/android/app/admin/PackagePermissionPolicyKey.java
+++ b/core/java/android/app/admin/PackagePermissionPolicyKey.java
@@ -59,7 +59,7 @@ public final class PackagePermissionPolicyKey extends PolicyKey {
     public PackagePermissionPolicyKey(@NonNull String identifier, @NonNull String packageName,
             @NonNull String permissionName) {
         super(identifier);
-        if (Flags.devicePolicySizeTrackingInternalEnabled()) {
+        if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
             PolicySizeVerifier.enforceMaxPackageNameLength(packageName);
             PolicySizeVerifier.enforceMaxStringLength(permissionName, "permissionName");
         }
diff --git a/core/java/android/app/admin/PackagePolicyKey.java b/core/java/android/app/admin/PackagePolicyKey.java
index 68dc797f6513..9e31a23aec91 100644
--- a/core/java/android/app/admin/PackagePolicyKey.java
+++ b/core/java/android/app/admin/PackagePolicyKey.java
@@ -55,7 +55,7 @@ public final class PackagePolicyKey extends PolicyKey {
     @TestApi
     public PackagePolicyKey(@NonNull String key, @NonNull String packageName) {
         super(key);
-        if (Flags.devicePolicySizeTrackingInternalEnabled()) {
+        if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
             PolicySizeVerifier.enforceMaxPackageNameLength(packageName);
         }
         mPackageName = Objects.requireNonNull((packageName));
diff --git a/core/java/android/app/admin/StringPolicyValue.java b/core/java/android/app/admin/StringPolicyValue.java
index 8995c0f20de8..6efe9ad0dbed 100644
--- a/core/java/android/app/admin/StringPolicyValue.java
+++ b/core/java/android/app/admin/StringPolicyValue.java
@@ -30,7 +30,7 @@ public final class StringPolicyValue extends PolicyValue<String> {
 
     public StringPolicyValue(@NonNull String value) {
         super(value);
-        if (Flags.devicePolicySizeTrackingInternalEnabled()) {
+        if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
             PolicySizeVerifier.enforceMaxStringLength(value, "policyValue");
         }
     }
diff --git a/core/java/android/app/admin/StringSetPolicyValue.java b/core/java/android/app/admin/StringSetPolicyValue.java
index f37dfee0f9dc..12b11f4ba687 100644
--- a/core/java/android/app/admin/StringSetPolicyValue.java
+++ b/core/java/android/app/admin/StringSetPolicyValue.java
@@ -32,7 +32,7 @@ public final class StringSetPolicyValue extends PolicyValue<Set<String>> {
 
     public StringSetPolicyValue(@NonNull Set<String> value) {
         super(value);
-        if (Flags.devicePolicySizeTrackingInternalEnabled()) {
+        if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
             for (String str : value) {
                 PolicySizeVerifier.enforceMaxStringLength(str, "policyValue");
             }
diff --git a/core/java/android/app/admin/UserRestrictionPolicyKey.java b/core/java/android/app/admin/UserRestrictionPolicyKey.java
index ee90ccd9417f..9054287cb7a0 100644
--- a/core/java/android/app/admin/UserRestrictionPolicyKey.java
+++ b/core/java/android/app/admin/UserRestrictionPolicyKey.java
@@ -45,7 +45,7 @@ public final class UserRestrictionPolicyKey extends PolicyKey {
     @TestApi
     public UserRestrictionPolicyKey(@NonNull String identifier, @NonNull String restriction) {
         super(identifier);
-        if (Flags.devicePolicySizeTrackingInternalEnabled()) {
+        if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
             PolicySizeVerifier.enforceMaxStringLength(restriction, "restriction");
         }
         mRestriction = Objects.requireNonNull(restriction);
diff --git a/core/java/android/app/admin/flags/flags.aconfig b/core/java/android/app/admin/flags/flags.aconfig
index 56fb4aa45fb3..0bc64b0dc309 100644
--- a/core/java/android/app/admin/flags/flags.aconfig
+++ b/core/java/android/app/admin/flags/flags.aconfig
@@ -28,6 +28,17 @@ flag {
 }
 
 flag {
+  name: "device_policy_size_tracking_internal_bug_fix_enabled"
+  namespace: "enterprise"
+  description: "Bug fix for tracking the total policy size and have a max threshold"
+  bug: "281543351"
+  metadata {
+      purpose: PURPOSE_BUGFIX
+  }
+}
+
+
+flag {
   name: "onboarding_bugreport_v2_enabled"
   is_exported: true
   namespace: "enterprise"
diff --git a/core/res/AndroidManifest.xml b/core/res/AndroidManifest.xml
index f74329903690..913b63e71d68 100644
--- a/core/res/AndroidManifest.xml
+++ b/core/res/AndroidManifest.xml
@@ -3889,6 +3889,13 @@
     <permission android:name="android.permission.MANAGE_DEVICE_POLICY_ACROSS_USERS_FULL"
                 android:protectionLevel="internal|role" />
 
+    <!-- Allows the holder to manage and retrieve max storage limit for admin policies. This
+        permission is only grantable on rooted devices.
+        @TestAPI
+        @hide -->
+    <permission android:name="android.permission.MANAGE_DEVICE_POLICY_STORAGE_LIMIT"
+                android:protectionLevel="internal" />
+
     <!-- Allows an application to access EnhancedConfirmationManager.
         @SystemApi
         @FlaggedApi("android.permission.flags.enhanced_confirmation_mode_apis_enabled")
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyEngine.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyEngine.java
index f39d0193f28a..065c14e3f208 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyEngine.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyEngine.java
@@ -103,7 +103,7 @@ final class DevicePolicyEngine {
                     UserManager.DISALLOW_CELLULAR_2G);
 
     //TODO(b/295504706) : Speak to security team to decide what to set Policy_Size_Limit
-    private static final int DEFAULT_POLICY_SIZE_LIMIT = -1;
+    static final int DEFAULT_POLICY_SIZE_LIMIT = -1;
 
     private final Context mContext;
     private final UserManager mUserManager;
@@ -225,7 +225,7 @@ final class DevicePolicyEngine {
 
         synchronized (mLock) {
             PolicyState<V> localPolicyState = getLocalPolicyStateLocked(policyDefinition, userId);
-            if (Flags.devicePolicySizeTrackingInternalEnabled()) {
+            if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
                 if (!handleAdminPolicySizeLimit(localPolicyState, enforcingAdmin, value,
                         policyDefinition, userId)) {
                     return;
@@ -350,7 +350,7 @@ final class DevicePolicyEngine {
             }
             PolicyState<V> localPolicyState = getLocalPolicyStateLocked(policyDefinition, userId);
 
-            if (Flags.devicePolicySizeTrackingInternalEnabled()) {
+            if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
                 decreasePolicySizeForAdmin(localPolicyState, enforcingAdmin);
             }
 
@@ -496,7 +496,7 @@ final class DevicePolicyEngine {
 
         synchronized (mLock) {
             PolicyState<V> globalPolicyState = getGlobalPolicyStateLocked(policyDefinition);
-            if (Flags.devicePolicySizeTrackingInternalEnabled()) {
+            if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
                 if (!handleAdminPolicySizeLimit(globalPolicyState, enforcingAdmin, value,
                         policyDefinition, UserHandle.USER_ALL)) {
                     return;
@@ -568,7 +568,7 @@ final class DevicePolicyEngine {
         synchronized (mLock) {
             PolicyState<V> policyState = getGlobalPolicyStateLocked(policyDefinition);
 
-            if (Flags.devicePolicySizeTrackingInternalEnabled()) {
+            if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
                 decreasePolicySizeForAdmin(policyState, enforcingAdmin);
             }
 
@@ -1598,6 +1598,7 @@ final class DevicePolicyEngine {
             existingPolicySize = sizeOf(policyState.getPoliciesSetByAdmins().get(admin));
         }
         int policySize = sizeOf(value);
+
         // Policy size limit is disabled if mPolicySizeLimit is -1.
         if (mPolicySizeLimit == -1
                 || currentAdminPoliciesSize + policySize - existingPolicySize < mPolicySizeLimit) {
@@ -1657,10 +1658,6 @@ final class DevicePolicyEngine {
      * the limitation.
      */
     void setMaxPolicyStorageLimit(int storageLimit) {
-        if (storageLimit < DEFAULT_POLICY_SIZE_LIMIT && storageLimit != -1) {
-            throw new IllegalArgumentException("Can't set a size limit less than the minimum "
-                    + "allowed size.");
-        }
         mPolicySizeLimit = storageLimit;
     }
 
@@ -1672,6 +1669,15 @@ final class DevicePolicyEngine {
         return mPolicySizeLimit;
     }
 
+    int getPolicySizeForAdmin(EnforcingAdmin admin) {
+        if (mAdminPolicySize.contains(admin.getUserId())
+                && mAdminPolicySize.get(
+                admin.getUserId()).containsKey(admin)) {
+            return mAdminPolicySize.get(admin.getUserId()).get(admin);
+        }
+        return 0;
+    }
+
     public void dump(IndentingPrintWriter pw) {
         synchronized (mLock) {
             pw.println("Local Policies: ");
@@ -1906,7 +1912,7 @@ final class DevicePolicyEngine {
 
         private void writeEnforcingAdminSizeInner(TypedXmlSerializer serializer)
                 throws IOException {
-            if (Flags.devicePolicySizeTrackingInternalEnabled()) {
+            if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
                 if (mAdminPolicySize != null) {
                     for (int i = 0; i < mAdminPolicySize.size(); i++) {
                         int userId = mAdminPolicySize.keyAt(i);
@@ -1930,7 +1936,7 @@ final class DevicePolicyEngine {
 
         private void writeMaxPolicySizeInner(TypedXmlSerializer serializer)
                 throws IOException {
-            if (!Flags.devicePolicySizeTrackingInternalEnabled()) {
+            if (!Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
                 return;
             }
             serializer.startTag(/* namespace= */ null, TAG_MAX_POLICY_SIZE_LIMIT);
@@ -2095,7 +2101,7 @@ final class DevicePolicyEngine {
 
         private void readMaxPolicySizeInner(TypedXmlPullParser parser)
                 throws XmlPullParserException, IOException {
-            if (!Flags.devicePolicySizeTrackingInternalEnabled()) {
+            if (!Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
                 return;
             }
             mPolicySizeLimit = parser.getAttributeInt(/* namespace= */ null, ATTR_POLICY_SUM_SIZE);
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index 3dd7b5480da1..2bae414411dd 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -88,6 +88,7 @@ import static android.Manifest.permission.MANAGE_DEVICE_POLICY_WALLPAPER;
 import static android.Manifest.permission.MANAGE_DEVICE_POLICY_WIFI;
 import static android.Manifest.permission.MANAGE_DEVICE_POLICY_WINDOWS;
 import static android.Manifest.permission.MANAGE_DEVICE_POLICY_WIPE_DATA;
+import static android.Manifest.permission.MANAGE_DEVICE_POLICY_STORAGE_LIMIT;
 import static android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS;
 import static android.Manifest.permission.MASTER_CLEAR;
 import static android.Manifest.permission.NOTIFY_PENDING_SYSTEM_UPDATE;
@@ -268,6 +269,7 @@ import static com.android.internal.widget.LockPatternUtils.StrongAuthTracker.SOM
 import static com.android.internal.widget.LockPatternUtils.StrongAuthTracker.STRONG_AUTH_REQUIRED_AFTER_DPM_LOCK_NOW;
 import static com.android.server.SystemTimeZone.TIME_ZONE_CONFIDENCE_HIGH;
 import static com.android.server.am.ActivityManagerService.STOCK_PM_FLAGS;
+import static com.android.server.devicepolicy.DevicePolicyEngine.DEFAULT_POLICY_SIZE_LIMIT;
 import static com.android.server.devicepolicy.TransferOwnershipMetadataManager.ADMIN_TYPE_DEVICE_OWNER;
 import static com.android.server.devicepolicy.TransferOwnershipMetadataManager.ADMIN_TYPE_PROFILE_OWNER;
 import static com.android.server.pm.PackageManagerService.PLATFORM_PACKAGE_NAME;
@@ -12138,7 +12140,7 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
         }
 
         if (packageList != null) {
-            if (!Flags.devicePolicySizeTrackingInternalEnabled()) {
+            if (!Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
                 for (String pkg : packageList) {
                     PolicySizeVerifier.enforceMaxPackageNameLength(pkg);
                 }
@@ -13913,7 +13915,7 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
             return;
         }
 
-        if (!Flags.devicePolicySizeTrackingInternalEnabled()) {
+        if (!Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
             PolicySizeVerifier.enforceMaxStringLength(accountType, "account type");
         }
 
@@ -14527,7 +14529,7 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
     public void setLockTaskPackages(ComponentName who, String callerPackageName, String[] packages)
             throws SecurityException {
         Objects.requireNonNull(packages, "packages is null");
-        if (!Flags.devicePolicySizeTrackingInternalEnabled()) {
+        if (!Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
             for (String pkg : packages) {
                 PolicySizeVerifier.enforceMaxPackageNameLength(pkg);
             }
@@ -24536,19 +24538,23 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
 
     @Override
     public void setMaxPolicyStorageLimit(String callerPackageName, int storageLimit) {
-        if (!Flags.devicePolicySizeTrackingInternalEnabled()) {
+        if (!Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
             return;
         }
         CallerIdentity caller = getCallerIdentity(callerPackageName);
         enforcePermission(MANAGE_PROFILE_AND_DEVICE_OWNERS, caller.getPackageName(),
                 caller.getUserId());
 
+        if (storageLimit < DEFAULT_POLICY_SIZE_LIMIT && storageLimit != -1) {
+            throw new IllegalArgumentException("Can't set a size limit less than the minimum "
+                    + "allowed size.");
+        }
         mDevicePolicyEngine.setMaxPolicyStorageLimit(storageLimit);
     }
 
     @Override
     public int getMaxPolicyStorageLimit(String callerPackageName) {
-        if (!Flags.devicePolicySizeTrackingInternalEnabled()) {
+        if (!Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
             return -1;
         }
         CallerIdentity caller = getCallerIdentity(callerPackageName);
@@ -24559,6 +24565,32 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
     }
 
     @Override
+    public void forceSetMaxPolicyStorageLimit(String callerPackageName, int storageLimit) {
+        if (!Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
+            return;
+        }
+        CallerIdentity caller = getCallerIdentity(callerPackageName);
+        enforcePermission(MANAGE_DEVICE_POLICY_STORAGE_LIMIT, caller.getPackageName(),
+                caller.getUserId());
+
+        mDevicePolicyEngine.setMaxPolicyStorageLimit(storageLimit);
+    }
+
+    @Override
+    public int getPolicySizeForAdmin(
+            String callerPackageName, android.app.admin.EnforcingAdmin admin) {
+        if (!Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
+            return -1;
+        }
+        CallerIdentity caller = getCallerIdentity(callerPackageName);
+        enforcePermission(MANAGE_DEVICE_POLICY_STORAGE_LIMIT, caller.getPackageName(),
+                caller.getUserId());
+
+        return mDevicePolicyEngine.getPolicySizeForAdmin(
+                EnforcingAdmin.createEnforcingAdmin(admin));
+    }
+
+    @Override
     public int getHeadlessDeviceOwnerMode(String callerPackageName) {
         final CallerIdentity caller = getCallerIdentity(callerPackageName);
         enforcePermission(MANAGE_PROFILE_AND_DEVICE_OWNERS, caller.getPackageName(),
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/EnforcingAdmin.java b/services/devicepolicy/java/com/android/server/devicepolicy/EnforcingAdmin.java
index d234dee3c8f7..02590f97ab6b 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/EnforcingAdmin.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/EnforcingAdmin.java
@@ -21,6 +21,7 @@ import android.annotation.Nullable;
 import android.app.admin.Authority;
 import android.app.admin.DeviceAdminAuthority;
 import android.app.admin.DpcAuthority;
+import android.app.admin.PackagePermissionPolicyKey;
 import android.app.admin.RoleAuthority;
 import android.app.admin.UnknownAuthority;
 import android.content.ComponentName;
@@ -105,6 +106,32 @@ final class EnforcingAdmin {
                 userId, activeAdmin);
     }
 
+    static EnforcingAdmin createEnforcingAdmin(android.app.admin.EnforcingAdmin admin) {
+        Objects.requireNonNull(admin);
+        Authority authority = admin.getAuthority();
+        Set<String> internalAuthorities = new HashSet<>();
+        if (DpcAuthority.DPC_AUTHORITY.equals(authority)) {
+            return new EnforcingAdmin(
+                    admin.getPackageName(), admin.getComponentName(),
+                    Set.of(DPC_AUTHORITY), admin.getUserHandle().getIdentifier(),
+                    /* activeAdmin = */ null);
+        } else if (DeviceAdminAuthority.DEVICE_ADMIN_AUTHORITY.equals(authority)) {
+            return new EnforcingAdmin(
+                    admin.getPackageName(), admin.getComponentName(),
+                    Set.of(DEVICE_ADMIN_AUTHORITY), admin.getUserHandle().getIdentifier(),
+                    /* activeAdmin = */ null);
+        } else if (authority instanceof RoleAuthority roleAuthority) {
+            return new EnforcingAdmin(
+                    admin.getPackageName(), admin.getComponentName(),
+                    Set.of(DEVICE_ADMIN_AUTHORITY), admin.getUserHandle().getIdentifier(),
+                    /* activeAdmin = */ null,
+                    /* isRoleAuthority = */ true);
+        }
+        return new EnforcingAdmin(admin.getPackageName(), admin.getComponentName(),
+                Set.of(), admin.getUserHandle().getIdentifier(),
+                /* activeAdmin = */ null);
+    }
+
     static String getRoleAuthorityOf(String roleName) {
         return ROLE_AUTHORITY_PREFIX + roleName;
     }
@@ -154,6 +181,20 @@ final class EnforcingAdmin {
         mActiveAdmin = activeAdmin;
     }
 
+    private EnforcingAdmin(
+            String packageName, @Nullable ComponentName componentName, Set<String> authorities,
+            int userId, @Nullable ActiveAdmin activeAdmin, boolean isRoleAuthority) {
+        Objects.requireNonNull(packageName);
+        Objects.requireNonNull(authorities);
+
+        mIsRoleAuthority = isRoleAuthority;
+        mPackageName = packageName;
+        mComponentName = componentName;
+        mAuthorities = new HashSet<>(authorities);
+        mUserId = userId;
+        mActiveAdmin = activeAdmin;
+    }
+
     private static Set<String> getRoleAuthoritiesOrDefault(String packageName, int userId) {
         Set<String> roles = getRoles(packageName, userId);
         Set<String> authorities = new HashSet<>();