From ec5ccd1e3af9cf1d0cd26c6b434b9accb25c212c Mon Sep 17 00:00:00 2001
From: Rafael Winterhalter <rafael.wth@gmail.com>
Date: Mon, 9 Oct 2023 11:40:01 +0200
Subject: [PATCH] Adds premain method such that Mockito can be used as an agent
 and enable instrumentation, considering upcoming restrictions on dynamic
 attach.

---
 gradle/mockito-core/osgi.gradle               |  2 ++
 .../org/mockito/internal/MockitoAgent.java    | 35 +++++++++++++++++++
 .../InlineDelegateByteBuddyMockMaker.java     |  6 +++-
 .../InstrumentationMemberAccessor.java        |  6 +++-
 4 files changed, 47 insertions(+), 2 deletions(-)
 create mode 100644 src/main/java/org/mockito/internal/MockitoAgent.java

diff --git a/gradle/mockito-core/osgi.gradle b/gradle/mockito-core/osgi.gradle
index 934e3463cf..90f2a60ef6 100644
--- 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
new file mode 100644
index 0000000000..0c2070135b
--- /dev/null
+++ 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
index d40967f94e..94f217fcda 100644
--- 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
index 61b3f06c46..baa3c1ba68 100644
--- 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,