Delay loading InetAddressResolverProvider until after the agent has started (#11987)

Co-authored-by: Lauri Tulmin <[email protected]>

Author: serkan-ozal

javaagent-bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/AgentInitializer.java

@@ -26,6 +26,7 @@ public final class AgentInitializer {
   @Nullable private static ClassLoader agentClassLoader = null;
   @Nullable private static AgentStarter agentStarter = null;
   private static boolean isSecurityManagerSupportEnabled = false;
+  private static volatile boolean agentStarted = false;
 
   public static void initialize(Instrumentation inst, File javaagentFile, boolean fromPremain)
       throws Exception {
@@ -51,6 +52,7 @@ public Void run() throws Exception {
             agentStarter = createAgentStarter(agentClassLoader, inst, javaagentFile);
             if (!fromPremain || !delayAgentStart()) {
               agentStarter.start();
+              agentStarted = true;
             }
             return null;
           }
@@ -149,11 +151,27 @@ public static void delayedStartHook() throws Exception {
           @Override
           public Void run() {
             agentStarter.start();
+            agentStarted = true;
             return null;
           }
         });
   }
 
+  /**
+   * Check whether agent has started or not along with VM.
+   *
+   * <p>This method is used by
+   * io.opentelemetry.javaagent.tooling.AgentStarterImpl#InetAddressClassFileTransformer internally
+   * to check whether agent has started.
+   *
+   * @param vmStarted flag about whether VM has started or not.
+   * @return {@code true} if agent has started or not along with VM, {@code false} otherwise.
+   */
+  @SuppressWarnings("unused")
+  public static boolean isAgentStarted(boolean vmStarted) {
+    return vmStarted && agentStarted;
+  }
+
   public static ClassLoader getExtensionsClassLoader() {
     // agentStarter can be null when running tests
     return agentStarter != null ? agentStarter.getExtensionClassLoader() : null;

javaagent-tooling/jdk18-testing/build.gradle.kts

@@ -0,0 +1,11 @@
+plugins {
+  id("otel.javaagent-testing")
+}
+
+dependencies {
+  compileOnly("io.opentelemetry:opentelemetry-sdk-common")
+}
+
+otelJava {
+  minJavaVersionSupported.set(JavaVersion.VERSION_18)
+}

javaagent-tooling/jdk18-testing/src/main/java/testing/TestResourceProvider.java

@@ -0,0 +1,30 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package testing;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
+import io.opentelemetry.sdk.autoconfigure.spi.ResourceProvider;
+import io.opentelemetry.sdk.resources.Resource;
+import java.net.InetAddress;
+import java.net.UnknownHostException;
+
+@AutoService(ResourceProvider.class)
+public class TestResourceProvider implements ResourceProvider {
+
+  @Override
+  public Resource createResource(ConfigProperties config) {
+    // used in test to determine whether this method was called
+    System.setProperty("test.resource.provider.called", "true");
+    // this call trigger loading InetAddressResolverProvider SPI on jdk 18
+    try {
+      InetAddress.getLocalHost();
+    } catch (UnknownHostException e) {
+      throw new IllegalStateException(e);
+    }
+    return Resource.empty();
+  }
+}

javaagent-tooling/jdk18-testing/src/test/java/io/opentelemetry/javaagent/tooling/inetaddress/InetAddressResolverTest.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.inetaddress;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.net.InetAddress;
+import org.junit.jupiter.api.Test;
+
+public class InetAddressResolverTest {
+
+  @Test
+  void agentStartShouldNotTriggerLoadingCustomInetAddressResolvers() throws Exception {
+    // This system property is set in TestResourceProvider
+    assertThat(System.getProperty("test.resource.provider.called")).isEqualTo("true");
+    // Agent start should not trigger loading (and instantiating) custom InetAddress resolvers
+    assertThat(TestAddressResolver.isInstantiated()).isFalse();
+
+    // Trigger loading (and instantiating) custom InetAddress resolvers manually
+    InetAddress.getAllByName("test");
+
+    // Verify that custom InetAddress resolver loaded and instantiated
+    assertThat(TestAddressResolver.isInstantiated()).isTrue();
+  }
+}

javaagent-tooling/jdk18-testing/src/test/java/io/opentelemetry/javaagent/tooling/inetaddress/TestAddressResolver.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.inetaddress;
+
+import java.net.InetAddress;
+import java.net.UnknownHostException;
+import java.net.spi.InetAddressResolver;
+import java.util.stream.Stream;
+
+public class TestAddressResolver implements InetAddressResolver {
+
+  private static volatile boolean instantiated = false;
+
+  @SuppressWarnings("StaticAssignmentInConstructor")
+  public TestAddressResolver() {
+    TestAddressResolver.instantiated = true;
+  }
+
+  public static boolean isInstantiated() {
+    return instantiated;
+  }
+
+  @Override
+  public Stream<InetAddress> lookupByName(String host, LookupPolicy lookupPolicy)
+      throws UnknownHostException {
+    if (host.equals("test")) {
+      return Stream.of(InetAddress.getByAddress(new byte[] {127, 0, 0, 1}));
+    }
+    throw new UnknownHostException();
+  }
+
+  @Override
+  public String lookupByAddress(byte[] addr) {
+    throw new UnsupportedOperationException();
+  }
+}

javaagent-tooling/jdk18-testing/src/test/java/io/opentelemetry/javaagent/tooling/inetaddress/TestAddressResolverProvider.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.inetaddress;
+
+import java.net.spi.InetAddressResolver;
+import java.net.spi.InetAddressResolverProvider;
+
+public class TestAddressResolverProvider extends InetAddressResolverProvider {
+
+  @Override
+  public InetAddressResolver get(Configuration configuration) {
+    return new TestAddressResolver();
+  }
+
+  @Override
+  public String name() {
+    return "Test Internet Address Resolver Provider";
+  }
+}

javaagent-tooling/jdk18-testing/src/test/resources/META-INF/services/java.net.spi.InetAddressResolverProvider

@@ -0,0 +1 @@
+io.opentelemetry.javaagent.tooling.inetaddress.TestAddressResolverProvider

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/AgentStarterImpl.java

@@ -68,6 +68,8 @@ public boolean delayStart() {
 
   @Override
   public void start() {
+    installTransformers();
+
     EarlyInitAgentConfig earlyConfig = EarlyInitAgentConfig.create();
     extensionClassLoader = createExtensionClassLoader(getClass().getClassLoader(), earlyConfig);
 
@@ -115,6 +117,14 @@ public void start() {
     }
   }
 
+  private void installTransformers() {
+    // prevents loading InetAddressResolverProvider SPI before agent has started
+    // https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/7130
+    // https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/10921
+    InetAddressClassFileTransformer transformer = new InetAddressClassFileTransformer();
+    instrumentation.addTransformer(transformer, true);
+  }
+
   @SuppressWarnings("SystemOut")
   private static void logUnrecognizedLoggerImplWarning(String loggerImplementationName) {
     System.err.println(
@@ -180,4 +190,60 @@ public void visitCode() {
       return hookInserted ? cw.toByteArray() : null;
     }
   }
+
+  private static class InetAddressClassFileTransformer implements ClassFileTransformer {
+    boolean hookInserted = false;
+
+    @Override
+    public byte[] transform(
+        ClassLoader loader,
+        String className,
+        Class<?> classBeingRedefined,
+        ProtectionDomain protectionDomain,
+        byte[] classfileBuffer) {
+      if (!"java/net/InetAddress".equals(className)) {
+        return null;
+      }
+      ClassReader cr = new ClassReader(classfileBuffer);
+      ClassWriter cw = new ClassWriter(cr, 0);
+      ClassVisitor cv =
+          new ClassVisitor(AsmApi.VERSION, cw) {
+            @Override
+            public MethodVisitor visitMethod(
+                int access, String name, String descriptor, String signature, String[] exceptions) {
+              MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
+              if (!"resolver".equals(name)) {
+                return mv;
+              }
+              return new MethodVisitor(api, mv) {
+                @Override
+                public void visitMethodInsn(
+                    int opcode,
+                    String ownerClassName,
+                    String methodName,
+                    String descriptor,
+                    boolean isInterface) {
+                  super.visitMethodInsn(
+                      opcode, ownerClassName, methodName, descriptor, isInterface);
+                  // rewrite Vm.isBooted() to AgentInitializer.isAgentStarted(Vm.isBooted())
+                  if ("jdk/internal/misc/VM".equals(ownerClassName)
+                      && "isBooted".equals(methodName)) {
+                    super.visitMethodInsn(
+                        Opcodes.INVOKESTATIC,
+                        Type.getInternalName(AgentInitializer.class),
+                        "isAgentStarted",
+                        "(Z)Z",
+                        false);
+                    hookInserted = true;
+                  }
+                }
+              };
+            }
+          };
+
+      cr.accept(cv, 0);
+
+      return hookInserted ? cw.toByteArray() : null;
+    }
+  }
 }

settings.gradle.kts

@@ -121,6 +121,7 @@ include(":javaagent-bootstrap")
 include(":javaagent-extension-api")
 include(":javaagent-tooling")
 include(":javaagent-tooling:javaagent-tooling-java9")
+include(":javaagent-tooling:jdk18-testing")
 include(":javaagent-internal-logging-application")
 include(":javaagent-internal-logging-simple")
 include(":javaagent")