Support agent by attaching it at build time

substratevm/mx.substratevm/mx_substratevm.py

@@ -1736,10 +1736,15 @@ def cinterfacetutorial(args):
 @mx.command(suite.name, 'javaagenttest', 'Runs tests for java agent with native image')
 def java_agent_test(args):
     def build_and_run(args, binary_path, native_image, agents, agents_arg):
-        test_cp = os.pathsep.join([classpath('com.oracle.svm.test')] + agents)
+        mx.log('Run agent with JVM as baseline')
+        test_cp = os.pathsep.join([classpath('com.oracle.svm.test')])
+        java_run_cp = os.pathsep.join([test_cp, mx.dependency('org.graalvm.nativeimage').classpath_repr()])
+        mx.run_java( agents_arg +  ['--add-exports=java.base/jdk.internal.org.objectweb.asm=ALL-UNNAMED',
+                     '-cp', java_run_cp, 'com.oracle.svm.test.javaagent.AgentTest'])
+        test_cp = os.pathsep.join([test_cp] + agents)
         native_agent_premain_options = ['-XXpremain:com.oracle.svm.test.javaagent.agent1.TestJavaAgent1:test.agent1=true', '-XXpremain:com.oracle.svm.test.javaagent.agent2.TestJavaAgent2:test.agent2=true']
-        image_args = ['-cp', test_cp, '-J-ea', '-J-esa', '-H:+ReportExceptionStackTraces', '-H:Class=com.oracle.svm.test.javaagent.AgentTest']
-        native_image(image_args + svm_experimental_options(['-H:PremainClasses=' + agents_arg]) + ['-o', binary_path] + args)
+        image_args = ['-cp', test_cp, '-J-ea', '-J-esa', '-H:+ReportExceptionStackTraces', '-J--add-exports=java.base/jdk.internal.org.objectweb.asm=ALL-UNNAMED', '-H:Class=com.oracle.svm.test.javaagent.AgentTest']
+        native_image(image_args + svm_experimental_options(agents_arg) + ['-o', binary_path] + args)
         mx.run([binary_path] + native_agent_premain_options)
 
     def build_and_test_java_agent_image(native_image, args):
@@ -1756,18 +1761,23 @@ def build_and_test_java_agent_image(native_image, args):
             # Note: we are not using MX here to avoid polluting the suite.py and requiring extra build flags
             mx.log("Building agent jars from " + test_classpath)
             agents = []
-            for i in range(1, 2):
+            for i in range(1, 3):
                 agent = join(tmp_dir, "testagent%d.jar" % (i))
-                agent_test_classpath = join(test_classpath, 'com', 'oracle', 'svm', 'test', 'javaagent', 'agent' + str(i))
-                class_list = [join(test_classpath, 'com', 'oracle', 'svm', 'test', 'javaagent', 'agent' + str(i), f) for f in os.listdir(agent_test_classpath) if os.path.isfile(os.path.join(agent_test_classpath, f)) and f.endswith(".class")]
-                mx.run([mx.get_jdk().jar, 'cmf', join(test_classpath, 'resources', 'javaagent' + str(i), 'MANIFEST.MF'), agent] + class_list, cwd = tmp_dir)
+                current_dir = os.getcwd()
+                # Change to test classpath to create agent jar file
+                os.chdir(test_classpath)
+                agent_test_classpath = join('com', 'oracle', 'svm', 'test', 'javaagent', 'agent' + str(i))
+                class_list = [join('com', 'oracle', 'svm', 'test', 'javaagent', 'agent' + str(i), f) for f in os.listdir(agent_test_classpath) if os.path.isfile(os.path.join(agent_test_classpath, f)) and f.endswith(".class")]
+                class_list.append(join('com', 'oracle', 'svm', 'test', 'javaagent', 'AgentPremainHelper.class'))
+                mx.run([mx.get_jdk().jar, 'cmf', join(test_classpath, 'resources', 'javaagent' + str(i), 'MANIFEST.MF'), agent] + class_list, cwd = test_classpath)
                 agents.append(agent)
+                os.chdir(current_dir)
 
             mx.log("Building images with different agent orders ")
-            build_and_run(args, join(tmp_dir, 'agenttest1'), native_image, agents,'com.oracle.svm.test.javaagent.agent1.TestJavaAgent1,com.oracle.svm.test.javaagent.agent2.TestJavaAgent2')
+            build_and_run(args, join(tmp_dir, 'agenttest1'), native_image, agents,[f'-javaagent:{agents[0]}=test.agent1=true', f'-javaagent:{agents[1]}=test.agent2=true'])
 
             # Switch the premain sequence of agent1 and agent2
-            build_and_run(args, join(tmp_dir, 'agenttest2'), native_image, agents, 'com.oracle.svm.test.javaagent.agent2.TestJavaAgent2,com.oracle.svm.test.javaagent.agent1.TestJavaAgent1')
+            build_and_run(args, join(tmp_dir, 'agenttest2'), native_image, agents, [f'-javaagent:{agents[1]}=test.agent2=true', f'-javaagent:{agents[0]}=test.agent1=true'])
 
     native_image_context_run(build_and_test_java_agent_image, args)
 

substratevm/mx.substratevm/suite.py

@@ -994,6 +994,7 @@
                 "java.base" : [
                     "jdk.internal.misc",
                     "sun.security.jca",
+                    "jdk.internal.org.objectweb.asm"
                 ],
             },
             "checkstyle": "com.oracle.svm.test",

substratevm/src/com.oracle.graal.pointsto/src/com/oracle/graal/pointsto/api/HostVM.java

@@ -336,6 +336,11 @@ public boolean isClosedTypeWorld() {
         return true;
     }
 
+    @SuppressWarnings("unused")
+    public boolean isFromJavaAgent(Class<?> clazz) {
+        return false;
+    }
+
     /**
      * Helpers to determine what analysis actions should be taken for a given Multi-Method version.
      */

substratevm/src/com.oracle.graal.pointsto/src/com/oracle/graal/pointsto/infrastructure/WrappedConstantPool.java

@@ -26,10 +26,12 @@
 
 import static jdk.vm.ci.common.JVMCIError.unimplemented;
 
+import java.util.Arrays;
 import java.util.List;
 import java.util.stream.Collectors;
 
 import com.oracle.graal.pointsto.constraints.UnresolvedElementException;
+import com.oracle.graal.pointsto.meta.AnalysisUniverse;
 import com.oracle.graal.pointsto.util.GraalAccess;
 
 import jdk.graal.compiler.api.replacements.SnippetReflectionProvider;
@@ -41,6 +43,7 @@
 import jdk.vm.ci.meta.JavaType;
 import jdk.vm.ci.meta.ResolvedJavaMethod;
 import jdk.vm.ci.meta.ResolvedJavaType;
+import jdk.vm.ci.meta.UnresolvedJavaMethod;
 
 public class WrappedConstantPool implements ConstantPool, ConstantPoolPatch {
 
@@ -117,7 +120,26 @@ public JavaMethod lookupMethod(int cpi, int opcode) {
     @Override
     public JavaMethod lookupMethod(int cpi, int opcode, ResolvedJavaMethod caller) {
         try {
-            return universe.lookupAllowUnresolved(wrapped.lookupMethod(cpi, opcode, OriginalMethodProvider.getOriginalMethod(caller)));
+            JavaMethod ret = universe.lookupAllowUnresolved(wrapped.lookupMethod(cpi, opcode, OriginalMethodProvider.getOriginalMethod(caller)));
+            /**
+             * The java agent classes are loaded by appClassloader, but their dependencies could be
+             * loaded by NativeImageClassloader. So if the required method could not be resolved, we
+             * look further into classes loaded by nativeImageClassloader.
+             */
+            if (ret instanceof UnresolvedJavaMethod && universe.hostVM().isFromJavaAgent(OriginalClassProvider.getJavaClass(caller.getDeclaringClass()))) {
+                UnresolvedJavaMethod unresolvedResult = (UnresolvedJavaMethod) ret;
+                String className = unresolvedResult.format("%H");
+                String methodNameWithSignature = unresolvedResult.format("%n(%P)");
+                try {
+                    Class<?> loadedClass = ((AnalysisUniverse) universe).getConcurrentAnalysisAccess().findClassByName(className);
+                    ResolvedJavaType resolvedType = ((AnalysisUniverse) universe).getOriginalMetaAccess().lookupJavaType(loadedClass);
+                    ResolvedJavaMethod resolvedMethod = Arrays.stream(resolvedType.getDeclaredMethods(false)).filter(m -> m.format("%n(%P)").equals(methodNameWithSignature)).findFirst().get();
+                    return universe.lookupAllowUnresolved(resolvedMethod);
+                } catch (Exception e) {
+                    // Could not get the resolved method, get to the unresolved path
+                }
+            }
+            return ret;
         } catch (Throwable ex) {
             Throwable cause = ex;
             if (ex instanceof ExceptionInInitializerError && ex.getCause() != null) {

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/PreMainSupport.java

@@ -140,12 +140,12 @@ public String[] retrievePremainArgs(String[] args) {
 
     public void invokePremain() {
         for (PremainMethod premainMethod : premainMethods) {
-
-            Object[] args = premainMethod.args;
-            if (premainOptions.containsKey(premainMethod.className)) {
-                args[0] = premainOptions.get(premainMethod.className);
-            }
             try {
+                Object[] args = premainMethod.args;
+                // options set at runtime can override options set at build time
+                if (premainOptions.containsKey(premainMethod.className)) {
+                    args[0] = premainOptions.get(premainMethod.className);
+                }
                 // premain method must be static
                 premainMethod.method.invoke(null, args);
             } catch (Throwable t) {

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/SubstrateOptions.java

@@ -202,6 +202,10 @@ protected void onValueUpdate(EconomicMap<OptionKey<?>, Object> values, String ol
     public static OptionEnabledHandler<Boolean> imageLayerEnabledHandler;
     public static OptionEnabledHandler<Boolean> imageLayerCreateEnabledHandler;
 
+    @APIOption(name = "-javaagent", valueSeparator = ':')//
+    @Option(help = "Enable the specified java agent in native image. Usage: -javaagent:<jarpath>[=<options>]. It's the same as using javaagent in JVM", type = User, stability = OptionStability.EXPERIMENTAL)//
+    public static final HostedOptionKey<AccumulatingLocatableMultiOptionValue.Strings> JavaAgent = new HostedOptionKey<>(AccumulatingLocatableMultiOptionValue.Strings.build());
+
     @Fold
     public static boolean getSourceLevelDebug() {
         return SourceLevelDebug.getValue();

substratevm/src/com.oracle.svm.driver/src/com/oracle/svm/driver/NativeImage.java

@@ -272,6 +272,8 @@ private static <T> String oR(OptionKey<T> option) {
     final String oHInspectServerContentPath = oH(PointstoOptions.InspectServerContentPath);
     final String oHDeadlockWatchdogInterval = oH(SubstrateOptions.DeadlockWatchdogInterval);
 
+    final String oHJavaAgent = oH(SubstrateOptions.JavaAgent);
+
     final Map<String, String> imageBuilderEnvironment = new HashMap<>();
     private final ArrayList<String> imageBuilderArgs = new ArrayList<>();
     private final LinkedHashSet<Path> imageBuilderModulePath = new LinkedHashSet<>();
@@ -1415,6 +1417,7 @@ private List<String> getAgentArguments() {
         String agentOptions = "";
         List<ArgumentEntry> traceClassInitializationOpts = getHostedOptionArgumentValues(imageBuilderArgs, oHTraceClassInitialization);
         List<ArgumentEntry> traceObjectInstantiationOpts = getHostedOptionArgumentValues(imageBuilderArgs, oHTraceObjectInstantiation);
+        List<ArgumentEntry> javaAgentOpts = getHostedOptionArgumentValues(imageBuilderArgs, oHJavaAgent);
         if (!traceClassInitializationOpts.isEmpty()) {
             agentOptions = getAgentOptions(traceClassInitializationOpts, "c");
         }
@@ -1433,6 +1436,12 @@ private List<String> getAgentArguments() {
             args.add("-agentlib:native-image-diagnostics-agent=" + agentOptions);
         }
 
+        if (!javaAgentOpts.isEmpty()) {
+            for (ArgumentEntry javaAgentOpt : javaAgentOpts) {
+                args.add("-javaagent:" + javaAgentOpt.value);
+            }
+        }
+
         return args;
     }
 

substratevm/src/com.oracle.svm.hosted/src/com/oracle/svm/hosted/InstrumentFeature.java

@@ -27,6 +27,7 @@
 package com.oracle.svm.hosted;
 
 import com.oracle.svm.core.PreMainSupport;
+import com.oracle.svm.core.SubstrateOptions;
 import com.oracle.svm.core.feature.AutomaticallyRegisteredFeature;
 import com.oracle.svm.core.feature.InternalFeature;
 import com.oracle.svm.core.option.AccumulatingLocatableMultiOptionValue;
@@ -39,10 +40,12 @@
 import org.graalvm.nativeimage.ImageSingletons;
 import org.graalvm.nativeimage.hosted.Feature;
 
+import java.io.IOException;
 import java.lang.instrument.Instrumentation;
 import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.jar.JarFile;
 
 /**
  * This feature supports instrumentation in native image.
@@ -73,10 +76,10 @@ public void afterRegistration(AfterRegistrationAccess access) {
         FeatureImpl.AfterRegistrationAccessImpl a = (FeatureImpl.AfterRegistrationAccessImpl) access;
         cl = a.getImageClassLoader().getClassLoader();
         ImageSingletons.add(PreMainSupport.class, preMainSupport = new PreMainSupport());
-        if (Options.PremainClasses.hasBeenSet()) {
-            List<String> premains = Options.PremainClasses.getValue().values();
-            for (String premain : premains) {
-                addPremainClass(premain);
+        if (SubstrateOptions.JavaAgent.hasBeenSet()) {
+            List<String> agentOptions = SubstrateOptions.JavaAgent.getValue().values();
+            for (String agentOption : agentOptions) {
+                addPremainClass(agentOption);
             }
         }
     }
@@ -94,12 +97,34 @@ public void afterRegistration(AfterRegistrationAccess access) {
      * is absent. <br>
      * So this method looks for them in the same order.
      */
-    private void addPremainClass(String premainClass) {
+    private void addPremainClass(String javaagentOption) {
+        int separatorIndex = javaagentOption.indexOf("=");
+        String agent;
+        String premainClass = null;
+        String options = "";
+        // Get the agent file
+        if (separatorIndex == -1) {
+            agent = javaagentOption;
+        } else {
+            agent = javaagentOption.substring(0, separatorIndex);
+            options = javaagentOption.substring(separatorIndex + 1);
+        }
+        // Read MANIFEST in agent jar
+        try {
+            JarFile agentJarFile = new JarFile(agent);
+            premainClass = agentJarFile.getManifest().getMainAttributes().getValue("Premain-Class");
+        } catch (IOException e) {
+            // This shall not happen, because at this moment GraalVM is running with -javaagent.
+            // If the agent doesn't exist, the JVM shall fail to start.
+            UserError.abort(e, "Can't read the agent jar %s. Please check option %s", agent,
+                            SubstrateOptionsParser.commandArgument(SubstrateOptions.JavaAgent, ""));
+        }
+
         try {
             Class<?> clazz = Class.forName(premainClass, false, cl);
             Method premain = null;
             List<Object> args = new ArrayList<>();
-            args.add(""); // First argument is options which will be set at runtime
+            args.add(options); // First argument is options which will be set at runtime
             try {
                 premain = clazz.getDeclaredMethod("premain", String.class, Instrumentation.class);
                 args.add(new PreMainSupport.NativeImageNoOpRuntimeInstrumentation());
@@ -108,13 +133,13 @@ private void addPremainClass(String premainClass) {
                     premain = clazz.getDeclaredMethod("premain", String.class);
                 } catch (NoSuchMethodException e1) {
                     UserError.abort(e1, "Can't register agent premain method, because can't find the premain method from the given class %s. Please check your %s setting.", premainClass,
-                                    SubstrateOptionsParser.commandArgument(Options.PremainClasses, ""));
+                                    SubstrateOptionsParser.commandArgument(SubstrateOptions.JavaAgent, ""));
                 }
             }
             preMainSupport.registerPremainMethod(premainClass, premain, args.toArray(new Object[0]));
         } catch (ClassNotFoundException e) {
             UserError.abort(e, "Can't register agent premain method, because the given class %s is not found. Please check your %s setting.", premainClass,
-                            SubstrateOptionsParser.commandArgument(Options.PremainClasses, ""));
+                            SubstrateOptionsParser.commandArgument(SubstrateOptions.JavaAgent, ""));
         }
     }
 }

substratevm/src/com.oracle.svm.hosted/src/com/oracle/svm/hosted/SVMHost.java

@@ -1087,4 +1087,17 @@ public Set<AnalysisMethod> loadOpenTypeWorldDispatchTableMethods(AnalysisType ty
         // return OpenTypeWorldFeature.loadDispatchTable(type);
         return Set.of();
     }
+
+    @Override
+    public boolean isFromJavaAgent(Class<?> clazz) {
+        if (SubstrateOptions.JavaAgent.hasBeenSet()) {
+            try {
+                String classLocation = clazz.getProtectionDomain().getCodeSource().getLocation().getFile();
+                return SubstrateOptions.JavaAgent.getValue().values().stream().map(s -> s.split("=")[0]).anyMatch(s -> s.equals(classLocation));
+            } catch (Exception e) {
+                return false;
+            }
+        }
+        return false;
+    }
 }

substratevm/src/com.oracle.svm.test/src/com/oracle/svm/test/javaagent/AgentTest.java

@@ -58,10 +58,20 @@ private static void testPremainSequence() {
         }
     }
 
+    private static void testInstrumentation() {
+        // The return value of getCounter() should be changed by agent
+        Assert.assertEquals(11, getCounter());
+    }
+
+    private static int getCounter() {
+        return 10;
+    }
+
     public static void main(String[] args) {
         testPremain();
         testAgentOptions();
         testPremainSequence();
+        testInstrumentation();
         System.out.println("Finished running Agent test.");
     }
 }