[res] Duplicate AssetManager when changes are needed

When adding new apk paths to the ResourcesImpl's AssetManager,
we used to make it in a non-atomical way - first, the paths get
added without recalculating the internal indices, then later it
updates the indices and applies new configuration. This causes
some lookups to fail until the index recalculation, as apk
asset shift in the internal array, and the indices stop pointing
to the correct apks.

In this change, the asset paths update is now more parametrized:

- When creating a ResourcesImpl we know it's not being used yet,
  so we just update the paths in place, and use the same non-
  atomic, but the fastest, code

- When updating an existing object we instead create a new
  AssetManager object with the new paths, and then use that
  instead of the original one that could be in use

    - this helps to keep the currently running lookups to be
      consistent, as the object can't change under the caller
      and suddenly start returning different resources
    - an added bonus is that now we make sure that the paths
      actually change before running any update operations,
      potentially saving lots of time on heavy locking and JNI
      calls

+ Add a lock when accessing the registered shared libraries

Bug: 362700391
Test: manual, run a test app with a resources getter in a loop
      and verify it doesn't get NotFoundException
Flag: EXEMPT bugfix with no single point of entry
Change-Id: If1890a62b34e504fd87dc27ef9dd27863cc718ee
diff --git a/core/java/android/app/ResourcesManager.java b/core/java/android/app/ResourcesManager.java
index a458b4e..f702b85 100644
--- a/core/java/android/app/ResourcesManager.java
+++ b/core/java/android/app/ResourcesManager.java
@@ -174,22 +174,54 @@
     }
 
     /**
-     * Apply the registered library paths to the passed impl object
-     * @return the hash code for the current version of the registered paths
+     * Apply the registered library paths to the passed AssetManager. If may create a new
+     * AssetManager if any changes are needed and it isn't allowed to reuse the old one.
+     *
+     * @return new AssetManager and the hash code for the current version of the registered paths
      */
-    public int updateResourceImplWithRegisteredLibs(@NonNull ResourcesImpl impl) {
+    public @NonNull Pair<AssetManager, Integer> updateResourceImplAssetsWithRegisteredLibs(
+            @NonNull AssetManager assets, boolean reuseAssets) {
         if (!Flags.registerResourcePaths()) {
-            return 0;
+            return new Pair<>(assets, 0);
         }
 
-        final var collector = new PathCollector(null);
-        final int size = mSharedLibAssetsMap.size();
-        for (int i = 0; i < size; i++) {
-            final var libraryKey = mSharedLibAssetsMap.valueAt(i).getResourcesKey();
-            collector.appendKey(libraryKey);
+        final int size;
+        final PathCollector collector;
+
+        synchronized (mLock) {
+            size = mSharedLibAssetsMap.size();
+            if (assets == AssetManager.getSystem()) {
+                return new Pair<>(assets, size);
+            }
+            collector = new PathCollector(resourcesKeyFromAssets(assets));
+            for (int i = 0; i < size; i++) {
+                final var libraryKey = mSharedLibAssetsMap.valueAt(i).getResourcesKey();
+                collector.appendKey(libraryKey);
+            }
         }
-        impl.getAssets().addPresetApkKeys(extractApkKeys(collector.collectedKey()));
-        return size;
+        if (collector.isSameAsOriginal()) {
+            return new Pair<>(assets, size);
+        }
+        if (reuseAssets) {
+            assets.addPresetApkKeys(extractApkKeys(collector.collectedKey()));
+            return new Pair<>(assets, size);
+        }
+        final var newAssetsBuilder = new AssetManager.Builder();
+        for (final var asset : assets.getApkAssets()) {
+            if (!asset.isForLoader()) {
+                newAssetsBuilder.addApkAssets(asset);
+            }
+        }
+        for (final var key : extractApkKeys(collector.collectedKey())) {
+            try {
+                final var asset = loadApkAssets(key);
+                newAssetsBuilder.addApkAssets(asset);
+            } catch (IOException e) {
+                Log.e(TAG, "Couldn't load assets for key " + key, e);
+            }
+        }
+        assets.getLoaders().forEach(newAssetsBuilder::addLoader);
+        return new Pair<>(newAssetsBuilder.build(), size);
     }
 
     public static class ApkKey {
@@ -624,6 +656,23 @@
         return apkKeys;
     }
 
+    private ResourcesKey resourcesKeyFromAssets(@NonNull AssetManager assets) {
+        final var libs = new ArrayList<String>();
+        final var overlays = new ArrayList<String>();
+        for (final ApkAssets asset : assets.getApkAssets()) {
+            if (asset.isSystem() || asset.isForLoader()) {
+                continue;
+            }
+            if (asset.isOverlay()) {
+                overlays.add(asset.getAssetPath());
+            } else if (asset.isSharedLib()) {
+                libs.add(asset.getAssetPath());
+            }
+        }
+        return new ResourcesKey(null, null, overlays.toArray(new String[0]),
+                libs.toArray(new String[0]), 0, null, null);
+    }
+
     /**
      * Creates an AssetManager from the paths within the ResourcesKey.
      *
@@ -752,7 +801,7 @@
 
         final Configuration config = generateConfig(key);
         final DisplayMetrics displayMetrics = getDisplayMetrics(generateDisplayId(key), daj);
-        final ResourcesImpl impl = new ResourcesImpl(assets, displayMetrics, config, daj);
+        final ResourcesImpl impl = new ResourcesImpl(assets, displayMetrics, config, daj, true);
 
         if (DEBUG) {
             Slog.d(TAG, "- creating impl=" + impl + " with key: " + key);
@@ -1835,31 +1884,32 @@
         for (int i = 0; i < resourcesCount; i++) {
             final WeakReference<Resources> ref = mAllResourceReferences.get(i);
             final Resources r = ref != null ? ref.get() : null;
-            if (r != null) {
-                final ResourcesKey key = updatedResourceKeys.get(r.getImpl());
-                if (key != null) {
-                    final ResourcesImpl impl = findOrCreateResourcesImplForKeyLocked(key);
-                    if (impl == null) {
-                        throw new Resources.NotFoundException("failed to redirect ResourcesImpl");
-                    }
-                    r.setImpl(impl);
-                } else {
-                    // ResourcesKey is null which means the ResourcesImpl could belong to a
-                    // Resources created by application through Resources constructor and was not
-                    // managed by ResourcesManager, so the ResourcesImpl needs to be recreated to
-                    // have shared library asset paths appended if there are any.
-                    if (r.getImpl() != null) {
-                        final ResourcesImpl oldImpl = r.getImpl();
-                        final AssetManager oldAssets = oldImpl.getAssets();
-                        // ResourcesImpl constructor will help to append shared library asset paths.
-                        if (oldAssets != AssetManager.getSystem() && oldAssets.isUpToDate()) {
-                            final ResourcesImpl newImpl = new ResourcesImpl(oldAssets,
-                                    oldImpl.getMetrics(), oldImpl.getConfiguration(),
-                                    oldImpl.getDisplayAdjustments());
+            if (r == null) {
+                continue;
+            }
+            final ResourcesKey key = updatedResourceKeys.get(r.getImpl());
+            if (key != null) {
+                final ResourcesImpl impl = findOrCreateResourcesImplForKeyLocked(key);
+                if (impl == null) {
+                    throw new Resources.NotFoundException("failed to redirect ResourcesImpl");
+                }
+                r.setImpl(impl);
+            } else {
+                // ResourcesKey is null which means the ResourcesImpl could belong to a
+                // Resources created by application through Resources constructor and was not
+                // managed by ResourcesManager, so the ResourcesImpl needs to be recreated to
+                // have shared library asset paths appended if there are any.
+                final ResourcesImpl oldImpl = r.getImpl();
+                if (oldImpl != null) {
+                    final AssetManager oldAssets = oldImpl.getAssets();
+                    // ResourcesImpl constructor will help to append shared library asset paths.
+                    if (oldAssets != AssetManager.getSystem()) {
+                        if (oldAssets.isUpToDate()) {
+                            final ResourcesImpl newImpl = new ResourcesImpl(oldImpl);
                             r.setImpl(newImpl);
                         } else {
-                            Slog.w(TAG, "Skip appending shared library asset paths for the "
-                                    + "Resource as its assets are not up to date.");
+                            Slog.w(TAG, "Skip appending shared library asset paths for "
+                                    + "the Resources as its assets are not up to date.");
                         }
                     }
                 }
diff --git a/core/java/android/content/res/ApkAssets.java b/core/java/android/content/res/ApkAssets.java
index 68b5d78..908999b 100644
--- a/core/java/android/content/res/ApkAssets.java
+++ b/core/java/android/content/res/ApkAssets.java
@@ -124,11 +124,13 @@
 
     @Nullable
     @GuardedBy("this")
-    private final StringBlock mStringBlock;  // null or closed if mNativePtr = 0.
+    private StringBlock mStringBlock;  // null or closed if mNativePtr = 0.
 
     @PropertyFlags
     private final int mFlags;
 
+    private final boolean mIsOverlay;
+
     @Nullable
     private final AssetsProvider mAssets;
 
@@ -302,40 +304,43 @@
 
     private ApkAssets(@FormatType int format, @NonNull String path, @PropertyFlags int flags,
             @Nullable AssetsProvider assets) throws IOException {
+        this(format, flags, assets);
         Objects.requireNonNull(path, "path");
-        mFlags = flags;
         mNativePtr = nativeLoad(format, path, flags, assets);
         mStringBlock = new StringBlock(nativeGetStringBlock(mNativePtr), true /*useSparse*/);
-        mAssets = assets;
     }
 
     private ApkAssets(@FormatType int format, @NonNull FileDescriptor fd,
             @NonNull String friendlyName, @PropertyFlags int flags, @Nullable AssetsProvider assets)
             throws IOException {
+        this(format, flags, assets);
         Objects.requireNonNull(fd, "fd");
         Objects.requireNonNull(friendlyName, "friendlyName");
-        mFlags = flags;
         mNativePtr = nativeLoadFd(format, fd, friendlyName, flags, assets);
         mStringBlock = new StringBlock(nativeGetStringBlock(mNativePtr), true /*useSparse*/);
-        mAssets = assets;
     }
 
     private ApkAssets(@FormatType int format, @NonNull FileDescriptor fd,
             @NonNull String friendlyName, long offset, long length, @PropertyFlags int flags,
             @Nullable AssetsProvider assets) throws IOException {
+        this(format, flags, assets);
         Objects.requireNonNull(fd, "fd");
         Objects.requireNonNull(friendlyName, "friendlyName");
-        mFlags = flags;
         mNativePtr = nativeLoadFdOffsets(format, fd, friendlyName, offset, length, flags, assets);
         mStringBlock = new StringBlock(nativeGetStringBlock(mNativePtr), true /*useSparse*/);
-        mAssets = assets;
     }
 
     private ApkAssets(@PropertyFlags int flags, @Nullable AssetsProvider assets) {
-        mFlags = flags;
+        this(FORMAT_APK, flags, assets);
         mNativePtr = nativeLoadEmpty(flags, assets);
         mStringBlock = null;
+    }
+
+    private ApkAssets(@FormatType int format, @PropertyFlags int flags,
+            @Nullable AssetsProvider assets) {
+        mFlags = flags;
         mAssets = assets;
+        mIsOverlay = format == FORMAT_IDMAP;
     }
 
     @UnsupportedAppUsage
@@ -425,6 +430,18 @@
         }
     }
 
+    public boolean isSystem() {
+        return (mFlags & PROPERTY_SYSTEM) != 0;
+    }
+
+    public boolean isSharedLib() {
+        return (mFlags & PROPERTY_DYNAMIC) != 0;
+    }
+
+    public boolean isOverlay() {
+        return mIsOverlay;
+    }
+
     @Override
     public String toString() {
         return "ApkAssets{path=" + getDebugName() + "}";
diff --git a/core/java/android/content/res/ResourcesImpl.java b/core/java/android/content/res/ResourcesImpl.java
index e6b9342..bcaceb2 100644
--- a/core/java/android/content/res/ResourcesImpl.java
+++ b/core/java/android/content/res/ResourcesImpl.java
@@ -203,9 +203,25 @@
     @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.R, trackingBug = 170729553)
     public ResourcesImpl(@NonNull AssetManager assets, @Nullable DisplayMetrics metrics,
             @Nullable Configuration config, @NonNull DisplayAdjustments displayAdjustments) {
-        mAssets = assets;
-        mAppliedSharedLibsHash =
-                ResourcesManager.getInstance().updateResourceImplWithRegisteredLibs(this);
+        // Don't reuse assets by default as we have no control over whether they're already
+        // inside some other ResourcesImpl.
+        this(assets, metrics, config, displayAdjustments, false);
+    }
+
+    public ResourcesImpl(@NonNull ResourcesImpl orig) {
+        // We know for sure that the other assets are in use, so can't reuse the object here.
+        this(orig.getAssets(), orig.getMetrics(), orig.getConfiguration(),
+                orig.getDisplayAdjustments(), false);
+    }
+
+    public ResourcesImpl(@NonNull AssetManager assets, @Nullable DisplayMetrics metrics,
+            @Nullable Configuration config, @NonNull DisplayAdjustments displayAdjustments,
+            boolean reuseAssets) {
+        final var assetsAndHash =
+                ResourcesManager.getInstance().updateResourceImplAssetsWithRegisteredLibs(assets,
+                        reuseAssets);
+        mAssets = assetsAndHash.first;
+        mAppliedSharedLibsHash = assetsAndHash.second;
         mMetrics.setToDefaults();
         mDisplayAdjustments = displayAdjustments;
         mConfiguration.setToDefaults();