Add a jpms opener for HostNameResolver cache

Author: amarziali

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/net/HostNameResolver.java

@@ -7,39 +7,47 @@
 import java.net.InetAddress;
 
 public final class HostNameResolver {
-  private static final MethodHandle HOLDER_GET;
-  private static final MethodHandle HOSTNAME_GET;
+  private static volatile MethodHandle HOLDER_GET;
+  private static volatile MethodHandle HOSTNAME_GET;
 
   private static final DDCache<String, String> HOSTNAME_CACHE = DDCaches.newFixedSizeCache(64);
 
-  static {
-    MethodHandle holderTmp = null, hostnameTmp = null;
-    try {
-      final ClassLoader cl = HostNameResolver.class.getClassLoader();
-      final MethodHandles methodHandles = new MethodHandles(cl);
+  private HostNameResolver() {}
+
+  public static void tryInitialize() {
+    if (HOLDER_GET != null) {
+      return; // fast path: already initialized
+    }
+    synchronized (HostNameResolver.class) {
+      if (HOLDER_GET != null) {
+        return; // double-check: another thread just succeeded
+      }
+      MethodHandle holderTmp = null, hostnameTmp = null;
+      try {
+        final ClassLoader cl = HostNameResolver.class.getClassLoader();
+        final MethodHandles methodHandles = new MethodHandles(cl);
 
-      final Class<?> holderClass =
-          Class.forName("java.net.InetAddress$InetAddressHolder", false, cl);
-      holderTmp = methodHandles.method(InetAddress.class, "holder");
-      if (holderTmp != null) {
-        hostnameTmp = methodHandles.method(holderClass, "getHostName");
+        final Class<?> holderClass =
+            Class.forName("java.net.InetAddress$InetAddressHolder", false, cl);
+        holderTmp = methodHandles.method(InetAddress.class, "holder");
+        if (holderTmp != null) {
+          hostnameTmp = methodHandles.method(holderClass, "getHostName");
+        }
+      } catch (Throwable ignored) {
+        holderTmp = null;
       }
-    } catch (Throwable ignored) {
-      holderTmp = null;
-    } finally {
+      // volatile writes ensure visibility to other threads
       if (holderTmp != null && hostnameTmp != null) {
-        HOLDER_GET = holderTmp;
         HOSTNAME_GET = hostnameTmp;
-      } else {
-        HOLDER_GET = null;
-        HOSTNAME_GET = null;
+        HOLDER_GET = holderTmp; // written last: signals successful initialization
       }
     }
   }
 
-  private HostNameResolver() {}
-
   static String getAlreadyResolvedHostName(InetAddress address) {
+    if (HOLDER_GET == null) {
+      tryInitialize();
+    }
     if (HOLDER_GET == null) {
       return null;
     }

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/net/JpmsInetAddressHelper.java

@@ -0,0 +1,9 @@
+package datadog.trace.bootstrap.instrumentation.java.net;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+
+public class JpmsInetAddressHelper {
+  public static final AtomicBoolean OPENED = new AtomicBoolean(false);
+
+  private JpmsInetAddressHelper() {}
+}

dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie

@@ -59,6 +59,7 @@
 0 java.lang.VirtualThread
 0 java.net.http.*
 0 java.net.HttpURLConnection
+0 java.net.InetAddress
 0 java.net.Socket
 0 java.net.URL
 0 java.nio.DirectByteBuffer

dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java/datadog/trace/instrumentation/httpclient/JpmsInetAddressInstrumentation.java

@@ -0,0 +1,33 @@
+package datadog.trace.instrumentation.httpclient;
+
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+
+import com.google.auto.service.AutoService;
+import datadog.environment.JavaVirtualMachine;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+
+@AutoService(InstrumenterModule.class)
+public class JpmsInetAddressInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.HasMethodAdvice, Instrumenter.ForSingleType {
+
+  public JpmsInetAddressInstrumentation() {
+    super("java-net");
+  }
+
+  @Override
+  public boolean isEnabled() {
+    return super.isEnabled() && JavaVirtualMachine.isJavaVersionAtLeast(9);
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "java.net.InetAddress";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    // it does not work with typeInitializer()
+    transformer.applyAdvice(isConstructor(), packageName + ".JpmsInetAddressClearanceAdvice");
+  }
+}

dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/JpmsInetAddressClearanceAdvice.java

@@ -0,0 +1,21 @@
+package datadog.trace.instrumentation.httpclient;
+
+import datadog.trace.bootstrap.instrumentation.java.net.HostNameResolver;
+import datadog.trace.bootstrap.instrumentation.java.net.JpmsInetAddressHelper;
+import java.net.InetAddress;
+import net.bytebuddy.asm.Advice;
+
+public class JpmsInetAddressClearanceAdvice {
+  @Advice.OnMethodExit(suppress = Throwable.class)
+  public static void openOnReturn() {
+    if (JpmsInetAddressHelper.OPENED.compareAndSet(false, true)) {
+      // This call needs imperatively to be done from the same module we're adding opens to,
+      // because the JDK checks that the caller belongs to the same module.
+      // The code of this advice is inlined into the constructor of InetAddress (java.base),
+      // so it will work. Moving the same call to a helper class won't.
+      InetAddress.class.getModule().addOpens("java.net", HostNameResolver.class.getModule());
+      // Now that java.net is open for deep reflection, initialize the HostNameResolver handles
+      HostNameResolver.tryInitialize();
+    }
+  }
+}

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

@@ -0,0 +1,40 @@
+package datadog.trace.instrumentation.httpclient
+
+import datadog.trace.agent.test.InstrumentationSpecification
+import datadog.trace.bootstrap.instrumentation.java.net.HostNameResolver
+
+class JpmsInetAddressDisabledForkedTest extends InstrumentationSpecification {
+
+  @Override
+  protected void configurePreAgent() {
+    super.configurePreAgent()
+    // Disable the JPMS instrumentation so java.net is NOT opened for deep reflection.
+    // HostNameResolver will be unable to bypass the IP→hostname cache and will fall back
+    // to the cache keyed by IP address.
+    injectSysConfig("dd.trace.java-net.enabled", "false")
+  }
+
+  /**
+   * 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.
+   *
+   * This is the broken behaviour that the JPMS instrumentation is designed to fix.
+   */
+  def "without JPMS instrumentation, IP cache causes stale hostname to be returned"() {
+    given:
+    def ip = [192, 0, 2, 2] as byte[]  // different subnet from the enabled-test to avoid cross-test cache pollution
+    def addr1 = InetAddress.getByAddress("service1.example.com", ip)
+    // Prime the IP→hostname cache with service1's hostname
+    HostNameResolver.hostName(addr1, "192.0.2.2")
+
+    when: "a second service with the same IP but a different hostname is resolved"
+    def addr2 = InetAddress.getByAddress("service2.example.com", ip)
+    def result = HostNameResolver.hostName(addr2, "192.0.2.2")
+
+    then: "the stale cached hostname of service1 is returned instead of service2's"
+    result == "service1.example.com"
+  }
+}

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

@@ -0,0 +1,33 @@
+package datadog.trace.instrumentation.httpclient
+
+import datadog.trace.agent.test.InstrumentationSpecification
+import datadog.trace.bootstrap.instrumentation.java.net.HostNameResolver
+
+class JpmsInetAddressForkedTest extends InstrumentationSpecification {
+
+  /**
+   * 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).
+   *
+   * 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.
+   */
+  def "instrumentation opens java.net so hostname is resolved correctly when IP is shared"() {
+    given:
+    def ip = [192, 0, 2, 1] as byte[]  // TEST-NET, will never appear in real DNS cache
+    def addr1 = InetAddress.getByAddress("service1.example.com", ip)
+    // Warm the IP→hostname cache with service1's hostname
+    HostNameResolver.hostName(addr1, "192.0.2.1")
+
+    when: "a second service with the same IP but different hostname is resolved"
+    def addr2 = InetAddress.getByAddress("service2.example.com", ip)
+    def result = HostNameResolver.hostName(addr2, "192.0.2.1")
+
+    then: "the hostname of addr2 is returned, not the cached hostname of addr1"
+    result == "service2.example.com"
+  }
+}

metadata/supported-configurations.json

@@ -6737,6 +6737,14 @@
         "aliases": ["DD_TRACE_INTEGRATION_JAVA_HTTP_CLIENT_ENABLED", "DD_INTEGRATION_JAVA_HTTP_CLIENT_ENABLED"]
       }
     ],
+    "DD_TRACE_JAVA_NET_ENABLED": [
+      {
+        "version": "A",
+        "type": "boolean",
+        "default": "true",
+        "aliases": ["DD_TRACE_INTEGRATION_JAVA_NET_ENABLED", "DD_INTEGRATION_JAVA_NET_ENABLED"]
+      }
+    ],
     "DD_TRACE_TRACE_FFM_ANALYTICS_ENABLED": [
       {
         "version": "A",