diff --git a/core/java/android/app/admin/DevicePolicyResourcesManager.java b/core/java/android/app/admin/DevicePolicyResourcesManager.java
index 06729222dea1c..e8eb792f7fd78 100644
--- a/core/java/android/app/admin/DevicePolicyResourcesManager.java
+++ b/core/java/android/app/admin/DevicePolicyResourcesManager.java
@@ -70,6 +70,7 @@ public class DevicePolicyResourcesManager {
      *
      * <p>Important notes to consider when using this API:
      * <ul>
+     * <li> Updated resources are persisted over reboots.
      * <li>{@link #getDrawable} references the resource
      * {@link DevicePolicyDrawableResource#getResourceIdInCallingPackage()} in the
      * calling package each time it gets called. You have to ensure that the resource is always
@@ -381,7 +382,9 @@ public class DevicePolicyResourcesManager {
      *
      * <p>Important notes to consider when using this API:
      * <ul>
-     * <li> {@link #getString} references the resource {@code callingPackageResourceId} in the
+     * <li> Updated resources are persisted over reboots.
+     * <li> {@link #getString} references the resource
+     * {@link DevicePolicyStringResource#getResourceIdInCallingPackage()} in the
      * calling package each time it gets called. You have to ensure that the resource is always
      * available in the calling package as long as it is used as an updated resource.
      * <li> You still have to re-call {@code setStrings} even if you only make changes to the
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DeviceManagementResourcesProvider.java b/services/devicepolicy/java/com/android/server/devicepolicy/DeviceManagementResourcesProvider.java
index 82fe8b9362b07..6aef90c797055 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DeviceManagementResourcesProvider.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DeviceManagementResourcesProvider.java
@@ -55,8 +55,6 @@ class DeviceManagementResourcesProvider {
     private static final String TAG_ROOT = "root";
     private static final String TAG_DRAWABLE_STYLE_ENTRY = "drawable-style-entry";
     private static final String TAG_DRAWABLE_SOURCE_ENTRY = "drawable-source-entry";
-    private static final String ATTR_DRAWABLE_STYLE_SIZE = "drawable-style-size";
-    private static final String ATTR_DRAWABLE_SOURCE_SIZE = "drawable-source-size";
     private static final String ATTR_DRAWABLE_STYLE = "drawable-style";
     private static final String ATTR_DRAWABLE_SOURCE = "drawable-source";
     private static final String ATTR_DRAWABLE_ID = "drawable-id";
@@ -70,9 +68,9 @@ class DeviceManagementResourcesProvider {
             mUpdatedDrawablesForStyle = new HashMap<>();
 
     /**
-     * Map of <drawable_id, <source_id, resource_value>>
+     * Map of <drawable_id, <source_id, <style_id, resource_value>>>
      */
-    private final Map<String, Map<String, ParcelableResource>>
+    private final Map<String, Map<String, Map<String, ParcelableResource>>>
             mUpdatedDrawablesForSource = new HashMap<>();
 
     /**
@@ -110,7 +108,8 @@ class DeviceManagementResourcesProvider {
             if (DevicePolicyResources.UNDEFINED.equals(drawableSource)) {
                 updated |= updateDrawable(drawableId, drawableStyle, resource);
             } else {
-                updated |= updateDrawableForSource(drawableId, drawableSource, resource);
+                updated |= updateDrawableForSource(
+                        drawableId, drawableSource, drawableStyle, resource);
             }
         }
         if (!updated) {
@@ -138,19 +137,23 @@ class DeviceManagementResourcesProvider {
         }
     }
 
-    // TODO(b/214576716): change this to respect style
     private boolean updateDrawableForSource(
-            String drawableId, String drawableSource, ParcelableResource updatableResource) {
+            String drawableId, String drawableSource, String drawableStyle,
+            ParcelableResource updatableResource) {
         synchronized (mLock) {
+            Map<String, Map<String, ParcelableResource>> drawablesForId =
+                    mUpdatedDrawablesForSource.get(drawableId);
             if (!mUpdatedDrawablesForSource.containsKey(drawableId)) {
                 mUpdatedDrawablesForSource.put(drawableId, new HashMap<>());
             }
-            ParcelableResource current = mUpdatedDrawablesForSource.get(drawableId).get(
-                    drawableSource);
+            if (!drawablesForId.containsKey(drawableSource)) {
+                mUpdatedDrawablesForSource.get(drawableId).put(drawableSource, new HashMap<>());
+            }
+            ParcelableResource current = drawablesForId.get(drawableSource).get(drawableStyle);
             if (updatableResource.equals(current)) {
                 return false;
             }
-            mUpdatedDrawablesForSource.get(drawableId).put(drawableSource, updatableResource);
+            drawablesForId.get(drawableSource).put(drawableStyle, updatableResource);
             return true;
         }
     }
@@ -175,23 +178,30 @@ class DeviceManagementResourcesProvider {
     }
 
     @Nullable
-    ParcelableResource getDrawable(
-            String drawableId, String drawableStyle, String drawableSource) {
+    ParcelableResource getDrawable(String drawableId, String drawableStyle, String drawableSource) {
         synchronized (mLock) {
-            if (mUpdatedDrawablesForSource.containsKey(drawableId)
-                    && mUpdatedDrawablesForSource.get(drawableId).containsKey(drawableSource)) {
-                return mUpdatedDrawablesForSource.get(drawableId).get(drawableSource);
+            ParcelableResource resource = getDrawableForSourceLocked(
+                    drawableId, drawableStyle, drawableSource);
+            if (resource != null) {
+                return resource;
             }
             if (!mUpdatedDrawablesForStyle.containsKey(drawableId)) {
-                Log.d(TAG, "No updated drawable found for drawable id " + drawableId);
                 return null;
             }
-            if (mUpdatedDrawablesForStyle.get(drawableId).containsKey(drawableStyle)) {
-                return mUpdatedDrawablesForStyle.get(drawableId).get(drawableStyle);
-            }
+            return mUpdatedDrawablesForStyle.get(drawableId).get(drawableStyle);
         }
-        Log.d(TAG, "No updated drawable found for drawable id " + drawableId);
-        return null;
+    }
+
+    @Nullable
+    ParcelableResource getDrawableForSourceLocked(
+            String drawableId, String drawableStyle, String drawableSource) {
+        if (!mUpdatedDrawablesForSource.containsKey(drawableId)) {
+            return null;
+        }
+        if (!mUpdatedDrawablesForSource.get(drawableId).containsKey(drawableSource)) {
+            return null;
+        }
+        return mUpdatedDrawablesForSource.get(drawableId).get(drawableSource).get(drawableStyle);
     }
 
     /**
@@ -249,12 +259,8 @@ class DeviceManagementResourcesProvider {
     @Nullable
     ParcelableResource getString(String stringId) {
         synchronized (mLock) {
-            if (mUpdatedStrings.containsKey(stringId)) {
-                return mUpdatedStrings.get(stringId);
-            }
+            return mUpdatedStrings.get(stringId);
         }
-        Log.d(TAG, "No updated string found for string id " + stringId);
-        return null;
     }
 
     private void write() {
@@ -359,50 +365,55 @@ class DeviceManagementResourcesProvider {
         }
 
         void writeInner(TypedXmlSerializer out) throws IOException {
+            writeDrawablesForStylesInner(out);
+            writeDrawablesForSourcesInner(out);
+            writeStringsInner(out);
+        }
+
+        private void writeDrawablesForStylesInner(TypedXmlSerializer out) throws IOException {
             if (mUpdatedDrawablesForStyle != null && !mUpdatedDrawablesForStyle.isEmpty()) {
                 for (Map.Entry<String, Map<String, ParcelableResource>> drawableEntry
                         : mUpdatedDrawablesForStyle.entrySet()) {
-                    out.startTag(/* namespace= */ null, TAG_DRAWABLE_STYLE_ENTRY);
-                    out.attribute(
-                            /* namespace= */ null, ATTR_DRAWABLE_ID, drawableEntry.getKey());
-                    out.attributeInt(
-                            /* namespace= */ null,
-                            ATTR_DRAWABLE_STYLE_SIZE,
-                            drawableEntry.getValue().size());
-                    int counter = 0;
                     for (Map.Entry<String, ParcelableResource> styleEntry
                             : drawableEntry.getValue().entrySet()) {
+                        out.startTag(/* namespace= */ null, TAG_DRAWABLE_STYLE_ENTRY);
+                        out.attribute(
+                                /* namespace= */ null, ATTR_DRAWABLE_ID, drawableEntry.getKey());
                         out.attribute(
                                 /* namespace= */ null,
-                                ATTR_DRAWABLE_STYLE + (counter++),
+                                ATTR_DRAWABLE_STYLE,
                                 styleEntry.getKey());
                         styleEntry.getValue().writeToXmlFile(out);
+                        out.endTag(/* namespace= */ null, TAG_DRAWABLE_STYLE_ENTRY);
                     }
-                    out.endTag(/* namespace= */ null, TAG_DRAWABLE_STYLE_ENTRY);
                 }
             }
+        }
+
+        private void writeDrawablesForSourcesInner(TypedXmlSerializer out) throws IOException {
             if (mUpdatedDrawablesForSource != null && !mUpdatedDrawablesForSource.isEmpty()) {
-                for (Map.Entry<String, Map<String, ParcelableResource>> drawableEntry
+                for (Map.Entry<String, Map<String, Map<String, ParcelableResource>>> drawableEntry
                         : mUpdatedDrawablesForSource.entrySet()) {
-                    out.startTag(/* namespace= */ null, TAG_DRAWABLE_SOURCE_ENTRY);
-                    out.attribute(
-                            /* namespace= */ null, ATTR_DRAWABLE_ID, drawableEntry.getKey());
-                    out.attributeInt(
-                            /* namespace= */ null,
-                            ATTR_DRAWABLE_SOURCE_SIZE,
-                            drawableEntry.getValue().size());
-                    int counter = 0;
-                    for (Map.Entry<String, ParcelableResource> sourceEntry
+                    for (Map.Entry<String, Map<String, ParcelableResource>> sourceEntry
                             : drawableEntry.getValue().entrySet()) {
-                        out.attribute(
-                                /* namespace= */ null,
-                                ATTR_DRAWABLE_SOURCE + (counter++),
-                                sourceEntry.getKey());
-                        sourceEntry.getValue().writeToXmlFile(out);
+                        for (Map.Entry<String, ParcelableResource> styleEntry
+                                : sourceEntry.getValue().entrySet()) {
+                            out.startTag(/* namespace= */ null, TAG_DRAWABLE_SOURCE_ENTRY);
+                            out.attribute(/* namespace= */ null, ATTR_DRAWABLE_ID,
+                                    drawableEntry.getKey());
+                            out.attribute(/* namespace= */ null, ATTR_DRAWABLE_SOURCE,
+                                    sourceEntry.getKey());
+                            out.attribute(/* namespace= */ null, ATTR_DRAWABLE_STYLE,
+                                    styleEntry.getKey());
+                            styleEntry.getValue().writeToXmlFile(out);
+                            out.endTag(/* namespace= */ null, TAG_DRAWABLE_SOURCE_ENTRY);
+                        }
                     }
-                    out.endTag(/* namespace= */ null, TAG_DRAWABLE_SOURCE_ENTRY);
                 }
             }
+        }
+
+        private void writeStringsInner(TypedXmlSerializer out) throws IOException {
             if (mUpdatedStrings != null && !mUpdatedStrings.isEmpty()) {
                 for (Map.Entry<String, ParcelableResource> entry
                         : mUpdatedStrings.entrySet()) {
@@ -417,52 +428,48 @@ class DeviceManagementResourcesProvider {
             }
         }
 
-        private boolean readInner(
-                TypedXmlPullParser parser, int depth, String tag)
+        private boolean readInner(TypedXmlPullParser parser, int depth, String tag)
                 throws XmlPullParserException, IOException {
             if (depth > 2) {
                 return true; // Ignore
             }
             switch (tag) {
-                case TAG_DRAWABLE_STYLE_ENTRY:
-                    String drawableId = parser.getAttributeValue(
-                            /* namespace= */ null, ATTR_DRAWABLE_ID);
-                    mUpdatedDrawablesForStyle.put(
-                            drawableId,
-                            new HashMap<>());
-                    int size = parser.getAttributeInt(
-                            /* namespace= */ null, ATTR_DRAWABLE_STYLE_SIZE);
-                    for (int i = 0; i < size; i++) {
-                        String style = parser.getAttributeValue(
-                                /* namespace= */ null, ATTR_DRAWABLE_STYLE + i);
-                        mUpdatedDrawablesForStyle.get(drawableId).put(
-                                style,
-                                ParcelableResource.createFromXml(parser));
+                case TAG_DRAWABLE_STYLE_ENTRY: {
+                    String id = parser.getAttributeValue(/* namespace= */ null, ATTR_DRAWABLE_ID);
+                    String style = parser.getAttributeValue(
+                            /* namespace= */ null, ATTR_DRAWABLE_STYLE);
+                    ParcelableResource resource = ParcelableResource.createFromXml(parser);
+                    if (!mUpdatedDrawablesForStyle.containsKey(id)) {
+                        mUpdatedDrawablesForStyle.put(id, new HashMap<>());
                     }
+                    mUpdatedDrawablesForStyle.get(id).put(style, resource);
                     break;
-                case TAG_DRAWABLE_SOURCE_ENTRY:
-                    drawableId = parser.getAttributeValue(
-                            /* namespace= */ null, ATTR_DRAWABLE_ID);
-                    mUpdatedDrawablesForSource.put(drawableId, new HashMap<>());
-                    size = parser.getAttributeInt(
-                            /* namespace= */ null, ATTR_DRAWABLE_SOURCE_SIZE);
-                    for (int i = 0; i < size; i++) {
-                        String source = parser.getAttributeValue(
-                                /* namespace= */ null, ATTR_DRAWABLE_SOURCE + i);
-                        mUpdatedDrawablesForSource.get(drawableId).put(
-                                source,
-                                ParcelableResource.createFromXml(parser));
+                }
+                case TAG_DRAWABLE_SOURCE_ENTRY: {
+                    String id = parser.getAttributeValue(/* namespace= */ null, ATTR_DRAWABLE_ID);
+                    String source = parser.getAttributeValue(
+                            /* namespace= */ null, ATTR_DRAWABLE_SOURCE);
+                    String style = parser.getAttributeValue(
+                            /* namespace= */ null, ATTR_DRAWABLE_STYLE);
+                    ParcelableResource resource = ParcelableResource.createFromXml(parser);
+                    if (!mUpdatedDrawablesForSource.containsKey(id)) {
+                        mUpdatedDrawablesForSource.put(id, new HashMap<>());
                     }
+                    if (!mUpdatedDrawablesForSource.get(id).containsKey(source)) {
+                        mUpdatedDrawablesForSource.get(id).put(source, new HashMap<>());
+                    }
+                    mUpdatedDrawablesForSource.get(id).get(source).put(style, resource);
                     break;
-                case TAG_STRING_ENTRY:
-                    String sourceId = parser.getAttributeValue(
-                            /* namespace= */ null, ATTR_SOURCE_ID);
-                    mUpdatedStrings.put(
-                            sourceId, ParcelableResource.createFromXml(parser));
+                }
+                case TAG_STRING_ENTRY: {
+                    String id = parser.getAttributeValue(/* namespace= */ null, ATTR_SOURCE_ID);
+                    mUpdatedStrings.put(id, ParcelableResource.createFromXml(parser));
                     break;
-                default:
+                }
+                default: {
                     Log.e(TAG, "Unexpected tag: " + tag);
                     return false;
+                }
             }
             return true;
         }