Fix LazyValue ClassLoader issue

LazyValue's class loader is set when the bundle that contains it
is first unparceled. This is when the LazyValue object is created.
But between when the container bundle is unparceled and when the
LazyValue is unparceled (LazyValue.apply() is called), the bundle's
ClassLoader could have changed. LazyValue unparcel should use the
new ClassLoader.

To fix this issue, replace the ClassLoader member of the LazyValue
with a ClassLoaderProvider, which could simply point back to the
containe bundle as the ClassLoaderProvider. Thus, updates of the
ClassLoader of the container bundle is always effective when the
LazyValue.apply is called.

Also, when a bundle is retrieved from a bundle, the parent's
classLoader needs to be set to the child. This is because if a
bundle contains another nested bundle, when the top level bundle
is unparceled, the nested bundle's classLoader is set as the same
as the top level bundle's

Bug: 382633789
Test: manual test
Flag: EXEMPT bug fix
Change-Id: Id775366506d101b618b214837df1a61f1db380d7

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);
     }
 
     /**