Initialize java.net.InetAddress after VM and agent are initialized

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 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,17 @@ public static void delayedStartHook() throws Exception {
           @Override
           public Void run() {
             agentStarter.start();
+            agentStarted = true;
             return null;
           }
         });
   }
 
+  @SuppressWarnings("unused")
+  public static boolean isAgentStarted() {
+    return agentStarted;
+  }
+
   public static ClassLoader getExtensionsClassLoader() {
     // agentStarter can be null when running tests
     return agentStarter != null ? agentStarter.getExtensionClassLoader() : null;

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

@@ -20,6 +20,7 @@
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.ClassWriter;
+import org.objectweb.asm.Label;
 import org.objectweb.asm.MethodVisitor;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.Type;
@@ -68,6 +69,8 @@ public boolean delayStart() {
 
   @Override
   public void start() {
+    init();
+
     EarlyInitAgentConfig earlyConfig = EarlyInitAgentConfig.create();
     extensionClassLoader = createExtensionClassLoader(getClass().getClassLoader(), earlyConfig);
 
@@ -115,6 +118,29 @@ public void start() {
     }
   }
 
+  private void init() {
+    instrumentInetAddress();
+  }
+
+  private void instrumentInetAddress() {
+    InetAddressClassFileTransformer transformer = new InetAddressClassFileTransformer();
+    instrumentation.addTransformer(transformer, true);
+
+    try {
+      Class<?> clazz = Class.forName("java.net.InetAddress", false, null);
+      if (transformer.transformed) {
+        // InetAddress was loaded and got transformed
+        return;
+      }
+      // InetAddress was already loaded before we set up transformer
+      instrumentation.retransformClasses(clazz);
+    } catch (ClassNotFoundException | UnmodifiableClassException ignore) {
+      // ignore
+    } finally {
+      instrumentation.removeTransformer(transformer);
+    }
+  }
+
   @SuppressWarnings("SystemOut")
   private static void logUnrecognizedLoggerImplWarning(String loggerImplementationName) {
     System.err.println(
@@ -180,4 +206,117 @@ public void visitCode() {
       return hookInserted ? cw.toByteArray() : null;
     }
   }
+
+  private static class InetAddressClassFileTransformer implements ClassFileTransformer {
+    boolean hookInserted = false;
+    boolean transformed = false;
+    boolean wrapperMethodCreated = false;
+
+    private static void createWrapperMethod(ClassWriter cw) {
+      /*
+       private static boolean isAgentAndVmBooted();
+       Code:
+          0: invokestatic  #X                 // Method io/opentelemetry/javaagent/bootstrap/AgentInitializer/isAgentStarted:()Z
+          3: ifeq          16
+          6: invokestatic  #Y                 // Method jdk/internal/misc/VM.isBooted:()Z
+          9: ifeq          16
+         12: iconst_1
+         13: goto          17
+         16: iconst_0
+         17: ireturn
+      */
+
+      String descriptor = Type.getMethodDescriptor(Type.BOOLEAN_TYPE);
+      Label elseLabel = new Label();
+      Label endLabel = new Label();
+
+      MethodVisitor mv =
+          cw.visitMethod(
+              Opcodes.ACC_PRIVATE | Opcodes.ACC_STATIC,
+              "isAgentAndVmBooted",
+              descriptor,
+              null,
+              null);
+      mv.visitCode();
+
+      mv.visitMethodInsn(
+          Opcodes.INVOKESTATIC,
+          Type.getInternalName(AgentInitializer.class),
+          "isAgentStarted",
+          descriptor,
+          false);
+      mv.visitJumpInsn(Opcodes.IFEQ, elseLabel);
+      mv.visitMethodInsn(
+          Opcodes.INVOKESTATIC, "jdk/internal/misc/VM", "isBooted", descriptor, false);
+      mv.visitJumpInsn(Opcodes.IFEQ, elseLabel);
+      mv.visitInsn(Opcodes.ICONST_1);
+      mv.visitJumpInsn(Opcodes.GOTO, endLabel);
+      mv.visitLabel(elseLabel);
+      mv.visitInsn(Opcodes.ICONST_0);
+      mv.visitLabel(endLabel);
+      mv.visitInsn(Opcodes.IRETURN);
+
+      mv.visitMaxs(0, 0);
+      mv.visitEnd();
+    }
+
+    @Override
+    public byte[] transform(
+        ClassLoader loader,
+        String className,
+        Class<?> classBeingRedefined,
+        ProtectionDomain protectionDomain,
+        byte[] classfileBuffer) {
+      if (!"java/net/InetAddress".equals(className)) {
+        return null;
+      }
+      transformed = true;
+      ClassReader cr = new ClassReader(classfileBuffer);
+      ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_MAXS);
+      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);
+              // We don't want to patch "jdk.internal.misc.VM.isBooted" call in our wrapper
+              if ("isAgentAndVmBooted".equals(name)) {
+                return mv;
+              }
+              return new MethodVisitor(api, mv) {
+                @Override
+                public void visitMethodInsn(
+                    int opcode,
+                    String ownerClassName,
+                    String methodName,
+                    String descriptor,
+                    boolean isInterface) {
+                  if ("jdk/internal/misc/VM".equals(ownerClassName)
+                      && "isBooted".equals(methodName)) {
+                    // Create wrapper method only once
+                    if (!wrapperMethodCreated) {
+                      createWrapperMethod(cw);
+                      wrapperMethodCreated = true;
+                    }
+                    super.visitMethodInsn(
+                        Opcodes.INVOKESTATIC,
+                        "java/net/InetAddress",
+                        "isAgentAndVmBooted",
+                        "()Z",
+                        isInterface);
+                    hookInserted = true;
+                  } else {
+                    super.visitMethodInsn(
+                        opcode, ownerClassName, methodName, descriptor, isInterface);
+                  }
+                }
+              };
+            }
+          };
+
+      cr.accept(cv, 0);
+
+      return hookInserted ? cw.toByteArray() : null;
+    }
+  }
 }