Support agent by attaching it at build time

GraalVM now can apply the class transformation in the native image by
attaching the java agent at build time.
User can attach the java agent to native-image with -javaagent option.

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,10 @@ public boolean isClosedTypeWorld() {
         return true;
     }
 
+    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.");
     }
 }