Implement class-pre-define time redefinition.

This allows one to redefine classes as they are being loaded without
restriction as to the types of transformations that may be included.
For example one is allowed to add additional methods or fields to the
class being defined.

Bug: 31684920
Test: mma -j40 test-art-host
Change-Id: I671ee13444c05f28424e727fd80d9d46c78ca287
diff --git a/runtime/openjdkjvmti/art_jvmti.h b/runtime/openjdkjvmti/art_jvmti.h
index 256c3a6..91353e2 100644
--- a/runtime/openjdkjvmti/art_jvmti.h
+++ b/runtime/openjdkjvmti/art_jvmti.h
@@ -114,6 +114,21 @@
 }
 
 ALWAYS_INLINE
+static inline jvmtiError CopyDataIntoJvmtiBuffer(ArtJvmTiEnv* env,
+                                                 const unsigned char* source,
+                                                 jint len,
+                                                 /*out*/unsigned char** dest) {
+  jvmtiError res = env->Allocate(len, dest);
+  if (res != OK) {
+    return res;
+  }
+  memcpy(reinterpret_cast<void*>(*dest),
+         reinterpret_cast<const void*>(source),
+         len);
+  return OK;
+}
+
+ALWAYS_INLINE
 static inline jvmtiError CopyString(jvmtiEnv* env, const char* src, unsigned char** copy) {
   size_t len = strlen(src) + 1;
   unsigned char* buf;
diff --git a/runtime/openjdkjvmti/events-inl.h b/runtime/openjdkjvmti/events-inl.h
index 655a53a..4f5eb0c 100644
--- a/runtime/openjdkjvmti/events-inl.h
+++ b/runtime/openjdkjvmti/events-inl.h
@@ -131,8 +131,8 @@
   ArtJvmTiEnv* last_env = nullptr;
   for (ArtJvmTiEnv* env : envs) {
     if (ShouldDispatch<kEvent>(env, thread)) {
-      jint new_len;
-      unsigned char* new_data;
+      jint new_len = 0;
+      unsigned char* new_data = nullptr;
       auto callback = impl::GetCallback<kEvent>(env);
       callback(env,
                jnienv,
diff --git a/runtime/openjdkjvmti/ti_class.cc b/runtime/openjdkjvmti/ti_class.cc
index 6dbab86..b6de592 100644
--- a/runtime/openjdkjvmti/ti_class.cc
+++ b/runtime/openjdkjvmti/ti_class.cc
@@ -31,6 +31,8 @@
 
 #include "ti_class.h"
 
+#include "android-base/stringprintf.h"
+
 #include <mutex>
 #include <unordered_set>
 
@@ -38,20 +40,208 @@
 #include "base/macros.h"
 #include "class_table-inl.h"
 #include "class_linker.h"
+#include "common_throws.h"
 #include "events-inl.h"
 #include "handle.h"
 #include "jni_env_ext-inl.h"
 #include "jni_internal.h"
+#include "mirror/array-inl.h"
+#include "mirror/class-inl.h"
+#include "mirror/class_ext.h"
 #include "runtime.h"
 #include "runtime_callbacks.h"
 #include "ScopedLocalRef.h"
 #include "scoped_thread_state_change-inl.h"
 #include "thread-inl.h"
 #include "thread_list.h"
+#include "ti_redefine.h"
+#include "utils.h"
 
 namespace openjdkjvmti {
 
+using android::base::StringPrintf;
+
+static std::unique_ptr<const art::DexFile> MakeSingleDexFile(art::Thread* self,
+                                                             const char* descriptor,
+                                                             const std::string& orig_location,
+                                                             jint final_len,
+                                                             const unsigned char* final_dex_data)
+      REQUIRES_SHARED(art::Locks::mutator_lock_) {
+  // Make the mmap
+  std::string error_msg;
+  std::unique_ptr<art::MemMap> map(Redefiner::MoveDataToMemMap(orig_location,
+                                                               final_len,
+                                                               final_dex_data,
+                                                               &error_msg));
+  if (map.get() == nullptr) {
+    LOG(WARNING) << "Unable to allocate mmap for redefined dex file! Error was: " << error_msg;
+    self->ThrowOutOfMemoryError(StringPrintf(
+        "Unable to allocate dex file for transformation of %s", descriptor).c_str());
+    return nullptr;
+  }
+
+  // Make a dex-file
+  if (map->Size() < sizeof(art::DexFile::Header)) {
+    LOG(WARNING) << "Could not read dex file header because dex_data was too short";
+    art::ThrowClassFormatError(nullptr,
+                               "Unable to read transformed dex file of %s",
+                               descriptor);
+    return nullptr;
+  }
+  uint32_t checksum = reinterpret_cast<const art::DexFile::Header*>(map->Begin())->checksum_;
+  std::unique_ptr<const art::DexFile> dex_file(art::DexFile::Open(map->GetName(),
+                                                                  checksum,
+                                                                  std::move(map),
+                                                                  /*verify*/true,
+                                                                  /*verify_checksum*/true,
+                                                                  &error_msg));
+  if (dex_file.get() == nullptr) {
+    LOG(WARNING) << "Unable to load modified dex file for " << descriptor << ": " << error_msg;
+    art::ThrowClassFormatError(nullptr,
+                               "Unable to read transformed dex file of %s because %s",
+                               descriptor,
+                               error_msg.c_str());
+    return nullptr;
+  }
+  if (dex_file->NumClassDefs() != 1) {
+    LOG(WARNING) << "Dex file contains more than 1 class_def. Ignoring.";
+    // TODO Throw some other sort of error here maybe?
+    art::ThrowClassFormatError(
+        nullptr,
+        "Unable to use transformed dex file of %s because it contained too many classes",
+        descriptor);
+    return nullptr;
+  }
+  return dex_file;
+}
+
 struct ClassCallback : public art::ClassLoadCallback {
+  void ClassPreDefine(const char* descriptor,
+                      art::Handle<art::mirror::Class> klass,
+                      art::Handle<art::mirror::ClassLoader> class_loader,
+                      const art::DexFile& initial_dex_file,
+                      const art::DexFile::ClassDef& initial_class_def ATTRIBUTE_UNUSED,
+                      /*out*/art::DexFile const** final_dex_file,
+                      /*out*/art::DexFile::ClassDef const** final_class_def)
+      OVERRIDE REQUIRES_SHARED(art::Locks::mutator_lock_) {
+    bool is_enabled =
+        event_handler->IsEventEnabledAnywhere(ArtJvmtiEvent::kClassFileLoadHookRetransformable) ||
+        event_handler->IsEventEnabledAnywhere(ArtJvmtiEvent::kClassFileLoadHookNonRetransformable);
+    if (!is_enabled) {
+      return;
+    }
+    if (descriptor[0] != 'L') {
+      // It is a primitive or array. Just return
+      return;
+    }
+    std::string name(art::PrettyDescriptor(descriptor));
+
+    art::Thread* self = art::Thread::Current();
+    art::JNIEnvExt* env = self->GetJniEnv();
+    ScopedLocalRef<jobject> loader(
+        env, class_loader.IsNull() ? nullptr : env->AddLocalReference<jobject>(class_loader.Get()));
+    // Go back to native.
+    art::ScopedThreadSuspension sts(self, art::ThreadState::kNative);
+    // Call all Non-retransformable agents.
+    jint post_no_redefine_len = 0;
+    unsigned char* post_no_redefine_dex_data = nullptr;
+    std::unique_ptr<const unsigned char> post_no_redefine_unique_ptr(nullptr);
+    event_handler->DispatchEvent<ArtJvmtiEvent::kClassFileLoadHookNonRetransformable>(
+        self,
+        static_cast<JNIEnv*>(env),
+        static_cast<jclass>(nullptr),  // The class doesn't really exist yet so send null.
+        loader.get(),
+        name.c_str(),
+        static_cast<jobject>(nullptr),  // Android doesn't seem to have protection domains
+        static_cast<jint>(initial_dex_file.Size()),
+        static_cast<const unsigned char*>(initial_dex_file.Begin()),
+        static_cast<jint*>(&post_no_redefine_len),
+        static_cast<unsigned char**>(&post_no_redefine_dex_data));
+    if (post_no_redefine_dex_data == nullptr) {
+      DCHECK_EQ(post_no_redefine_len, 0);
+      post_no_redefine_dex_data = const_cast<unsigned char*>(initial_dex_file.Begin());
+      post_no_redefine_len = initial_dex_file.Size();
+    } else {
+      post_no_redefine_unique_ptr = std::unique_ptr<const unsigned char>(post_no_redefine_dex_data);
+      DCHECK_GT(post_no_redefine_len, 0);
+    }
+    // Call all retransformable agents.
+    jint final_len = 0;
+    unsigned char* final_dex_data = nullptr;
+    std::unique_ptr<const unsigned char> final_dex_unique_ptr(nullptr);
+    event_handler->DispatchEvent<ArtJvmtiEvent::kClassFileLoadHookRetransformable>(
+        self,
+        static_cast<JNIEnv*>(env),
+        static_cast<jclass>(nullptr),  // The class doesn't really exist yet so send null.
+        loader.get(),
+        name.c_str(),
+        static_cast<jobject>(nullptr),  // Android doesn't seem to have protection domains
+        static_cast<jint>(post_no_redefine_len),
+        static_cast<const unsigned char*>(post_no_redefine_dex_data),
+        static_cast<jint*>(&final_len),
+        static_cast<unsigned char**>(&final_dex_data));
+    if (final_dex_data == nullptr) {
+      DCHECK_EQ(final_len, 0);
+      final_dex_data = post_no_redefine_dex_data;
+      final_len = post_no_redefine_len;
+    } else {
+      final_dex_unique_ptr = std::unique_ptr<const unsigned char>(final_dex_data);
+      DCHECK_GT(final_len, 0);
+    }
+
+    if (final_dex_data != initial_dex_file.Begin()) {
+      LOG(WARNING) << "Changing class " << descriptor;
+      art::ScopedObjectAccess soa(self);
+      art::StackHandleScope<2> hs(self);
+      // Save the results of all the non-retransformable agents.
+      // First allocate the ClassExt
+      art::Handle<art::mirror::ClassExt> ext(hs.NewHandle(klass->EnsureExtDataPresent(self)));
+      // Make sure we have a ClassExt. This is fine even though we are a temporary since it will
+      // get copied.
+      if (ext.IsNull()) {
+        // We will just return failure if we fail to allocate
+        LOG(WARNING) << "Could not allocate ext-data for class '" << descriptor << "'. "
+                     << "Aborting transformation since we will be unable to store it.";
+        self->AssertPendingOOMException();
+        return;
+      }
+
+      // 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)));
+      if (arr.IsNull()) {
+        LOG(WARNING) << "Unable to allocate byte array for initial dex-file bytes. Aborting "
+                     << "transformation";
+        self->AssertPendingOOMException();
+        return;
+      }
+
+      std::unique_ptr<const art::DexFile> dex_file(MakeSingleDexFile(self,
+                                                                     descriptor,
+                                                                     initial_dex_file.GetLocation(),
+                                                                     final_len,
+                                                                     final_dex_data));
+      if (dex_file.get() == nullptr) {
+        return;
+      }
+
+      // TODO Check Redefined dex file for invariants.
+      LOG(WARNING) << "Dex file created by class-definition time transformation of "
+                   << descriptor << " is not checked for all retransformation invariants.";
+      // TODO Put it in classpath
+      LOG(WARNING) << "Dex file created for class-definition time transformation of "
+                   << descriptor << " was not added to any classpaths!";
+      // Actually set the ClassExt's original bytes once we have actually succeeded.
+      ext->SetOriginalDexFileBytes(arr.Get());
+      // Set the return values
+      *final_class_def = &dex_file->GetClassDef(0);
+      *final_dex_file = dex_file.release();
+    }
+  }
+
   void ClassLoad(art::Handle<art::mirror::Class> klass) REQUIRES_SHARED(art::Locks::mutator_lock_) {
     if (event_handler->IsEventEnabledAnywhere(ArtJvmtiEvent::kClassLoad)) {
       art::Thread* thread = art::Thread::Current();
diff --git a/runtime/openjdkjvmti/ti_redefine.cc b/runtime/openjdkjvmti/ti_redefine.cc
index 34efc50..d2ddc21 100644
--- a/runtime/openjdkjvmti/ti_redefine.cc
+++ b/runtime/openjdkjvmti/ti_redefine.cc
@@ -520,28 +520,13 @@
   result_ = result;
 }
 
-// Allocates a ByteArray big enough to store the given number of bytes and copies them from the
-// bytes pointer.
-static art::mirror::ByteArray* AllocateAndFillBytes(art::Thread* self,
-                                                    const uint8_t* bytes,
-                                                    int32_t num_bytes)
-    REQUIRES_SHARED(art::Locks::mutator_lock_) {
-  art::StackHandleScope<1> hs(self);
-  art::Handle<art::mirror::ByteArray> arr(hs.NewHandle(
-      art::mirror::ByteArray::Alloc(self, num_bytes)));
-  if (!arr.IsNull()) {
-    // Copy it in. Just skip if it's null
-    memcpy(arr->GetData(), bytes, num_bytes);
-  }
-  return arr.Get();
-}
-
 art::mirror::ByteArray* Redefiner::ClassRedefinition::AllocateOrGetOriginalDexFileBytes() {
   // If we have been specifically given a new set of bytes use that
   if (original_dex_file_.size() != 0) {
-    return AllocateAndFillBytes(driver_->self_,
-                                &original_dex_file_.At(0),
-                                original_dex_file_.size());
+    return art::mirror::ByteArray::AllocateAndFill(
+        driver_->self_,
+        reinterpret_cast<const signed char*>(&original_dex_file_.At(0)),
+        original_dex_file_.size());
   }
 
   // See if we already have one set.
@@ -561,7 +546,10 @@
     LOG(WARNING) << "Current dex file has more than one class in it. Calling RetransformClasses "
                  << "on this class might fail if no transformations are applied to it!";
   }
-  return AllocateAndFillBytes(driver_->self_, current_dex_file.Begin(), current_dex_file.Size());
+  return art::mirror::ByteArray::AllocateAndFill(
+      driver_->self_,
+      reinterpret_cast<const signed char*>(current_dex_file.Begin()),
+      current_dex_file.Size());
 }
 
 struct CallbackCtx {
diff --git a/runtime/openjdkjvmti/ti_redefine.h b/runtime/openjdkjvmti/ti_redefine.h
index 29a7e1f..fdc13ee 100644
--- a/runtime/openjdkjvmti/ti_redefine.h
+++ b/runtime/openjdkjvmti/ti_redefine.h
@@ -96,6 +96,11 @@
 
   static jvmtiError IsModifiableClass(jvmtiEnv* env, jclass klass, jboolean* is_redefinable);
 
+  static std::unique_ptr<art::MemMap> MoveDataToMemMap(const std::string& original_location,
+                                                       jint data_len,
+                                                       const unsigned char* dex_data,
+                                                       std::string* error_msg);
+
  private:
   class ClassRedefinition {
    public:
@@ -234,11 +239,6 @@
                                               /*out*/std::string* error_msg)
       REQUIRES_SHARED(art::Locks::mutator_lock_);
 
-  static std::unique_ptr<art::MemMap> MoveDataToMemMap(const std::string& original_location,
-                                                       jint data_len,
-                                                       const unsigned char* dex_data,
-                                                       std::string* error_msg);
-
   // TODO Put on all the lock qualifiers.
   jvmtiError Run() REQUIRES_SHARED(art::Locks::mutator_lock_);
 
diff --git a/runtime/openjdkjvmti/transform.cc b/runtime/openjdkjvmti/transform.cc
index 745c0f5..3c4cfea 100644
--- a/runtime/openjdkjvmti/transform.cc
+++ b/runtime/openjdkjvmti/transform.cc
@@ -137,20 +137,6 @@
   return OK;
 }
 
-static jvmtiError CopyDataIntoJvmtiBuffer(ArtJvmTiEnv* env,
-                                          const unsigned char* source,
-                                          jint len,
-                                          /*out*/unsigned char** dest) {
-  jvmtiError res = env->Allocate(len, dest);
-  if (res != OK) {
-    return res;
-  }
-  memcpy(reinterpret_cast<void*>(*dest),
-         reinterpret_cast<const void*>(source),
-         len);
-  return OK;
-}
-
 jvmtiError Transformer::GetDexDataForRetransformation(ArtJvmTiEnv* env,
                                                       art::Handle<art::mirror::Class> klass,
                                                       /*out*/jint* dex_data_len,