Merge "Fix LazyValue ClassLoader issue" into main

3 files changed, 71 insertions, 19 deletions

diff --git a/core/java/android/os/BaseBundle.java b/core/java/android/os/BaseBundle.java
index 50121278f0e6..8a0adfaede8a 100644
--- a/core/java/android/os/BaseBundle.java
+++ b/core/java/android/os/BaseBundle.java
@@ -45,7 +45,8 @@ import java.util.function.BiFunction;
  * {@link PersistableBundle} subclass.
  */
 @android.ravenwood.annotation.RavenwoodKeepWholeClass
-public class BaseBundle {
+@SuppressWarnings("HiddenSuperclass")
+public class BaseBundle implements Parcel.ClassLoaderProvider {
     /** @hide */
     protected static final String TAG = "Bundle";
     static final boolean DEBUG = false;
@@ -299,8 +300,9 @@ public class BaseBundle {
 
     /**
      * Return the ClassLoader currently associated with this Bundle.
+     * @hide
      */
-    ClassLoader getClassLoader() {
+    public ClassLoader getClassLoader() {
         return mClassLoader;
     }
 
@@ -405,6 +407,9 @@ public class BaseBundle {
             if ((mFlags & Bundle.FLAG_VERIFY_TOKENS_PRESENT) != 0) {
                 Intent.maybeMarkAsMissingCreatorToken(object);
             }
+        } else if (object instanceof Bundle) {
+            Bundle bundle = (Bundle) object;
+            bundle.setClassLoaderSameAsContainerBundleWhenRetrievedFirstTime(this);
         }
         return (clazz != null) ? clazz.cast(object) : (T) object;
     }
@@ -478,7 +483,7 @@ public class BaseBundle {
         int numLazyValues = 0;
         try {
             numLazyValues = parcelledData.readArrayMap(map, count, !parcelledByNative,
-                    /* lazy */ ownsParcel, mClassLoader);
+                    /* lazy */ ownsParcel, this);
         } catch (BadParcelableException e) {
             if (sShouldDefuse) {
                 Log.w(TAG, "Failed to parse Bundle, but defusing quietly", e);
diff --git a/core/java/android/os/Bundle.java b/core/java/android/os/Bundle.java
index 819d58d9f059..55bfd451d97a 100644
--- a/core/java/android/os/Bundle.java
+++ b/core/java/android/os/Bundle.java
@@ -141,6 +141,8 @@ public final class Bundle extends BaseBundle implements Cloneable, Parcelable {
         STRIPPED.putInt("STRIPPED", 1);
     }
 
+    private boolean isFirstRetrievedFromABundle = false;
+
     /**
      * Constructs a new, empty Bundle.
      */
@@ -1020,7 +1022,9 @@ public final class Bundle extends BaseBundle implements Cloneable, Parcelable {
             return null;
         }
         try {
-            return (Bundle) o;
+            Bundle bundle = (Bundle) o;
+            bundle.setClassLoaderSameAsContainerBundleWhenRetrievedFirstTime(this);
+            return bundle;
         } catch (ClassCastException e) {
             typeWarning(key, o, "Bundle", e);
             return null;
@@ -1028,6 +1032,21 @@ public final class Bundle extends BaseBundle implements Cloneable, Parcelable {
     }
 
     /**
+     * Set the ClassLoader of a bundle to its container bundle. This is necessary so that when a
+     * bundle's ClassLoader is changed, it can be propagated to its children. Do this only when it
+     * is retrieved from the container bundle first time though. Once it is accessed outside of its
+     * container, its ClassLoader should no longer be changed by its container anymore.
+     *
+     * @param containerBundle the bundle this bundle is retrieved from.
+     */
+    void setClassLoaderSameAsContainerBundleWhenRetrievedFirstTime(BaseBundle containerBundle) {
+        if (!isFirstRetrievedFromABundle) {
+            setClassLoader(containerBundle.getClassLoader());
+            isFirstRetrievedFromABundle = true;
+        }
+    }
+
+    /**
      * Returns the value associated with the given key, or {@code null} if
      * no mapping of the desired type exists for the given key or a {@code null}
      * value is explicitly associated with the key.
diff --git a/core/java/android/os/Parcel.java b/core/java/android/os/Parcel.java
index cf473ec9c3ea..4bbc61c9f115 100644
--- a/core/java/android/os/Parcel.java
+++ b/core/java/android/os/Parcel.java
@@ -4650,7 +4650,7 @@ public final class Parcel {
      * @hide
      */
     @Nullable
-    public Object readLazyValue(@Nullable ClassLoader loader) {
+    private Object readLazyValue(@Nullable ClassLoaderProvider loaderProvider) {
         int start = dataPosition();
         int type = readInt();
         if (isLengthPrefixed(type)) {
@@ -4661,12 +4661,17 @@ public final class Parcel {
             int end = MathUtils.addOrThrow(dataPosition(), objectLength);
             int valueLength = end - start;
             setDataPosition(end);
-            return new LazyValue(this, start, valueLength, type, loader);
+            return new LazyValue(this, start, valueLength, type, loaderProvider);
         } else {
-            return readValue(type, loader, /* clazz */ null);
+            return readValue(type, getClassLoader(loaderProvider), /* clazz */ null);
         }
     }
 
+    @Nullable
+    private static ClassLoader getClassLoader(@Nullable ClassLoaderProvider loaderProvider) {
+        return loaderProvider == null ? null : loaderProvider.getClassLoader();
+    }
+
 
     private static final class LazyValue implements BiFunction<Class<?>, Class<?>[], Object> {
         /**
@@ -4680,7 +4685,12 @@ public final class Parcel {
         private final int mPosition;
         private final int mLength;
         private final int mType;
-        @Nullable private final ClassLoader mLoader;
+        // this member is set when a bundle that includes a LazyValue is unparceled. But it is used
+        // when apply method is called. Between these 2 events, the bundle's ClassLoader could have
+        // changed. Let the bundle be a ClassLoaderProvider allows the bundle provides its current
+        // ClassLoader at the time apply method is called.
+        @NonNull
+        private final ClassLoaderProvider mLoaderProvider;
         @Nullable private Object mObject;
 
         /**
@@ -4691,12 +4701,13 @@ public final class Parcel {
          */
         @Nullable private volatile Parcel mSource;
 
-        LazyValue(Parcel source, int position, int length, int type, @Nullable ClassLoader loader) {
+        LazyValue(Parcel source, int position, int length, int type,
+                @NonNull ClassLoaderProvider loaderProvider) {
             mSource = requireNonNull(source);
             mPosition = position;
             mLength = length;
             mType = type;
-            mLoader = loader;
+            mLoaderProvider = loaderProvider;
         }
 
         @Override
@@ -4709,7 +4720,8 @@ public final class Parcel {
                         int restore = source.dataPosition();
                         try {
                             source.setDataPosition(mPosition);
-                            mObject = source.readValue(mLoader, clazz, itemTypes);
+                            mObject = source.readValue(mLoaderProvider.getClassLoader(), clazz,
+                                    itemTypes);
                         } finally {
                             source.setDataPosition(restore);
                         }
@@ -4782,7 +4794,8 @@ public final class Parcel {
                 return Objects.equals(mObject, value.mObject);
             }
             // Better safely fail here since this could mean we get different objects.
-            if (!Objects.equals(mLoader, value.mLoader)) {
+            if (!Objects.equals(mLoaderProvider.getClassLoader(),
+                    value.mLoaderProvider.getClassLoader())) {
                 return false;
             }
             // Otherwise compare metadata prior to comparing payload.
@@ -4796,10 +4809,24 @@ public final class Parcel {
         @Override
         public int hashCode() {
             // Accessing mSource first to provide memory barrier for mObject
-            return Objects.hash(mSource == null, mObject, mLoader, mType, mLength);
+            return Objects.hash(mSource == null, mObject, mLoaderProvider.getClassLoader(), mType,
+                    mLength);
         }
     }
 
+    /**
+     * Provides a ClassLoader.
+     * @hide
+     */
+    public interface ClassLoaderProvider {
+        /**
+         * Returns a ClassLoader.
+         *
+         * @return ClassLoader
+         */
+        ClassLoader getClassLoader();
+    }
+
     /** Same as {@link #readValue(ClassLoader, Class, Class[])} without any item types. */
     private <T> T readValue(int type, @Nullable ClassLoader loader, @Nullable Class<T> clazz) {
         // Avoids allocating Class[0] array
@@ -5537,8 +5564,8 @@ public final class Parcel {
     }
 
     private void readArrayMapInternal(@NonNull ArrayMap<? super String, Object> outVal,
-            int size, @Nullable ClassLoader loader) {
-        readArrayMap(outVal, size, /* sorted */ true, /* lazy */ false, loader);
+            int size, @Nullable ClassLoaderProvider loaderProvider) {
+        readArrayMap(outVal, size, /* sorted */ true, /* lazy */ false, loaderProvider);
     }
 
     /**
@@ -5552,11 +5579,12 @@ public final class Parcel {
      * @hide
      */
     int readArrayMap(ArrayMap<? super String, Object> map, int size, boolean sorted,
-            boolean lazy, @Nullable ClassLoader loader) {
+            boolean lazy, @Nullable ClassLoaderProvider loaderProvider) {
         int lazyValues = 0;
         while (size > 0) {
             String key = readString();
-            Object value = (lazy) ? readLazyValue(loader) : readValue(loader);
+            Object value = (lazy) ? readLazyValue(loaderProvider) : readValue(
+                    getClassLoader(loaderProvider));
             if (value instanceof LazyValue) {
                 lazyValues++;
             }
@@ -5578,12 +5606,12 @@ public final class Parcel {
      */
     @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.R, trackingBug = 170729553)
     public void readArrayMap(@NonNull ArrayMap<? super String, Object> outVal,
-            @Nullable ClassLoader loader) {
+            @Nullable ClassLoaderProvider loaderProvider) {
         final int N = readInt();
         if (N < 0) {
             return;
         }
-        readArrayMapInternal(outVal, N, loader);
+        readArrayMapInternal(outVal, N, loaderProvider);
     }
 
     /**