Replace unwanted class instance from agent

As the java agent is running at build time, some classes that should
never run in native runtime could get instantiated and referred by classes about to build in native image.

Agent developer can provide a black list by setting -H:ExcludeAgentClasses, so that each one of them
will be replaced with a safe object at build time.

Author: ziyilin

compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/java/LambdaUtils.java

@@ -81,9 +81,10 @@ private LambdaUtils() {
     public static String findStableLambdaName(ResolvedJavaType lambdaType) {
         ResolvedJavaMethod[] lambdaProxyMethods = Arrays.stream(lambdaType.getDeclaredMethods(false)).filter(m -> !m.isBridge() && m.isPublic()).toArray(ResolvedJavaMethod[]::new);
         /*
-         * Traverse all methods in lambda class for invokes because it is possible a javaagent may inject new methods into
-         * the lambda class. For example, Byte-buddy used by OTele can transform all classes that implements
-         * {@link java.util.concurrent.Callable} by injecting new methods that may not have any invokes.
+         * Traverse all methods in lambda class for invokes because it is possible a javaagent may
+         * inject new methods into the lambda class. For example, Byte-buddy used by OTele can
+         * transform all classes that implements {@link java.util.concurrent.Callable} by injecting
+         * new methods that may not have any invokes.
          */
         Set<JavaMethod> invokedMethods = new HashSet<>();
         for (int i = 0; i < lambdaProxyMethods.length; i++) {

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

@@ -26,6 +26,8 @@
 
 package com.oracle.svm.hosted;
 
+import java.io.BufferedReader;
+import java.io.FileReader;
 import java.lang.instrument.Instrumentation;
 import java.lang.reflect.Method;
 import java.util.ArrayList;
@@ -50,23 +52,26 @@
 
 import java.io.IOException;
 import java.util.jar.JarFile;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 
 /**
  * This feature supports instrumentation in native image.
  */
 @AutomaticallyRegisteredFeature
 public class InstrumentFeature implements InternalFeature {
     public static class Options {
-        @Option(help = "Specify premain-class list. Multiple classes are separated by comma, and order matters. This is an experimental option.", stability = OptionStability.EXPERIMENTAL)//
-        public static final HostedOptionKey<AccumulatingLocatableMultiOptionValue.Strings> PremainClasses = new HostedOptionKey<>(
+        @Option(help = "The classes should be excluded from native image but actually get instantiated at build time." +
+                        " Multiple values are separated by comma. It accepts Java regular expression." +
+                        " Using @ to specify a file in which each line is an item to take.", stability = OptionStability.EXPERIMENTAL)//
+        public static final HostedOptionKey<AccumulatingLocatableMultiOptionValue.Strings> ExcludeAgentClasses = new HostedOptionKey<>(
                         AccumulatingLocatableMultiOptionValue.Strings.buildWithCommaDelimiter());
 
     }
 
-    @Override
-    public boolean isInConfiguration(IsInConfigurationAccess access) {
-        return !Options.PremainClasses.getValue().values().isEmpty();
-    }
+    private static final NativeImageSystemClassLoader nativeImageSystemClassLoader = NativeImageSystemClassLoader.singleton();
+    private static List<Pattern> excludeClasses;
+    private static final Object DUMMY_OBJECT = new Object();
 
     @Override
     public List<Class<? extends Feature>> getRequiredFeatures() {
@@ -143,4 +148,57 @@ private static Method findPremainMethod(String premainClass, Class<?> javaAgentC
 
         throw UserError.abort("Could not register agent premain method: class %s neither declares 'premain(String, Instrumentation)' nor 'premain(String)'.", premainClass);
     }
+
+    @Override
+    public void duringSetup(DuringSetupAccess a) {
+        FeatureImpl.DuringSetupAccessImpl access = (FeatureImpl.DuringSetupAccessImpl) a;
+        /**
+         * As the java agent is running at build time, some classes that should never run in native
+         * runtime could get instantiated and referred by classes in native image. Agent developer
+         * can provide a black list by setting {@link Options.ExcludeAgentClasses}, so that each one
+         * of them will be replaced with a safe object at build time.
+         */
+        getExcludes();
+        access.registerObjectReplacer(original -> {
+            String className = original.getClass().getName();
+            // Replace the classloader with nativeImageSystemClassLoader.defaultSystemClassLoader
+            if (original instanceof ClassLoader && matchExclude(className)) {
+                return nativeImageSystemClassLoader.defaultSystemClassLoader;
+            }
+            // Replace other unwanted objects with an object
+            if (matchExclude(className)) {
+                return DUMMY_OBJECT;
+            }
+            return original;
+        });
+    }
+
+    private static void getExcludes() {
+        excludeClasses = new ArrayList<>();
+        for (String excludePattern : Options.ExcludeAgentClasses.getValue().values()) {
+            if (excludePattern.startsWith("@")) {
+                String optionFile = excludePattern.substring(1);
+                try (BufferedReader br = new BufferedReader(new FileReader(optionFile))) {
+                    String line;
+                    while ((line = br.readLine()) != null) {
+                        excludeClasses.add(Pattern.compile(line));
+                    }
+                } catch (IOException e) {
+                    throw new RuntimeException("Failed to read file  " + optionFile + "specified in -H:" + Options.ExcludeAgentClasses.getName(), e);
+                }
+            } else {
+                excludeClasses.add(Pattern.compile(excludePattern));
+            }
+        }
+    }
+
+    private boolean matchExclude(String input) {
+        for (Pattern excludeClass : excludeClasses) {
+            Matcher matcher = excludeClass.matcher(input);
+            if (matcher.matches()) {
+                return true;
+            }
+        }
+        return false;
+    }
 }

substratevm/src/com.oracle.svm.hosted/src/com/oracle/svm/hosted/classinitialization/ClassInitializationFeature.java

@@ -56,6 +56,7 @@
 import com.oracle.svm.core.option.SubstrateOptionsParser;
 import com.oracle.svm.core.snippets.SnippetRuntime;
 import com.oracle.svm.core.util.UserError;
+import com.oracle.svm.hosted.InstrumentFeature;
 import com.oracle.svm.hosted.FeatureImpl;
 import com.oracle.svm.hosted.FeatureImpl.AfterAnalysisAccessImpl;
 import com.oracle.svm.hosted.FeatureImpl.BeforeAnalysisAccessImpl;
@@ -66,6 +67,8 @@
 import jdk.graal.compiler.options.OptionValues;
 import jdk.graal.compiler.phases.util.Providers;
 
+import org.graalvm.nativeimage.hosted.Feature;
+
 @AutomaticallyRegisteredFeature
 public class ClassInitializationFeature implements InternalFeature {
     private static final String NATIVE_IMAGE_CLASS_REASON = "Native Image classes are always initialized at build time";
@@ -137,6 +140,15 @@ private static void initializeNativeImagePackagesAtBuildTime(ClassInitialization
         initializationSupport.initializeAtBuildTime("org.graalvm.jniutils", NATIVE_IMAGE_CLASS_REASON);
     }
 
+    @Override
+    public List<Class<? extends Feature>> getRequiredFeatures() {
+        /**
+         * {@link InstrumentFeature} replace the objects that should not appear in runtime. It must
+         * come before this feature, otherwise shouldInitializeAtRuntime will be triggered first.
+         */
+        return List.of(InstrumentFeature.class);
+    }
+
     @Override
     public void duringSetup(DuringSetupAccess a) {
         FeatureImpl.DuringSetupAccessImpl access = (FeatureImpl.DuringSetupAccessImpl) a;