Improve JPMS module opener

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/module/JpmsHelper.java

@@ -17,23 +17,44 @@ private JpmsHelper() {}
 
   private static final Set<String> TRIGGERS = new HashSet<>();
 
+  private static final Set<String> TRIGGERS_VIEW = unmodifiableSet(TRIGGERS);
+
   private static final ClassValue<AtomicBoolean> HAS_FIRED =
       GenericClassValue.constructing(AtomicBoolean.class);
 
-  public static final Logger LOGGER = LoggerFactory.getLogger(JpmsHelper.class);
+  private static final Logger LOGGER = LoggerFactory.getLogger(JpmsHelper.class);
 
+  /**
+   * Registers trigger class names whose constructors will open their enclosing module. Must be
+   * called at agent startup before instrumentation is applied; not thread-safe.
+   */
   public static void addTriggers(Collection<String> classes) {
     if (classes == null) {
       return;
     }
     TRIGGERS.addAll(classes);
   }
 
+  /** Returns an unmodifiable view of all registered trigger class names. */
   public static Set<String> getAllTriggers() {
-    return unmodifiableSet(TRIGGERS);
+    return TRIGGERS_VIEW;
   }
 
+  /**
+   * Returns {@code true} and atomically marks {@code cls} as opened the first time this is called
+   * for a given class; returns {@code false} on all subsequent calls for the same class.
+   */
   public static boolean shouldBeOpened(Class<?> cls) {
     return HAS_FIRED.get(cls).compareAndSet(false, true);
   }
+
+  /** Called from inlined ByteBuddy advice; logs when module opening fails. */
+  public static void logFailedToOpen(String pkg, Throwable t) {
+    LOGGER.debug("Unable to open package {} to the agent module or unnamed module", pkg, t);
+  }
+
+  /** Called from inlined ByteBuddy advice; logs when a class has no named module. */
+  public static void logNoNamedModule(Class<?> cls) {
+    LOGGER.debug("{} has no named module; skipping module open", cls);
+  }
 }

dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java

@@ -193,19 +193,22 @@ public static ClassFileTransformer installBytebuddyAgent(
     // immediately and we don't have the ability to express dependencies between different
     // instrumentations to control the load order.
     for (InstrumenterModule module : instrumenterModules) {
-      boolean filterAdded = false;
       if (module instanceof ExcludeFilterProvider) {
         ExcludeFilterProvider provider = (ExcludeFilterProvider) module;
         ExcludeFilter.add(provider.excludedClasses());
-        filterAdded = true;
+        if (DEBUG) {
+          log.debug(
+              "Registering exclude-filter module - instrumentation.class={}",
+              module.getClass().getName());
+        }
       }
       if (javaModuleSupported && module instanceof JavaModuleOpenProvider) {
         JpmsHelper.addTriggers(((JavaModuleOpenProvider) module).triggerClasses());
-        filterAdded = true;
-      }
-      if (DEBUG && filterAdded) {
-        log.debug(
-            "Adding filtered classes - instrumentation.class={}", module.getClass().getName());
+        if (DEBUG) {
+          log.debug(
+              "Registering JPMS trigger module - instrumentation.class={}",
+              module.getClass().getName());
+        }
       }
     }
 

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/JavaModuleOpenProvider.java

@@ -3,7 +3,8 @@
 import java.util.Collection;
 
 /**
- * Allows an {@link InstrumenterModule} to possibly open a java module to the unnamed module.
+ * Allows an {@link InstrumenterModule} to open a java module to the agent module and to the unnamed
+ * module of the trigger class's class loader.
  *
  * <p>This is typically used when reflective operations need to be done and the agent cannot assume
  * that the host application has permitted them.

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/InstrumenterIndexTest.groovy

@@ -83,4 +83,26 @@ class InstrumenterIndexTest extends DDSpecification {
     index.instrumentationId(unknownInstrumentation) == -1
     index.transformationId(unknownTransformation) == -1
   }
+
+  def "JavaModuleOpenProvider-only module gets NEEDS_EARLY_LOAD_FLAG"() {
+    given:
+    def jpmsOnlyModule = new TestJpmsOnlyModule()
+
+    when:
+    byte flags = InstrumenterIndex.encodeModuleFlags(jpmsOnlyModule, false)
+
+    then:
+    InstrumenterIndex.decodeModuleNeedsEarlyLoad(flags)
+  }
+}
+
+class TestJpmsOnlyModule extends InstrumenterModule implements JavaModuleOpenProvider {
+  TestJpmsOnlyModule() {
+    super('jpms-test-only')
+  }
+
+  @Override
+  Collection<String> triggerClasses() {
+    return []
+  }
 }

dd-java-agent/instrumentation/java/java-lang/java-lang-9.0/src/main/java/datadog/trace/instrumentation/java/lang/module/JpmsClearanceInstrumentation.java

@@ -1,5 +1,8 @@
 package datadog.trace.instrumentation.java.lang.module;
 
+import static datadog.trace.bootstrap.instrumentation.java.module.JpmsHelper.logFailedToOpen;
+import static datadog.trace.bootstrap.instrumentation.java.module.JpmsHelper.logNoNamedModule;
+import static datadog.trace.bootstrap.instrumentation.java.module.JpmsHelper.shouldBeOpened;
 import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 
 import com.google.auto.service.AutoService;
@@ -13,10 +16,10 @@
 import net.bytebuddy.implementation.bytecode.assign.Assigner;
 
 /**
- * Generic instrumenter module that will advice the constructor of each known class in order to open
- * once their module. This is marked for bootstrap even if it's not for sure, but we cannot know in
- * advance (depends to the instrumented types and today we are instrumenting InetAddress that's in
- * the bootstrap).
+ * Generic instrumenter module that advises the constructor of each registered trigger class to open
+ * its enclosing module once. Marked {@link Instrumenter.ForBootstrap} because some trigger classes
+ * (e.g. {@code InetAddress}) reside in the bootstrap classloader; the annotation is applied
+ * conservatively since the set of trigger classes is not known until runtime.
  */
 @AutoService(InstrumenterModule.class)
 public class JpmsClearanceInstrumentation extends InstrumenterModule
@@ -34,7 +37,7 @@ public boolean isEnabled() {
 
   @Override
   public boolean isApplicable(Set<TargetSystem> enabledSystems) {
-    return true; // not directly linked ot a target system
+    return true; // not directly linked to a target system
   }
 
   @Override
@@ -51,24 +54,24 @@ public static class OpenModuleAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void onExit(@Advice.This(typing = Assigner.Typing.DYNAMIC) Object self) {
       final Class<?> cls = self.getClass();
-      if (JpmsHelper.shouldBeOpened(cls)) {
+      if (shouldBeOpened(cls)) {
         final Module module = cls.getModule();
+        final String pkg = cls.getPackageName();
         if (module != null) {
           try {
-            // This call needs imperatively to be done from the same module we're adding exports
-            // because the jdk is checking that the caller belongs to the same module.
-            // The code of this advice is getting inlined into the constructor of the class
-            // belonging
-            // to that package so it will work. Moving the same to a helper won't.
-            module.addOpens(cls.getPackageName(), JpmsHelper.class.getModule());
+            // This call must be inlined into the constructor of the class belonging to that
+            // package because the JDK verifies the caller belongs to the module being opened.
+            // Moving this to a helper method will not work.
+            module.addOpens(pkg, JpmsHelper.class.getModule());
             final ClassLoader loader = cls.getClassLoader();
             if (loader != null) {
-              module.addOpens(cls.getPackageName(), loader.getUnnamedModule());
+              module.addOpens(pkg, loader.getUnnamedModule());
             }
           } catch (Throwable t) {
-            JpmsHelper.LOGGER.debug(
-                "Unable to open package {} to the unnamed module", cls.getPackageName(), t);
+            logFailedToOpen(pkg, t);
           }
+        } else {
+          logNoNamedModule(cls);
         }
       }
     }

dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/java/datadog/trace/instrumentation/httpclient/JpmsInetAddressDisabledForkedTest.java

@@ -0,0 +1,46 @@
+package datadog.trace.instrumentation.httpclient;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assumptions.assumeFalse;
+
+import datadog.environment.JavaVirtualMachine;
+import datadog.trace.agent.test.AbstractInstrumentationTest;
+import datadog.trace.bootstrap.instrumentation.java.net.HostNameResolver;
+import datadog.trace.junit.utils.config.WithConfig;
+import java.net.InetAddress;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.EnabledForJreRange;
+import org.junit.jupiter.api.condition.JRE;
+
+// Forked test: runs in an isolated JVM with JPMS instrumentation DISABLED.
+// --illegal-access=deny is only enforced from Java 16 onward.
+@EnabledForJreRange(min = JRE.JAVA_16)
+@WithConfig(key = "trace.java-module.enabled", value = "false")
+class JpmsInetAddressDisabledForkedTest extends AbstractInstrumentationTest {
+
+  /**
+   * Verifies the fallback behaviour when the JPMS instrumentation is disabled: HostNameResolver
+   * cannot reflectively read the pre-set hostname from InetAddress and falls back to a cache keyed
+   * by IP address. As a result, once a hostname has been cached for a given IP, every subsequent
+   * lookup for that IP returns the first cached value, even when the InetAddress object carries a
+   * different hostname.
+   *
+   * <p>This is the broken behaviour that the JPMS instrumentation is designed to fix.
+   */
+  @Test
+  void withoutJpmsInstrumentationIpCausesStaleHostnameToBeReturned() throws Exception {
+    assumeFalse(JavaVirtualMachine.isJ9(), "Does not work on J9");
+    // different subnet from the enabled-test to avoid cross-test cache pollution
+    byte[] ip = {(byte) 192, 0, 2, 2};
+    InetAddress addr1 = InetAddress.getByAddress("service1.example.com", ip);
+    // Prime the IP→hostname cache with service1's hostname
+    HostNameResolver.hostName(addr1, "192.0.2.2");
+
+    // a second service with the same IP but a different hostname is resolved
+    InetAddress addr2 = InetAddress.getByAddress("service2.example.com", ip);
+    String result = HostNameResolver.hostName(addr2, "192.0.2.2");
+
+    // the stale cached hostname of service1 is returned instead of service2's
+    assertEquals("service1.example.com", result);
+  }
+}

dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/java/datadog/trace/instrumentation/httpclient/JpmsInetAddressForkedTest.java

@@ -0,0 +1,42 @@
+package datadog.trace.instrumentation.httpclient;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assumptions.assumeFalse;
+
+import datadog.environment.JavaVirtualMachine;
+import datadog.trace.agent.test.AbstractInstrumentationTest;
+import datadog.trace.bootstrap.instrumentation.java.net.HostNameResolver;
+import java.net.InetAddress;
+import org.junit.jupiter.api.Test;
+
+// Forked test: runs in an isolated JVM with JPMS instrumentation ENABLED (default).
+class JpmsInetAddressForkedTest extends AbstractInstrumentationTest {
+
+  /**
+   * Verifies that the JPMS instrumentation opens java.base/java.net so that HostNameResolver can
+   * bypass its IP→hostname cache and return the correct peer.hostname even when multiple services
+   * share a single IP address (e.g. services behind a reverse proxy).
+   *
+   * <p>Without the fix, HostNameResolver cannot reflectively access InetAddress$InetAddressHolder
+   * on Java 9+ and falls back to a cache keyed by IP, causing the first service's hostname to be
+   * returned for all subsequent services on the same IP.
+   */
+  @Test
+  void instrumentationOpensJavaNetSoHostnameIsResolvedCorrectlyWhenIpIsShared() throws Exception {
+    assumeFalse(JavaVirtualMachine.isJ9(), "Does not work on J9");
+
+    // emulate an early initialisation
+    HostNameResolver.hostName(null, "192.0.2.1");
+    byte[] ip = {(byte) 192, 0, 2, 1}; // TEST-NET, will never appear in real DNS cache
+    InetAddress addr1 = InetAddress.getByAddress("service1.example.com", ip);
+    // Warm the IP→hostname cache with service1's hostname
+    HostNameResolver.hostName(addr1, "192.0.2.1");
+
+    // a second service with the same IP but different hostname is resolved
+    InetAddress addr2 = InetAddress.getByAddress("service2.example.com", ip);
+    String result = HostNameResolver.hostName(addr2, "192.0.2.1");
+
+    // the hostname of addr2 is returned, not the cached hostname of addr1
+    assertEquals("service2.example.com", result);
+  }
+}

dd-java-agent/instrumentation/mule-4.5/src/main/java/datadog/trace/instrumentation/mule4/JpmsMuleInstrumentation.java

@@ -7,14 +7,19 @@
 import datadog.trace.agent.tooling.JavaModuleOpenProvider;
 import datadog.trace.agent.tooling.muzzle.Reference;
 import java.util.Collection;
+import java.util.Set;
 
 @AutoService(InstrumenterModule.class)
-public class JpmsMuleInstrumentation extends InstrumenterModule.Tracing
-    implements JavaModuleOpenProvider {
+public class JpmsMuleInstrumentation extends InstrumenterModule implements JavaModuleOpenProvider {
   public JpmsMuleInstrumentation() {
     super("mule", "mule-jpms");
   }
 
+  @Override
+  public boolean isApplicable(Set<TargetSystem> enabledSystems) {
+    return true;
+  }
+
   @Override
   public Reference[] additionalMuzzleReferences() {
     return new Reference[] {