Review feedback

DataDog · Aug 2, 2024 · c426b28 · c426b28
1 parent 92df099
commit c426b28
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 10 deletions.
diff --git a/...ent/agent-jmxfetch/src/main/java/datadog/trace/agent/jmxfetch/AgentConnectionFactory.java b/...ent/agent-jmxfetch/src/main/java/datadog/trace/agent/jmxfetch/AgentConnectionFactory.java
@@ -19,8 +19,8 @@ public class AgentConnectionFactory implements ConnectionFactory {
   public AgentConnectionFactory() {}
 
   @Override
-  public Connection createConnection(Map<String, Object> map) throws IOException {
-    Object mbeanServerClass = map.get("mbean_server_class");
+  public Connection createConnection(Map<String, Object> connectionParams) throws IOException {
+    Object mbeanServerClass = connectionParams.get("mbean_server_class");
     if (mbeanServerClass != null) {
       MBeanServer mbeanServer = MBeanServerRegistry.getServer(mbeanServerClass.toString());
       if (mbeanServer != null) {
@@ -31,6 +31,6 @@ public Connection createConnection(Map<String, Object> map) throws IOException {
           mbeanServerClass);
       return new JvmDirectConnection();
     }
-    return defaultConnectionFactory.createConnection(map);
+    return defaultConnectionFactory.createConnection(connectionParams);
   }
 }
diff --git a/...ent-jmxfetch/src/main/java/datadog/trace/agent/jmxfetch/InitialMBeanServerConnection.java b/...ent-jmxfetch/src/main/java/datadog/trace/agent/jmxfetch/InitialMBeanServerConnection.java
@@ -12,5 +12,7 @@ public InitialMBeanServerConnection(@Nonnull final MBeanServerConnection mbs) th
   }
 
   @Override
-  protected void createConnection() throws IOException {}
+  protected void createConnection() {
+    // already connected
+  }
 }
diff --git a/...in/java/datadog/trace/instrumentation/websphere_jmx/WebsphereSecurityInstrumentation.java b/...in/java/datadog/trace/instrumentation/websphere_jmx/WebsphereSecurityInstrumentation.java
@@ -13,15 +13,18 @@
 import net.bytebuddy.asm.Advice;
 
 /**
- * This instrumentation is needed in order to allow JMXFetch accessing to Admin mbean without
- * needing to change the server security configuration.
+ * Grant JMXFetch access to the WebSphere Admin MBean without changing the server security config.
  */
 @AutoService(InstrumenterModule.class)
 public class WebsphereSecurityInstrumentation extends InstrumenterModule.Tracing
     implements Instrumenter.ForSingleType {
 
+  private final String customBuilder;
+
   public WebsphereSecurityInstrumentation() {
     super("websphere-jmx");
+
+    customBuilder = System.getProperty("javax.management.builder.initial");
   }
 
   @Override
@@ -32,10 +35,9 @@ public String instrumentedType() {
   @Override
   public boolean isEnabled() {
     return super.isEnabled()
-        && "com.ibm.ws.management.PlatformMBeanServerBuilder"
-            .equals(System.getProperty("javax.management.builder.initial"))
-        // a shortcut to Config.get().isJmxFetchIntegrationEnabled("websphere", false) to avoid
-        // loading a Config instance from here
+        && "com.ibm.ws.management.PlatformMBeanServerBuilder".equals(customBuilder)
+        // we must avoid loading the global Config while setting up instrumentation, so use the same
+        // underlying provider call as Config.get().isJmxFetchIntegrationEnabled("websphere", false)
         && ConfigProvider.getInstance()
             .isEnabled(Collections.singletonList("websphere"), "jmxfetch.", ".enabled", false);
   }
@@ -50,6 +52,7 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class DisableSecurityAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void after(@Advice.Return(readOnly = false) boolean securityEnabled) {
+      // only grant access when we know the call is coming from one of our agent threads
       if (AgentThreadFactory.AGENT_THREAD_GROUP == Thread.currentThread().getThreadGroup()) {
         securityEnabled = false;
       }