Revert^2 "Fix LazyValue ClassLoader issue"

This reverts commit 2fd48af4c5f2cec41b8296b30836eed8c58408cb.

Reason for revert: The revert broke Settings Add Account again.
The original CL is to fix GmsCore side load module issue (Add
Account and Wallet tap to pay) where the side loaded module relies
on intent.setExtrasClassLoader to their modules classLoader and
then unparcel values from the intent's extra bundle and use the
new classLoader. But since collectExtraIntentKeys and
checkCreatorToken now unparcel the values before the side loaded
module retrieve their values, the original behavior would set
the classLoader when they are first unparceled and wrong. This
CL's original version fixed the issue so that retrieve value uses
the most current classLoader of the bundle, instead of when the
bundle first unparceled. But it missed a case where
Bundle.putAll(bundle) method is used. In case of
b1.putAll(b2), b2's nested bundles classLoaders are lost.see
https://b.corp.google.com/issues/387716166#comment39 for detailed
analysis of this case. That's why the CL was reverted.
But revert broke Add Account feature again.
So revert^2 and fix the putAll issue here by setting
isFristRetrivedFromBundle to true for b2's nested bundles so that
b2's nested bundles' classLoaders won't be changed again.

Bug: 393185674
Bug: 387716166
Bug: 382633789

Change-Id: Ifa63d73e2b993708d8f14a356c4f53fba91c8698

4 files changed, 131 insertions, 20 deletions

diff --git a/core/java/android/os/BaseBundle.java b/core/java/android/os/BaseBundle.java
index e79b2e7becce..26044545b8d1 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;
@@ -311,8 +312,9 @@ public class BaseBundle {
 
     /**
      * Return the ClassLoader currently associated with this Bundle.
+     * @hide
      */
-    ClassLoader getClassLoader() {
+    public ClassLoader getClassLoader() {
         return mClassLoader;
     }
 
@@ -426,6 +428,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;
     }
@@ -499,7 +504,7 @@ public class BaseBundle {
         int[] numLazyValues = new int[]{0};
         try {
             parcelledData.readArrayMap(map, count, !parcelledByNative,
-                    /* lazy */ ownsParcel, mClassLoader, numLazyValues);
+                    /* lazy */ ownsParcel, this, numLazyValues);
         } 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 a24dc5739b7e..c0591e6899b6 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.
      */
@@ -382,7 +384,15 @@ public final class Bundle extends BaseBundle implements Cloneable, Parcelable {
         bundle.unparcel();
         mOwnsLazyValues = false;
         bundle.mOwnsLazyValues = false;
-        mMap.putAll(bundle.mMap);
+        int N = bundle.mMap.size();
+        for (int i = 0; i < N; i++) {
+            String key = bundle.mMap.keyAt(i);
+            Object value = bundle.mMap.valueAt(i);
+            if (value instanceof Bundle) {
+                ((Bundle) value).isFirstRetrievedFromABundle = true;
+            }
+            mMap.put(key, value);
+        }
 
         // FD and Binders state is now known if and only if both bundles already knew
         if ((bundle.mFlags & FLAG_HAS_FDS) != 0) {
@@ -592,6 +602,8 @@ public final class Bundle extends BaseBundle implements Cloneable, Parcelable {
         mFlags &= ~FLAG_HAS_BINDERS_KNOWN;
         if (intentClass != null && intentClass.isInstance(value)) {
             setHasIntent(true);
+        } else if (value instanceof Bundle) {
+            ((Bundle) value).isFirstRetrievedFromABundle = true;
         }
     }
 
@@ -793,6 +805,9 @@ public final class Bundle extends BaseBundle implements Cloneable, Parcelable {
      */
     public void putBundle(@Nullable String key, @Nullable Bundle value) {
         unparcel();
+        if (value != null) {
+            value.isFirstRetrievedFromABundle = true;
+        }
         mMap.put(key, value);
     }
 
@@ -1020,7 +1035,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 +1045,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 e58934746c14..49d3f06eb80f 100644
--- a/core/java/android/os/Parcel.java
+++ b/core/java/android/os/Parcel.java
@@ -46,6 +46,7 @@ import android.util.SparseBooleanArray;
 import android.util.SparseIntArray;
 
 import com.android.internal.annotations.GuardedBy;
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.ArrayUtils;
 
 import dalvik.annotation.optimization.CriticalNative;
@@ -4661,7 +4662,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)) {
@@ -4672,12 +4673,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> {
         /**
@@ -4691,7 +4697,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;
 
         /**
@@ -4702,12 +4713,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
@@ -4720,7 +4732,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);
                         }
@@ -4758,6 +4771,12 @@ public final class Parcel {
             return Parcel.hasFileDescriptors(mObject);
         }
 
+        /** @hide */
+        @VisibleForTesting
+        public ClassLoader getClassLoader() {
+            return mLoaderProvider.getClassLoader();
+        }
+
         @Override
         public String toString() {
             return (mSource != null)
@@ -4793,7 +4812,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.
@@ -4807,10 +4827,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
@@ -5551,8 +5585,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, null);
+            int size, @Nullable ClassLoaderProvider loaderProvider) {
+        readArrayMap(outVal, size, /* sorted */ true, /* lazy */ false, loaderProvider, null);
     }
 
     /**
@@ -5566,11 +5600,12 @@ public final class Parcel {
      * @hide
      */
     void readArrayMap(ArrayMap<? super String, Object> map, int size, boolean sorted,
-            boolean lazy, @Nullable ClassLoader loader, int[] lazyValueCount) {
+            boolean lazy, @Nullable ClassLoaderProvider loaderProvider, int[] lazyValueCount) {
         ensureWithinMemoryLimit(SIZE_COMPLEX_TYPE, size);
         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) {
                 lazyValueCount[0]++;
             }
@@ -5591,12 +5626,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);
     }
 
     /**
diff --git a/core/tests/coretests/src/android/content/IntentTest.java b/core/tests/coretests/src/android/content/IntentTest.java
index fdfb0c34cdeb..fa1948d9786c 100644
--- a/core/tests/coretests/src/android/content/IntentTest.java
+++ b/core/tests/coretests/src/android/content/IntentTest.java
@@ -23,6 +23,7 @@ import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
 import android.os.Binder;
+import android.os.Bundle;
 import android.os.IBinder;
 import android.os.Parcel;
 import android.platform.test.annotations.Presubmit;
@@ -238,4 +239,42 @@ public class IntentTest {
         // Not all keys from intent are kept - clip data keys are dropped.
         assertFalse(intent.getExtraIntentKeys().containsAll(originalIntentKeys));
     }
+
+    @Test
+    public void testSetIntentExtrasClassLoaderEffectiveAfterExtraBundleUnparcel() {
+        Intent intent = new Intent();
+        intent.putExtra("bundle", new Bundle());
+
+        final Parcel parcel = Parcel.obtain();
+        intent.writeToParcel(parcel, 0);
+        parcel.setDataPosition(0);
+        final Intent target = new Intent();
+        target.readFromParcel(parcel);
+        target.collectExtraIntentKeys();
+        ClassLoader cl = new ClassLoader() {
+        };
+        target.setExtrasClassLoader(cl);
+        assertThat(target.getBundleExtra("bundle").getClassLoader()).isEqualTo(cl);
+    }
+
+    @Test
+    public void testBundlePutAllClassLoader() {
+        Intent intent = new Intent();
+        Bundle b1 = new Bundle();
+        b1.putBundle("bundle", new Bundle());
+        intent.putExtra("b1", b1);
+        final Parcel parcel = Parcel.obtain();
+        intent.writeToParcel(parcel, 0);
+        parcel.setDataPosition(0);
+        final Intent target = new Intent();
+        target.readFromParcel(parcel);
+
+        ClassLoader cl = new ClassLoader() {
+        };
+        target.setExtrasClassLoader(cl);
+        Bundle b2 = new Bundle();
+        b2.putAll(target.getBundleExtra("b1"));
+        assertThat(b2.getBundle("bundle").getClassLoader()).isEqualTo(cl);
+    }
+
 }