Java Agents and Instrumentation: Bytecode Manipulation at Runtime Explained

import java.lang.instrument.Instrumentation;

/**
 * Entry point for a startup-time Java agent.
 *
 * Compile this class, package it in a JAR, and declare it in MANIFEST.MF:
 *   Premain-Class: TimingAgent
 *   Can-Retransform-Classes: true
 *
 * Launch with:
 *   java -javaagent:timing-agent.jar -jar your-app.jar
 */
public class TimingAgent {

    /**
     * The JVM calls this BEFORE your application's main() method.
     *
     * @param agentArgs  Any string args passed after '=' in -javaagent:jar=args
     * @param instrumentation  The gateway to all bytecode transformation power
     */
    public static void premain(String agentArgs, Instrumentation instrumentation) {
        System.out.println("[TimingAgent] Agent loaded. JVM classes currently loaded: "
                + instrumentation.getAllLoadedClasses().length);

        // Register our transformer — it will be called for EVERY class the JVM loads
        // from this point forward, including JDK classes if we don't filter carefully
        MethodTimingTransformer transformer = new MethodTimingTransformer(agentArgs);
        instrumentation.addTransformer(transformer, /* canRetransform= */ true);

        System.out.println("[TimingAgent] ClassFileTransformer registered. Ready.");
    }

    /**
     * Called when the agent is attached to a RUNNING JVM via the Attach API.
     * The signature is identical to premain but the JVM lifecycle moment differs.
     */
    public static void agentmain(String agentArgs, Instrumentation instrumentation) {
        // Reuse the same startup logic — safe because we're stateless here
        premain(agentArgs, instrumentation);
    }
}

import org.objectweb.asm.*;
import org.objectweb.asm.commons.AdviceAdapter;
import java.lang.instrument.ClassFileTransformer;
import java.security.ProtectionDomain;

/**
 * Intercepts class loading and injects timing code into every method
 * of classes whose package matches our target prefix.
 *
 * ASM version: 9.6  (org.ow2.asm:asm-commons:9.6)
 */
public class MethodTimingTransformer implements ClassFileTransformer {

    // Only instrument classes under this package to avoid touching JDK internals
    private final String targetPackagePrefix;

    public MethodTimingTransformer(String agentArgs) {
        // agentArgs might be something like "io/codeforge/demo"
        this.targetPackagePrefix = (agentArgs != null && !agentArgs.isBlank())
                ? agentArgs.replace('.', '/')
                : "io/codeforge"; // default to our own package
    }

    @Override
    public byte[] transform(
            ClassLoader loader,
            String className,           // Internally formatted: "com/example/MyClass"
            Class<?> classBeingRedefined,
            ProtectionDomain protectionDomain,
            byte[] originalClassBytes) {

        // Fast-path: skip anything that's not our target to keep overhead minimal
        if (className == null || !className.startsWith(targetPackagePrefix)) {
            return null; // null = "I didn't change anything, use original bytes"
        }

        try {
            ClassReader classReader = new ClassReader(originalClassBytes);
            // COMPUTE_FRAMES tells ASM to recalculate stack frames automatically
            // This is essential — manual frame calculation is error-prone
            ClassWriter classWriter = new ClassWriter(
                    classReader, ClassWriter.COMPUTE_FRAMES);

            // Wrap the writer with our visitor that injects timing into each method
            TimingClassVisitor timingVisitor = new TimingClassVisitor(
                    Opcodes.ASM9, classWriter, className);

            // EXPAND_FRAMES is required when using AdviceAdapter with COMPUTE_FRAMES
            classReader.accept(timingVisitor, ClassReader.EXPAND_FRAMES);

            System.out.println("[TimingAgent] Instrumented: " + className);
            return classWriter.toByteArray(); // Return the rewritten bytecode

        } catch (Exception transformationError) {
            // NEVER let a transformer exception crash the JVM — log and return null
            System.err.println("[TimingAgent] Failed to instrument " + className
                    + ": " + transformationError.getMessage());
            return null; // Fall back to the original unmodified class
        }
    }

    // ── Inner class: visits each class and delegates method visiting ─────────

    static class TimingClassVisitor extends ClassVisitor {

        private final String ownerClassName;

        TimingClassVisitor(int api, ClassVisitor delegate, String ownerClassName) {
            super(api, delegate);
            this.ownerClassName = ownerClassName;
        }

        @Override
        public MethodVisitor visitMethod(
                int access, String methodName, String descriptor,
                String signature, String[] exceptions) {

            // Get the default MethodVisitor from the ClassWriter chain
            MethodVisitor originalVisitor = super.visitMethod(
                    access, methodName, descriptor, signature, exceptions);

            // Skip abstract and native methods — they have no bytecode body to visit
            boolean isAbstract = (access & Opcodes.ACC_ABSTRACT) != 0;
            boolean isNative   = (access & Opcodes.ACC_NATIVE) != 0;
            if (isAbstract || isNative || originalVisitor == null) {
                return originalVisitor;
            }

            // Wrap with our timing advisor
            return new TimingMethodAdvisor(
                    Opcodes.ASM9, originalVisitor, access,
                    methodName, descriptor, ownerClassName);
        }
    }

    // ── Inner class: injects nanoTime calls at method entry and exit ──────────

    static class TimingMethodAdvisor extends AdviceAdapter {

        private final String ownerClass;
        private final String methodName;
        private int startTimeLocalVarIndex; // index of our injected local variable

        TimingMethodAdvisor(int api, MethodVisitor delegate, int access,
                            String methodName, String descriptor, String ownerClass) {
            super(api, delegate, access, methodName, descriptor);
            this.ownerClass  = ownerClass;
            this.methodName  = methodName;
        }

        @Override
        protected void onMethodEnter() {
            // Inject: long __startTime = System.nanoTime();
            // INVOKESTATIC pushes the long result onto the stack
            mv.visitMethodInsn(
                    Opcodes.INVOKESTATIC,
                    "java/lang/System",
                    "nanoTime",
                    "()J",
                    false);

            // Store that long into a new local variable
            // newLocal() allocates the next available slot in the local variable table
            startTimeLocalVarIndex = newLocal(Type.LONG_TYPE);
            mv.visitVarInsn(Opcodes.LSTORE, startTimeLocalVarIndex);
        }

        @Override
        protected void onMethodExit(int opcode) {
            // This is called for EVERY exit: RETURN, IRETURN, ARETURN, ATHROW, etc.
            // AdviceAdapter handles all the variants — we just implement this once

            // Inject: long __elapsed = System.nanoTime() - __startTime;
            mv.visitMethodInsn(
                    Opcodes.INVOKESTATIC, "java/lang/System",
                    "nanoTime", "()J", false);
            mv.visitVarInsn(Opcodes.LLOAD, startTimeLocalVarIndex);
            mv.visitInsn(Opcodes.LSUB); // stack now has elapsed nanoseconds (long)

            // Inject: TimingReporter.record("io/codeforge/Foo#doWork", elapsed);
            // We push the method label as a constant string first
            mv.visitLdcInsn(ownerClass + "#" + methodName);
            // Swap so the stack order is (String label, long elapsed) matching the descriptor
            // Actually we need to restructure: push label first, then elapsed
            // Easiest: store elapsed, push label, reload elapsed
            int elapsedLocalVarIndex = newLocal(Type.LONG_TYPE);
            mv.visitVarInsn(Opcodes.LSTORE, elapsedLocalVarIndex);
            mv.visitLdcInsn(ownerClass + "#" + methodName);
            mv.visitVarInsn(Opcodes.LLOAD, elapsedLocalVarIndex);

            mv.visitMethodInsn(
                    Opcodes.INVOKESTATIC,
                    "io/codeforge/agent/TimingReporter",
                    "record",
                    "(Ljava/lang/String;J)V",
                    false);
        }
    }
}

import com.sun.tools.attach.VirtualMachine;
import com.sun.tools.attach.VirtualMachineDescriptor;
import java.util.List;

/**
 * Attaches a Java agent to a RUNNING JVM without restarting it.
 *
 * Requires: JDK (not just JRE), and the target JVM must allow self-attach or
 * be a separate process. Run this as a standalone program.
 *
 * Compile: javac --add-exports java.base/sun.nio.ch=ALL-UNNAMED LiveAttachDemo.java
 * Run:     java -cp .:$JAVA_HOME/lib/tools.jar LiveAttachDemo <target-pid> <agent-jar-path>
 *
 * Java 9+ note: tools.jar is gone; the Attach API is in jdk.attach module:
 *   java --add-modules jdk.attach LiveAttachDemo <pid> <agent.jar>
 */
public class LiveAttachDemo {

    public static void main(String[] args) throws Exception {
        if (args.length < 2) {
            System.err.println("Usage: LiveAttachDemo <target-pid> <path-to-agent.jar>");
            System.exit(1);
        }

        String targetPid   = args[0];
        String agentJarPath = args[1];

        // List all JVMs on this machine so the user can verify the PID
        System.out.println("\n=== JVMs currently running on this machine ===");
        List<VirtualMachineDescriptor> jvmList = VirtualMachine.list();
        for (VirtualMachineDescriptor descriptor : jvmList) {
            System.out.printf("  PID: %-8s  Display: %s%n",
                    descriptor.id(), descriptor.displayName());
        }
        System.out.println();

        // Attach to the target JVM — this opens a socket to its attach listener
        System.out.println("[Attacher] Attaching to PID: " + targetPid);
        VirtualMachine targetJvm = VirtualMachine.attach(targetPid);

        try {
            // The JVM will call agentmain() in our TimingAgent class
            // The second argument (agentOptions) becomes the agentArgs param in agentmain
            String agentOptions = "io/codeforge/demo"; // package prefix to instrument
            targetJvm.loadAgent(agentJarPath, agentOptions);
            System.out.println("[Attacher] Agent loaded successfully into PID " + targetPid);

        } finally {
            // Always detach to close the socket — a leaked attach can prevent
            // future attaches and consume file descriptors
            targetJvm.detach();
            System.out.println("[Attacher] Detached cleanly.");
        }
    }
}

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0
             http://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>

    <groupId>io.codeforge</groupId>
    <artifactId>timing-agent</artifactId>
    <version>1.0.0</version>
    <packaging>jar</packaging>

    <dependencies>
        <!-- ASM core for ClassReader / ClassWriter -->
        <dependency>
            <groupId>org.ow2.asm</groupId>
            <artifactId>asm</artifactId>
            <version>9.6</version>
        </dependency>
        <!-- ASM Commons for AdviceAdapter (handles all exit opcodes for us) -->
        <dependency>
            <groupId>org.ow2.asm</groupId>
            <artifactId>asm-commons</artifactId>
            <version>9.6</version>
        </dependency>
    </dependencies>

    <build>
        <plugins>
            <plugin>
                <groupId>org.apache.maven.plugins</groupId>
                <artifactId>maven-shade-plugin</artifactId>
                <version>3.5.1</version>
                <executions>
                    <execution>
                        <phase>package</phase>
                        <goals><goal>shade</goal></goals>
                        <configuration>
                            <relocations>
                                <!--
                                  Relocate ASM so it doesn't clash with
                                  the target app's own ASM version.
                                  Spring, Hibernate and CGLib all bundle ASM.
                                -->
                                <relocation>
                                    <pattern>org.objectweb.asm</pattern>
                                    <shadedPattern>io.codeforge.agent.shaded.asm</shadedPattern>
                                </relocation>
                            </relocations>
                            <transformers>
                                <transformer implementation="org.apache.maven.plugins.shade.resource.ManifestResourceTransformer">
                                    <manifestEntries>
                                        <!-- The class containing premain() -->
                                        <Premain-Class>io.codeforge.agent.TimingAgent</Premain-Class>
                                        <!-- The class containing agentmain() for attach-time loading -->
                                        <Agent-Class>io.codeforge.agent.TimingAgent</Agent-Class>
                                        <!--
                                          Must be true to call instrumentation.retransformClasses()
                                          on already-loaded classes after attach
                                        -->
                                        <Can-Retransform-Classes>true</Can-Retransform-Classes>
                                        <!--
                                          Must be true to call instrumentation.redefineClasses()
                                          (full class replacement, more disruptive than retransform)
                                        -->
                                        <Can-Redefine-Classes>false</Can-Redefine-Classes>
                                    </manifestEntries>
                                </transformer>
                            </transformers>
                        </configuration>
                    </execution>
                </executions>
            </plugin>
        </plugins>
    </build>
</project>