Dedup original dex file for classes transformed on first load.

We set the originalDexFile field to a native pointer to the
art::DexFile (as a boxed Long). In the case where there were no
non-can_retransform_classes agents this DexFile will be owned by the
transformed classes ClassLoader and is guaranteed to be valid for as
long as the Class is. This means we can simply keep track of that
pointer to hold onto the original Dex file.

Test: ./test.py --host -j40
Bug: 31455788
Change-Id: I0e5a74f20f86d32dc73babe64b81469b3155c4dd

3 files changed, 95 insertions, 11 deletions

diff --git a/runtime/openjdkjvmti/ti_class.cc b/runtime/openjdkjvmti/ti_class.cc
index 907ab0252d..4cdb148ff4 100644
--- a/runtime/openjdkjvmti/ti_class.cc
+++ b/runtime/openjdkjvmti/ti_class.cc
@@ -56,6 +56,8 @@
 #include "mirror/object_reference.h"
 #include "mirror/object-inl.h"
 #include "mirror/reference.h"
+#include "primitive.h"
+#include "reflection.h"
 #include "runtime.h"
 #include "runtime_callbacks.h"
 #include "ScopedLocalRef.h"
@@ -231,14 +233,22 @@ struct ClassCallback : public art::ClassLoadCallback {
       }
 
       // Allocate the byte array to store the dex file bytes in.
-      art::Handle<art::mirror::ByteArray> arr(hs.NewHandle(
-          art::mirror::ByteArray::AllocateAndFill(
-              self,
-              reinterpret_cast<const signed char*>(post_no_redefine_dex_data),
-              post_no_redefine_len)));
+      art::MutableHandle<art::mirror::Object> arr(hs.NewHandle<art::mirror::Object>(nullptr));
+      if (post_no_redefine_dex_data == dex_file_copy->Begin() && name != "java/lang/Long") {
+        // we didn't have any non-retransformable agents. We can just cache a pointer to the
+        // initial_dex_file. It will be kept live by the class_loader.
+        jlong dex_ptr = reinterpret_cast<uintptr_t>(&initial_dex_file);
+        art::JValue val;
+        val.SetJ(dex_ptr);
+        arr.Assign(art::BoxPrimitive(art::Primitive::kPrimLong, val));
+      } else {
+        arr.Assign(art::mirror::ByteArray::AllocateAndFill(
+            self,
+            reinterpret_cast<const signed char*>(post_no_redefine_dex_data),
+            post_no_redefine_len));
+      }
       if (arr.IsNull()) {
-        LOG(WARNING) << "Unable to allocate byte array for initial dex-file bytes. Aborting "
-                     << "transformation";
+        LOG(WARNING) << "Unable to allocate memory for initial dex-file. Aborting transformation";
         self->AssertPendingOOMException();
         return;
       }
diff --git a/runtime/openjdkjvmti/transform.cc b/runtime/openjdkjvmti/transform.cc
index a1883b414c..8e38a364ff 100644
--- a/runtime/openjdkjvmti/transform.cc
+++ b/runtime/openjdkjvmti/transform.cc
@@ -43,6 +43,7 @@
 #include "gc_root-inl.h"
 #include "globals.h"
 #include "jni_env_ext-inl.h"
+#include "jvalue.h"
 #include "jvmti.h"
 #include "linear_alloc.h"
 #include "mem_map.h"
@@ -52,6 +53,7 @@
 #include "mirror/class_loader-inl.h"
 #include "mirror/string-inl.h"
 #include "oat_file.h"
+#include "reflection.h"
 #include "scoped_thread_state_change-inl.h"
 #include "stack.h"
 #include "thread_list.h"
@@ -167,18 +169,27 @@ jvmtiError Transformer::GetDexDataForRetransformation(ArtJvmTiEnv* env,
             reinterpret_cast<const unsigned char*>(orig_dex_bytes->GetData()),
             *dex_data_len,
             /*out*/dex_data);
-      } else {
-        DCHECK(orig_dex->IsDexCache());
+      } else if (orig_dex->IsDexCache()) {
         dex_file = orig_dex->AsDexCache()->GetDexFile();
-        *dex_data_len = static_cast<jint>(dex_file->Size());
+      } else {
+        DCHECK_EQ(orig_dex->GetClass()->GetPrimitiveType(), art::Primitive::kPrimLong);
+        art::ObjPtr<art::mirror::Class> prim_long_class(
+            art::Runtime::Current()->GetClassLinker()->GetClassRoot(
+                art::ClassLinker::kPrimitiveLong));
+        art::JValue val;
+        if (!art::UnboxPrimitiveForResult(orig_dex.Get(), prim_long_class, &val)) {
+          // This should never happen.
+          return ERR(INTERNAL);
+        }
+        dex_file = reinterpret_cast<const art::DexFile*>(static_cast<uintptr_t>(val.GetJ()));
       }
     }
   }
   if (dex_file == nullptr) {
     dex_file = &klass->GetDexFile();
-    *dex_data_len = static_cast<jint>(dex_file->Size());
   }
   std::unique_ptr<FixedUpDexFile> fixed_dex_file(FixedUpDexFile::Create(*dex_file));
+  *dex_data_len = static_cast<jint>(fixed_dex_file->Size());
   return CopyDataIntoJvmtiBuffer(env,
                                  fixed_dex_file->Begin(),
                                  fixed_dex_file->Size(),
diff --git a/test/981-dedup-original-dex/src/Main.java b/test/981-dedup-original-dex/src/Main.java
index cd3f007532..1e063cf626 100644
--- a/test/981-dedup-original-dex/src/Main.java
+++ b/test/981-dedup-original-dex/src/Main.java
@@ -16,8 +16,10 @@
 
 import java.lang.reflect.Field;
 import java.util.Base64;
+import java.nio.ByteBuffer;
 
 import dalvik.system.ClassExt;
+import dalvik.system.InMemoryDexClassLoader;
 
 public class Main {
 
@@ -67,6 +69,53 @@ public class Main {
     "AAQAAADgAAAABgAAAAEAAAAAAQAAASAAAAIAAAAgAQAAARAAAAEAAABcAQAAAiAAAA4AAABiAQAA" +
     "AyAAAAIAAAAWAgAAACAAAAEAAAAhAgAAABAAAAEAAAAwAgAA");
 
+
+  /**
+   * base64 encoded class/dex file for
+   * class Transform3 {
+   *   public void sayHi() {
+   *    System.out.println("hello3");
+   *   }
+   * }
+   */
+  private static final byte[] DEX_BYTES_3_INITIAL = Base64.getDecoder().decode(
+    "ZGV4CjAzNQC2W2fBsAeLNAwWYlG8FVigzfsV7nBWITzQAgAAcAAAAHhWNBIAAAAAAAAAADACAAAO" +
+    "AAAAcAAAAAYAAACoAAAAAgAAAMAAAAABAAAA2AAAAAQAAADgAAAAAQAAAAABAACwAQAAIAEAAGIB" +
+    "AABqAQAAeAEAAI8BAACjAQAAtwEAAMsBAADcAQAA3wEAAOMBAAD3AQAA/wEAAAQCAAANAgAAAQAA" +
+    "AAIAAAADAAAABAAAAAUAAAAHAAAABwAAAAUAAAAAAAAACAAAAAUAAABcAQAABAABAAsAAAAAAAAA" +
+    "AAAAAAAAAAANAAAAAQABAAwAAAACAAAAAAAAAAAAAAAAAAAAAgAAAAAAAAAGAAAAAAAAAB8CAAAA" +
+    "AAAAAQABAAEAAAAUAgAABAAAAHAQAwAAAA4AAwABAAIAAAAZAgAACQAAAGIAAAAbAQoAAABuIAIA" +
+    "EAAOAAAAAQAAAAMABjxpbml0PgAMTFRyYW5zZm9ybTM7ABVMamF2YS9pby9QcmludFN0cmVhbTsA" +
+    "EkxqYXZhL2xhbmcvT2JqZWN0OwASTGphdmEvbGFuZy9TdHJpbmc7ABJMamF2YS9sYW5nL1N5c3Rl" +
+    "bTsAD1RyYW5zZm9ybTMuamF2YQABVgACVkwAEmVtaXR0ZXI6IGphY2stNC4zMAAGaGVsbG8zAANv" +
+    "dXQAB3ByaW50bG4ABXNheUhpAAIABw4ABAAHDocAAAABAQCAgASgAgEBuAIAAAANAAAAAAAAAAEA" +
+    "AAAAAAAAAQAAAA4AAABwAAAAAgAAAAYAAACoAAAAAwAAAAIAAADAAAAABAAAAAEAAADYAAAABQAA" +
+    "AAQAAADgAAAABgAAAAEAAAAAAQAAASAAAAIAAAAgAQAAARAAAAEAAABcAQAAAiAAAA4AAABiAQAA" +
+    "AyAAAAIAAAAUAgAAACAAAAEAAAAfAgAAABAAAAEAAAAwAgAA");
+
+  /**
+   * base64 encoded class/dex file for
+   * class Transform3 {
+   *   public void sayHi() {
+   *    System.out.println("Goodbye3");
+   *   }
+   * }
+   */
+  private static final byte[] DEX_BYTES_3_FINAL = Base64.getDecoder().decode(
+    "ZGV4CjAzNQBAXE5GthgMydaFBuinf+ZBfXcBYIw2UlXQAgAAcAAAAHhWNBIAAAAAAAAAADACAAAO" +
+    "AAAAcAAAAAYAAACoAAAAAgAAAMAAAAABAAAA2AAAAAQAAADgAAAAAQAAAAABAACwAQAAIAEAAGIB" +
+    "AABqAQAAdAEAAIIBAACZAQAArQEAAMEBAADVAQAA5gEAAOkBAADtAQAAAQIAAAYCAAAPAgAAAgAA" +
+    "AAMAAAAEAAAABQAAAAYAAAAIAAAACAAAAAUAAAAAAAAACQAAAAUAAABcAQAABAABAAsAAAAAAAAA" +
+    "AAAAAAAAAAANAAAAAQABAAwAAAACAAAAAAAAAAAAAAAAAAAAAgAAAAAAAAAHAAAAAAAAACECAAAA" +
+    "AAAAAQABAAEAAAAWAgAABAAAAHAQAwAAAA4AAwABAAIAAAAbAgAACQAAAGIAAAAbAQEAAABuIAIA" +
+    "EAAOAAAAAQAAAAMABjxpbml0PgAIR29vZGJ5ZTMADExUcmFuc2Zvcm0zOwAVTGphdmEvaW8vUHJp" +
+    "bnRTdHJlYW07ABJMamF2YS9sYW5nL09iamVjdDsAEkxqYXZhL2xhbmcvU3RyaW5nOwASTGphdmEv" +
+    "bGFuZy9TeXN0ZW07AA9UcmFuc2Zvcm0zLmphdmEAAVYAAlZMABJlbWl0dGVyOiBqYWNrLTQuMzAA" +
+    "A291dAAHcHJpbnRsbgAFc2F5SGkAAgAHDgAEAAcOhwAAAAEBAICABKACAQG4AgANAAAAAAAAAAEA" +
+    "AAAAAAAAAQAAAA4AAABwAAAAAgAAAAYAAACoAAAAAwAAAAIAAADAAAAABAAAAAEAAADYAAAABQAA" +
+    "AAQAAADgAAAABgAAAAEAAAAAAQAAASAAAAIAAAAgAQAAARAAAAEAAABcAQAAAiAAAA4AAABiAQAA" +
+    "AyAAAAIAAAAWAgAAACAAAAEAAAAhAgAAABAAAAEAAAAwAgAA");
+
   public static void main(String[] args) {
     try {
       doTest();
@@ -125,6 +174,20 @@ public class Main {
     enableCommonRetransformation(false);
     doCommonClassRedefinition(Transform.class, new byte[0], DEX_BYTES_1);
     assertSame((new byte[0]).getClass(), getOriginalDexFile(t1.getClass()).getClass());
+
+    // Check we don't have anything if we don't have any originalDexFile if the onload
+    // transformation doesn't do anything.
+    enableCommonRetransformation(true);
+    Class<?> transform3Class = new InMemoryDexClassLoader(
+        ByteBuffer.wrap(DEX_BYTES_3_INITIAL), Main.class.getClassLoader()).loadClass("Transform3");
+    assertSame(null, getOriginalDexFile(transform3Class));
+
+    // Check that we end up with a java.lang.Long pointer if we do an 'on-load' redefinition.
+    addCommonTransformationResult("Transform3", new byte[0], DEX_BYTES_3_FINAL);
+    enableCommonRetransformation(true);
+    Class<?> transform3ClassTransformed = new InMemoryDexClassLoader(
+        ByteBuffer.wrap(DEX_BYTES_3_INITIAL), Main.class.getClassLoader()).loadClass("Transform3");
+    assertSame(Long.class, getOriginalDexFile(transform3ClassTransformed).getClass());
   }
 
   // Transforms the class