1 files changed, 23 insertions, 8 deletions

diff --git a/dexmaker-mockito-inline/src/main/jni/dexmakerjvmtiagent/agent.cc b/dexmaker-mockito-inline/src/main/jni/dexmakerjvmtiagent/agent.cc
index e00ada8..358e375 100644
--- a/dexmaker-mockito-inline/src/main/jni/dexmakerjvmtiagent/agent.cc
+++ b/dexmaker-mockito-inline/src/main/jni/dexmakerjvmtiagent/agent.cc
@@ -779,6 +779,19 @@ Java_com_android_dx_mockito_inline_ClassTransformer_nativeRedefine(JNIEnv* env,
     return transformedArr;
 }
 
+// Register the ClassFileLoadHook hook. This causes Transform to be called for every class load and
+// for every trigger of RetransformClasses
+static jvmtiError registerClassFileLoadHook() {
+    return localJvmtiEnv->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_CLASS_FILE_LOAD_HOOK,
+                                                   NULL);
+}
+
+// Unregister the ClassFileLoadHook hook.
+static jvmtiError unregisterClassFileLoadHook() {
+    return localJvmtiEnv->SetEventNotificationMode(JVMTI_DISABLE, JVMTI_EVENT_CLASS_FILE_LOAD_HOOK,
+                                                   NULL);
+}
+
 // Initializes the agent
 extern "C" jint Agent_OnAttach(JavaVM* vm,
                                char* options,
@@ -806,12 +819,6 @@ extern "C" jint Agent_OnAttach(JavaVM* vm,
         return error;
     }
 
-    error = localJvmtiEnv->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_CLASS_FILE_LOAD_HOOK,
-                                                    NULL);
-    if (error != JVMTI_ERROR_NONE) {
-        return error;
-    }
-
     return JVMTI_ERROR_NONE;
 }
 
@@ -854,8 +861,16 @@ Java_com_android_dx_mockito_inline_JvmtiAgent_nativeRetransformClasses(JNIEnv* e
         transformedClasses[i] = (jclass) env->NewGlobalRef(env->GetObjectArrayElement(classes, i));
     }
 
-    jvmtiError error = localJvmtiEnv->RetransformClasses(numTransformedClasses,
-                                                         transformedClasses);
+    jvmtiError error = registerClassFileLoadHook();
+    if (error == JVMTI_ERROR_NONE) {
+        error = localJvmtiEnv->RetransformClasses(numTransformedClasses,
+                                                  transformedClasses);
+
+        jvmtiError unregisterError = unregisterClassFileLoadHook();
+        if (error == JVMTI_ERROR_NONE && unregisterError != JVMTI_ERROR_NONE) {
+            error = unregisterError;
+        }
+    }
 
     for (int i = 0; i < numTransformedClasses; i++) {
         env->DeleteGlobalRef(transformedClasses[i]);