Adds premain method such that Mockito can be used as an agent and ena…

…ble instrumentation, considering upcoming restrictions on dynamic attach.
mockito · Oct 9, 2023 · ec5ccd1 · ec5ccd1
1 parent 5d946b4
commit ec5ccd1
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 2 deletions.
diff --git a/gradle/mockito-core/osgi.gradle b/gradle/mockito-core/osgi.gradle
@@ -8,6 +8,8 @@ tasks.named("jar", Jar) {
             'Bundle-SymbolicName': 'org.mockito.mockito-core',
             'Bundle-Version': "\${version_cleanup;${project.version}}",
             '-versionpolicy': '[${version;==;${@}},${version;+;${@}})',
+            "Can-Retransform-Classes": "true",
+            "Premain-Class": "org.mockito.internal.MockitoAgent",
             'Export-Package': "org.mockito.internal.*;status=INTERNAL;mandatory:=status;version=${archiveVersion.get()},org.mockito.*;version=${archiveVersion.get()}",
             'Import-Package': [
                     'net.bytebuddy.*;version="[1.6.0,2.0)"',

diff --git a/src/main/java/org/mockito/internal/MockitoAgent.java b/src/main/java/org/mockito/internal/MockitoAgent.java
@@ -0,0 +1,35 @@
+/*
+ * Copyright (c) 2007 Mockito contributors
+ * This program is made available under the terms of the MIT License.
+ */
+package org.mockito.internal;
+
+import java.lang.instrument.Instrumentation;
+
+public class MockitoAgent {
+
+    private static Instrumentation instrumentation;
+
+    public static void premain(String arg, Instrumentation instrumentation) {
+        MockitoAgent.instrumentation = instrumentation;
+    }
+
+    public static Instrumentation getInstrumentation() {
+        Class<?> caller =
+                StackWalker.getInstance(StackWalker.Option.RETAIN_CLASS_REFERENCE).getCallerClass();
+        try {
+            if (caller
+                            != Class.forName(
+                                    "org.mockito.internal.creation.bytebuddy.InlineDelegateByteBuddyMockMaker")
+                    && caller
+                            != Class.forName(
+                                    "org.mockito.internal.util.reflection.InstrumentationMemberAccessor")) {
+                throw new RuntimeException(
+                        "Cannot access Mockito agent from unknown class: " + caller);
+            }
+        } catch (ClassNotFoundException e) {
+            throw new IllegalStateException(e);
+        }
+        return instrumentation;
+    }
+}
diff --git a/src/main/java/org/mockito/internal/creation/bytebuddy/InlineDelegateByteBuddyMockMaker.java b/src/main/java/org/mockito/internal/creation/bytebuddy/InlineDelegateByteBuddyMockMaker.java
@@ -11,6 +11,7 @@
 import org.mockito.exceptions.base.MockitoException;
 import org.mockito.exceptions.base.MockitoInitializationException;
 import org.mockito.exceptions.misusing.MockitoConfigurationException;
+import org.mockito.internal.MockitoAgent;
 import org.mockito.internal.SuppressSignatureCheck;
 import org.mockito.internal.configuration.plugins.Plugins;
 import org.mockito.internal.creation.instance.ConstructorInstantiator;
@@ -112,7 +113,10 @@ class InlineDelegateByteBuddyMockMaker
         Throwable initializationError = null;
         try {
             try {
-                instrumentation = ByteBuddyAgent.install();
+                instrumentation = MockitoAgent.getInstrumentation();
+                if (instrumentation == null) {
+                    instrumentation = ByteBuddyAgent.install();
+                }
                 if (!instrumentation.isRetransformClassesSupported()) {
                     throw new IllegalStateException(
                             join(

diff --git a/src/main/java/org/mockito/internal/util/reflection/InstrumentationMemberAccessor.java b/src/main/java/org/mockito/internal/util/reflection/InstrumentationMemberAccessor.java
@@ -9,6 +9,7 @@
 import net.bytebuddy.dynamic.loading.ClassLoadingStrategy;
 import net.bytebuddy.implementation.MethodCall;
 import org.mockito.exceptions.base.MockitoInitializationException;
+import org.mockito.internal.MockitoAgent;
 import org.mockito.internal.SuppressSignatureCheck;
 import org.mockito.plugins.MemberAccessor;
 
@@ -46,7 +47,10 @@ class InstrumentationMemberAccessor implements MemberAccessor {
         Dispatcher dispatcher;
         Throwable throwable;
         try {
-            instrumentation = ByteBuddyAgent.install();
+            instrumentation = MockitoAgent.getInstrumentation();
+            if (instrumentation == null) {
+                instrumentation = ByteBuddyAgent.install();
+            }
             // We need to generate a dispatcher instance that is located in a distinguished class
             // loader to create a unique (unnamed) module to which we can open other packages to.
             // This way, we assure that classes within Mockito's module (which might be a shared,