Allow overriding enterprise badge drawables

Bug: 203548565
Bug: 188410712
Test: manual testing
Change-Id: Ibee6c28e8a231eab9c370af3065619c2dffa706a
diff --git a/core/java/android/app/ApplicationPackageManager.java b/core/java/android/app/ApplicationPackageManager.java
index 7b55b6c..b1956ef 100644
--- a/core/java/android/app/ApplicationPackageManager.java
+++ b/core/java/android/app/ApplicationPackageManager.java
@@ -16,6 +16,11 @@
 
 package android.app;
 
+import static android.app.admin.DevicePolicyResources.Drawables.Style.SOLID_COLORED;
+import static android.app.admin.DevicePolicyResources.Drawables.Style.SOLID_NOT_COLORED;
+import static android.app.admin.DevicePolicyResources.Drawables.UNDEFINED;
+import static android.app.admin.DevicePolicyResources.Drawables.WORK_PROFILE_ICON;
+import static android.app.admin.DevicePolicyResources.Drawables.WORK_PROFILE_ICON_BADGE;
 import static android.content.pm.Checksum.TYPE_PARTIAL_MERKLE_ROOT_1M_SHA256;
 import static android.content.pm.Checksum.TYPE_PARTIAL_MERKLE_ROOT_1M_SHA512;
 import static android.content.pm.Checksum.TYPE_WHOLE_MD5;
@@ -31,6 +36,7 @@
 import android.annotation.StringRes;
 import android.annotation.UserIdInt;
 import android.annotation.XmlRes;
+import android.app.admin.DevicePolicyManager;
 import android.app.role.RoleManager;
 import android.compat.annotation.UnsupportedAppUsage;
 import android.content.ComponentName;
@@ -62,7 +68,6 @@
 import android.content.pm.PackageInstaller;
 import android.content.pm.PackageItemInfo;
 import android.content.pm.PackageManager;
-import android.content.pm.PackageParser;
 import android.content.pm.ParceledListSlice;
 import android.content.pm.PermissionGroupInfo;
 import android.content.pm.PermissionInfo;
@@ -74,7 +79,6 @@
 import android.content.pm.VerifierDeviceIdentity;
 import android.content.pm.VersionedPackage;
 import android.content.pm.dex.ArtManager;
-import android.content.pm.pkg.FrameworkPackageUserState;
 import android.content.res.Configuration;
 import android.content.res.Resources;
 import android.content.res.XmlResourceParser;
@@ -122,7 +126,6 @@
 
 import libcore.util.EmptyArray;
 
-import java.io.File;
 import java.lang.ref.WeakReference;
 import java.security.cert.Certificate;
 import java.security.cert.CertificateEncodingException;
@@ -172,6 +175,8 @@
     private PackageInstaller mInstaller;
     @GuardedBy("mLock")
     private ArtManager mArtManager;
+    @GuardedBy("mLock")
+    private DevicePolicyManager mDevicePolicyManager;
 
     @GuardedBy("mDelegates")
     private final ArrayList<MoveCallbackDelegate> mDelegates = new ArrayList<>();
@@ -188,6 +193,15 @@
         }
     }
 
+    DevicePolicyManager getDevicePolicyManager() {
+        synchronized (mLock) {
+            if (mDevicePolicyManager == null) {
+                mDevicePolicyManager = mContext.getSystemService(DevicePolicyManager.class);
+            }
+            return mDevicePolicyManager;
+        }
+    }
+
     private PermissionManager getPermissionManager() {
         synchronized (mLock) {
             if (mPermissionManager == null) {
@@ -1876,12 +1890,27 @@
         if (!hasUserBadge(user.getIdentifier())) {
             return icon;
         }
+
+        final Drawable badgeForeground = getDevicePolicyManager().getDrawable(
+                getUpdatableUserIconBadgeId(user),
+                SOLID_COLORED,
+                () -> getDefaultUserIconBadge(user));
+
         Drawable badge = new LauncherIcons(mContext).getBadgeDrawable(
-                getUserManager().getUserIconBadgeResId(user.getIdentifier()),
+                badgeForeground,
                 getUserBadgeColor(user, false));
         return getBadgedDrawable(icon, badge, null, true);
     }
 
+    private String getUpdatableUserIconBadgeId(UserHandle user) {
+        return getUserManager().isManagedProfile(user.getIdentifier())
+                ? WORK_PROFILE_ICON_BADGE : UNDEFINED;
+    }
+
+    private Drawable getDefaultUserIconBadge(UserHandle user) {
+        return mContext.getDrawable(getUserManager().getUserIconBadgeResId(user.getIdentifier()));
+    }
+
     @Override
     public Drawable getUserBadgedDrawableForDensity(Drawable drawable, UserHandle user,
             Rect badgeLocation, int badgeDensity) {
@@ -1913,13 +1942,28 @@
         if (badgeColor == null) {
             return null;
         }
-        Drawable badgeForeground = getDrawableForDensity(
-                getUserManager().getUserBadgeResId(user.getIdentifier()), density);
+
+        final Drawable badgeForeground = getDevicePolicyManager().getDrawableForDensity(
+                getUpdatableUserBadgeId(user),
+                SOLID_COLORED,
+                density,
+                () -> getDefaultUserBadgeForDensity(user, density));
+
         badgeForeground.setTint(getUserBadgeColor(user, false));
         Drawable badge = new LayerDrawable(new Drawable[] {badgeColor, badgeForeground });
         return badge;
     }
 
+    private String getUpdatableUserBadgeId(UserHandle user) {
+        return getUserManager().isManagedProfile(user.getIdentifier())
+                ? WORK_PROFILE_ICON : UNDEFINED;
+    }
+
+    private Drawable getDefaultUserBadgeForDensity(UserHandle user, int density) {
+        return getDrawableForDensity(
+                getUserManager().getUserBadgeResId(user.getIdentifier()), density);
+    }
+
     /**
      * Returns the badge color based on whether device has dark theme enabled or not.
      */
@@ -1928,14 +1972,24 @@
         if (!hasUserBadge(user.getIdentifier())) {
             return null;
         }
-        Drawable badge = getDrawableForDensity(
-                getUserManager().getUserBadgeNoBackgroundResId(user.getIdentifier()), density);
+
+        final Drawable badge = getDevicePolicyManager().getDrawableForDensity(
+                getUpdatableUserBadgeId(user),
+                SOLID_NOT_COLORED,
+                density,
+                () -> getDefaultUserBadgeNoBackgroundForDensity(user, density));
+
         if (badge != null) {
             badge.setTint(getUserBadgeColor(user, true));
         }
         return badge;
     }
 
+    private Drawable getDefaultUserBadgeNoBackgroundForDensity(UserHandle user, int density) {
+        return getDrawableForDensity(
+                getUserManager().getUserBadgeNoBackgroundResId(user.getIdentifier()), density);
+    }
+
     private Drawable getDrawableForDensity(int drawableId, int density) {
         if (density <= 0) {
             density = mContext.getResources().getDisplayMetrics().densityDpi;
diff --git a/core/java/android/util/IconDrawableFactory.java b/core/java/android/util/IconDrawableFactory.java
index b5e8dd7..5bb263a 100644
--- a/core/java/android/util/IconDrawableFactory.java
+++ b/core/java/android/util/IconDrawableFactory.java
@@ -15,7 +15,12 @@
  */
 package android.util;
 
+import static android.app.admin.DevicePolicyResources.Drawables.Style.SOLID_COLORED;
+import static android.app.admin.DevicePolicyResources.Drawables.UNDEFINED;
+import static android.app.admin.DevicePolicyResources.Drawables.WORK_PROFILE_ICON_BADGE;
+
 import android.annotation.UserIdInt;
+import android.app.admin.DevicePolicyManager;
 import android.compat.annotation.UnsupportedAppUsage;
 import android.content.Context;
 import android.content.pm.ApplicationInfo;
@@ -37,6 +42,7 @@
     protected final Context mContext;
     protected final PackageManager mPm;
     protected final UserManager mUm;
+    protected final DevicePolicyManager mDpm;
     protected final LauncherIcons mLauncherIcons;
     protected final boolean mEmbedShadow;
 
@@ -44,6 +50,7 @@
         mContext = context;
         mPm = context.getPackageManager();
         mUm = context.getSystemService(UserManager.class);
+        mDpm = context.getSystemService(DevicePolicyManager.class);
         mLauncherIcons = new LauncherIcons(context);
         mEmbedShadow = embedShadow;
     }
@@ -73,18 +80,32 @@
         if (appInfo.isInstantApp()) {
             int badgeColor = Resources.getSystem().getColor(
                     com.android.internal.R.color.instant_app_badge, null);
+            Drawable badge = mContext.getDrawable(
+                    com.android.internal.R.drawable.ic_instant_icon_badge_bolt);
             icon = mLauncherIcons.getBadgedDrawable(icon,
-                    com.android.internal.R.drawable.ic_instant_icon_badge_bolt,
+                    badge,
                     badgeColor);
         }
         if (mUm.hasBadge(userId)) {
-            icon = mLauncherIcons.getBadgedDrawable(icon,
-                    mUm.getUserIconBadgeResId(userId),
-                    mUm.getUserBadgeColor(userId));
+
+            Drawable badge = mDpm.getDrawable(
+                    getUpdatableUserIconBadgeId(userId),
+                    SOLID_COLORED,
+                    () -> getDefaultUserIconBadge(userId));
+
+            icon = mLauncherIcons.getBadgedDrawable(icon, badge, mUm.getUserBadgeColor(userId));
         }
         return icon;
     }
 
+    private String getUpdatableUserIconBadgeId(int userId) {
+        return mUm.isManagedProfile(userId) ? WORK_PROFILE_ICON_BADGE : UNDEFINED;
+    }
+
+    private Drawable getDefaultUserIconBadge(int userId) {
+        return mContext.getResources().getDrawable(mUm.getUserIconBadgeResId(userId));
+    }
+
     /**
      * Add shadow to the icon if {@link AdaptiveIconDrawable}
      */
diff --git a/core/java/android/util/LauncherIcons.java b/core/java/android/util/LauncherIcons.java
index e652e17..355b2e9 100644
--- a/core/java/android/util/LauncherIcons.java
+++ b/core/java/android/util/LauncherIcons.java
@@ -45,10 +45,12 @@
     private final SparseArray<Bitmap> mShadowCache = new SparseArray<>();
     private final int mIconSize;
     private final Resources mRes;
+    private final Context mContext;
 
     public LauncherIcons(Context context) {
         mRes = context.getResources();
         mIconSize = mRes.getDimensionPixelSize(android.R.dimen.app_icon_size);
+        mContext = context;
     }
 
     public Drawable wrapIconDrawableWithShadow(Drawable drawable) {
@@ -98,14 +100,14 @@
         return shadow;
     }
 
-    public Drawable getBadgeDrawable(int foregroundRes, int backgroundColor) {
-        return getBadgedDrawable(null, foregroundRes, backgroundColor);
+    public Drawable getBadgeDrawable(Drawable badgeForeground, int backgroundColor) {
+        return getBadgedDrawable(null, badgeForeground, backgroundColor);
     }
 
-    public Drawable getBadgedDrawable(Drawable base, int foregroundRes, int backgroundColor) {
+    public Drawable getBadgedDrawable(
+            Drawable base, Drawable badgeForeground, int backgroundColor) {
         Resources overlayableRes =
                 ActivityThread.currentActivityThread().getApplication().getResources();
-
         // ic_corp_icon_badge_shadow is not work-profile-specific.
         Drawable badgeShadow = overlayableRes.getDrawable(
                 com.android.internal.R.drawable.ic_corp_icon_badge_shadow);
@@ -115,7 +117,6 @@
                 com.android.internal.R.drawable.ic_corp_icon_badge_color)
                 .getConstantState().newDrawable().mutate();
 
-        Drawable badgeForeground = overlayableRes.getDrawable(foregroundRes);
         badgeForeground.setTint(backgroundColor);
 
         Drawable[] drawables = base == null