Add API to allow DMRH to set max policy storage limit per admin

Bug: 325036682
Test: btest android.devicepolicy.cts.DevicePolicyManagerTest
Change-Id: I1e408fbf418abe9adc7a1de7d7e3030a97496f0e
diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index 88f9aff..4719561 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -1307,6 +1307,7 @@
     method @Nullable @RequiresPermission(android.Manifest.permission.MANAGE_USERS) public android.os.UserHandle getDeviceOwnerUser();
     method @NonNull @RequiresPermission(android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS) public android.app.admin.DevicePolicyState getDevicePolicyState();
     method @Nullable @RequiresPermission(android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS) public String getFinancedDeviceKioskRoleHolder();
+    method @FlaggedApi("android.app.admin.flags.device_policy_size_tracking_enabled") @RequiresPermission(android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS) public int getMaxPolicyStorageLimit();
     method @Nullable @RequiresPermission(anyOf={android.Manifest.permission.MANAGE_USERS, android.Manifest.permission.QUERY_ADMIN_POLICY}) public java.util.List<java.lang.String> getPermittedAccessibilityServices(int);
     method @Nullable @RequiresPermission(anyOf={android.Manifest.permission.MANAGE_USERS, android.Manifest.permission.QUERY_ADMIN_POLICY}) public java.util.List<java.lang.String> getPermittedInputMethodsForCurrentUser();
     method @NonNull @RequiresPermission(android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS) public java.util.List<android.os.UserHandle> getPolicyManagedProfiles(@NonNull android.os.UserHandle);
@@ -1329,6 +1330,7 @@
     method @RequiresPermission(android.Manifest.permission.MANAGE_DEVICE_POLICY_APP_EXEMPTIONS) public void setApplicationExemptions(@NonNull String, @NonNull java.util.Set<java.lang.Integer>) throws android.content.pm.PackageManager.NameNotFoundException;
     method @RequiresPermission(android.Manifest.permission.MANAGE_USERS) public void setDeviceProvisioningConfigApplied();
     method @RequiresPermission(android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS) public void setDpcDownloaded(boolean);
+    method @FlaggedApi("android.app.admin.flags.device_policy_size_tracking_enabled") @RequiresPermission(android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS) public void setMaxPolicyStorageLimit(int);
     method @Deprecated @RequiresPermission(value=android.Manifest.permission.GRANT_PROFILE_OWNER_DEVICE_IDS_ACCESS, conditional=true) public void setProfileOwnerCanAccessDeviceIds(@NonNull android.content.ComponentName);
     method public void setSecondaryLockscreenEnabled(@NonNull android.content.ComponentName, boolean);
     method @RequiresPermission(android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS) public void setUserProvisioningState(int, @NonNull android.os.UserHandle);
diff --git a/core/java/android/app/admin/DevicePolicyManager.java b/core/java/android/app/admin/DevicePolicyManager.java
index 5c6308f..c53b54c 100644
--- a/core/java/android/app/admin/DevicePolicyManager.java
+++ b/core/java/android/app/admin/DevicePolicyManager.java
@@ -52,6 +52,7 @@
 import static android.Manifest.permission.SET_TIME;
 import static android.Manifest.permission.SET_TIME_ZONE;
 import static android.app.admin.flags.Flags.FLAG_ESIM_MANAGEMENT_ENABLED;
+import static android.app.admin.flags.Flags.FLAG_DEVICE_POLICY_SIZE_TRACKING_ENABLED;
 import static android.app.admin.flags.Flags.onboardingBugreportV2Enabled;
 import static android.content.Intent.LOCAL_FLAG_FROM_SYSTEM;
 import static android.net.NetworkCapabilities.NET_ENTERPRISE_ID_1;
@@ -17348,4 +17349,46 @@
         }
         return new HashSet<>();
     }
+
+    /**
+     * Controls the maximum storage size allowed for policies associated with an admin.
+     * Setting a limit of -1 effectively removes any storage restrictions.
+     *
+     * @param storageLimit Maximum storage allowed in bytes. Use -1 to disable limits.
+     *
+     * @hide
+     */
+    @SystemApi
+    @RequiresPermission(permission.MANAGE_PROFILE_AND_DEVICE_OWNERS)
+    @FlaggedApi(FLAG_DEVICE_POLICY_SIZE_TRACKING_ENABLED)
+    public void setMaxPolicyStorageLimit(int storageLimit) {
+        if (mService != null) {
+            try {
+                mService.setMaxPolicyStorageLimit(mContext.getPackageName(), storageLimit);
+            } catch (RemoteException e) {
+                throw e.rethrowFromSystemServer();
+            }
+        }
+    }
+
+    /**
+     * Retrieves the current maximum storage limit for policies associated with an admin.
+     *
+     * @return The maximum storage limit in bytes, or -1 if no limit is enforced.
+     *
+     * @hide
+     */
+    @SystemApi
+    @RequiresPermission(permission.MANAGE_PROFILE_AND_DEVICE_OWNERS)
+    @FlaggedApi(FLAG_DEVICE_POLICY_SIZE_TRACKING_ENABLED)
+    public int getMaxPolicyStorageLimit() {
+        if (mService != null) {
+            try {
+                return mService.getMaxPolicyStorageLimit(mContext.getPackageName());
+            } catch (RemoteException e) {
+                throw e.rethrowFromSystemServer();
+            }
+        }
+        return -1;
+    }
 }
\ No newline at end of file
diff --git a/core/java/android/app/admin/IDevicePolicyManager.aidl b/core/java/android/app/admin/IDevicePolicyManager.aidl
index f72fdc0..c4cbdd6 100644
--- a/core/java/android/app/admin/IDevicePolicyManager.aidl
+++ b/core/java/android/app/admin/IDevicePolicyManager.aidl
@@ -615,4 +615,7 @@
     int getContentProtectionPolicy(in ComponentName who, String callerPackageName);
 
     int[] getSubscriptionIds(String callerPackageName);
+
+    void setMaxPolicyStorageLimit(String packageName, int storageLimit);
+    int getMaxPolicyStorageLimit(String packageName);
 }
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyEngine.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyEngine.java
index 105dc88..0ccf810c 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyEngine.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyEngine.java
@@ -102,6 +102,9 @@
             DevicePolicyIdentifiers.getIdentifierForUserRestriction(
                     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;
+
     private final Context mContext;
     private final UserManager mUserManager;
 
@@ -122,10 +125,11 @@
      * Map containing the current set of admins in each user with active policies.
      */
     private final SparseArray<Set<EnforcingAdmin>> mEnforcingAdmins;
+
     private final SparseArray<HashMap<EnforcingAdmin, Integer>> mAdminPolicySize;
 
-    //TODO(b/295504706) : Speak to security team to decide what to set Policy_Size_Limit
-    private static final int POLICY_SIZE_LIMIT = 99999;
+    private int mPolicySizeLimit = DEFAULT_POLICY_SIZE_LIMIT;
+
     private final DeviceAdminServiceController mDeviceAdminServiceController;
 
     DevicePolicyEngine(
@@ -1594,7 +1598,9 @@
             existingPolicySize = sizeOf(policyState.getPoliciesSetByAdmins().get(admin));
         }
         int policySize = sizeOf(value);
-        if (currentAdminPoliciesSize + policySize - existingPolicySize < POLICY_SIZE_LIMIT) {
+        // Policy size limit is disabled if mPolicySizeLimit is -1.
+        if (mPolicySizeLimit == -1
+                || currentAdminPoliciesSize + policySize - existingPolicySize < mPolicySizeLimit) {
             increasePolicySizeForAdmin(
                     admin, /* policySizeDiff = */ policySize - existingPolicySize);
             return true;
@@ -1642,6 +1648,26 @@
         }
     }
 
+    /**
+     * Updates the max allowed size limit for policies per admin. Setting it to -1, disables
+     * 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;
+    }
+
+    /**
+     * Returns the max allowed size limit for policies per admin. -1 means the limitation is
+     * disabled.
+     */
+    int getMaxPolicyStorageLimit() {
+        return mPolicySizeLimit;
+    }
+
     public void dump(IndentingPrintWriter pw) {
         synchronized (mLock) {
             pw.println("Local Policies: ");
@@ -1761,6 +1787,7 @@
         private static final String TAG_ENFORCING_ADMIN_AND_SIZE = "enforcing-admin-and-size";
         private static final String TAG_ENFORCING_ADMIN = "enforcing-admin";
         private static final String TAG_POLICY_SUM_SIZE = "policy-sum-size";
+        private static final String TAG_MAX_POLICY_SIZE_LIMIT = "max-policy-size-limit";
         private static final String ATTR_USER_ID = "user-id";
         private static final String ATTR_POLICY_SUM_SIZE = "size";
 
@@ -1805,6 +1832,7 @@
             writeGlobalPoliciesInner(serializer);
             writeEnforcingAdminsInner(serializer);
             writeEnforcingAdminSizeInner(serializer);
+            writeMaxPolicySizeInner(serializer);
         }
 
         private void writeLocalPoliciesInner(TypedXmlSerializer serializer) throws IOException {
@@ -1886,6 +1914,17 @@
             }
         }
 
+        private void writeMaxPolicySizeInner(TypedXmlSerializer serializer)
+                throws IOException {
+            if (!devicePolicySizeTrackingEnabled()) {
+                return;
+            }
+            serializer.startTag(/* namespace= */ null, TAG_MAX_POLICY_SIZE_LIMIT);
+            serializer.attributeInt(
+                    /* namespace= */ null, ATTR_POLICY_SUM_SIZE, mPolicySizeLimit);
+            serializer.endTag(/* namespace= */ null, TAG_MAX_POLICY_SIZE_LIMIT);
+        }
+
         void readFromFileLocked() {
             if (!mFile.exists()) {
                 Log.d(TAG, "" + mFile + " doesn't exist");
@@ -1926,6 +1965,9 @@
                     case TAG_ENFORCING_ADMIN_AND_SIZE:
                         readEnforcingAdminAndSizeInner(parser);
                         break;
+                    case TAG_MAX_POLICY_SIZE_LIMIT:
+                        readMaxPolicySizeInner(parser);
+                        break;
                     default:
                         Slogf.wtf(TAG, "Unknown tag " + tag);
                 }
@@ -2036,5 +2078,13 @@
             }
             mAdminPolicySize.get(admin.getUserId()).put(admin, size);
         }
+
+        private void readMaxPolicySizeInner(TypedXmlPullParser parser)
+                throws XmlPullParserException, IOException {
+            if (!devicePolicySizeTrackingEnabled()) {
+                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 b7a2271..45f387e 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -12042,8 +12042,10 @@
         }
 
         if (packageList != null) {
-            for (String pkg : packageList) {
-                PolicySizeVerifier.enforceMaxPackageNameLength(pkg);
+            if (!devicePolicySizeTrackingEnabled()) {
+                for (String pkg : packageList) {
+                    PolicySizeVerifier.enforceMaxPackageNameLength(pkg);
+                }
             }
 
             List<InputMethodInfo> enabledImes = mInjector.binderWithCleanCallingIdentity(() ->
@@ -14362,8 +14364,10 @@
     public void setLockTaskPackages(ComponentName who, String callerPackageName, String[] packages)
             throws SecurityException {
         Objects.requireNonNull(packages, "packages is null");
-        for (String pkg : packages) {
-            PolicySizeVerifier.enforceMaxPackageNameLength(pkg);
+        if (!devicePolicySizeTrackingEnabled()) {
+            for (String pkg : packages) {
+                PolicySizeVerifier.enforceMaxPackageNameLength(pkg);
+            }
         }
 
         CallerIdentity caller = getCallerIdentity(who, callerPackageName);
@@ -24075,5 +24079,30 @@
             }
             return adminOwnedSubscriptions;
         });
+
+    }
+
+    @Override
+    public void setMaxPolicyStorageLimit(String callerPackageName, int storageLimit) {
+        if (!devicePolicySizeTrackingEnabled()) {
+            return;
+        }
+        CallerIdentity caller = getCallerIdentity(callerPackageName);
+        enforcePermission(MANAGE_PROFILE_AND_DEVICE_OWNERS, caller.getPackageName(),
+                caller.getUserId());
+
+        mDevicePolicyEngine.setMaxPolicyStorageLimit(storageLimit);
+    }
+
+    @Override
+    public int getMaxPolicyStorageLimit(String callerPackageName) {
+        if (!devicePolicySizeTrackingEnabled()) {
+            return -1;
+        }
+        CallerIdentity caller = getCallerIdentity(callerPackageName);
+        enforcePermission(MANAGE_PROFILE_AND_DEVICE_OWNERS, caller.getPackageName(),
+                caller.getUserId());
+
+        return mDevicePolicyEngine.getMaxPolicyStorageLimit();
     }
 }