How to Use Java Agents with ASM Bytecode Framework?

By Zhang Shuai (Xunche), Senior Development Engineer of International Mid-end Business Division.

The author is responsible for the research and development of logistics expression and performance and is keen on the research of middleware technology.

1) Overview

It is surely possible to modify the bytecode that is being loaded using Java agents. There are two ways to realize a Java agent: the premain method and the agentmain method. The first method is realized by the premain function that runs prior to the main function. The second method is realized using the attach API during program runtime.

Before we get into agents, let's take a brief look at Instrumentation. Provided by JDK 1.5, Instrumentation is an API for intercepting class loading events and modifying the bytecode. Its major functions are as follows.

java
public interface Instrumentation {
    //注册一个转换器，类加载事件会被注册的转换器所拦截
     void addTransformer(ClassFileTransformer transformer, boolean canRetransform);
    //重新触发类加载
     void retransformClasses(Class<?>... classes) throws UnmodifiableClassException;
    //直接替换类的定义
     void redefineClasses(ClassDefinition... definitions) throws  ClassNotFoundException, UnmodifiableClassException;
}

2) Premain

The premain method is the most common way to realize an agent. It runs prior to the main function. To run the agent, pack it into a JAR package and modify the command as follows.

java
java -javaagent:agent.jar=xunche HelloWorld

There are two overriding premain functions. During startup, JVM initially attempts to call the first function. If this function does not exist, JVM calls the second function.

java
public static void premain(String agentArgs, Instrumentation inst);
public static void premain(String agentArgs);

2.1 A Simple Example

The following example explains how to use the premain function in a program. First, prepare test code, which is designed to run the main function and output "hello world".

java
package org.xunche.app;
public class HelloWorld {
    public static void main(String[] args) {
        System.out.println("Hello World");
    }
}

Next, prepare agent code, which is designed to run the premain function and output the input parameters.

java
package org.xunche.agent;
public class HelloAgent {
  public static void premain(String args) {
    System.out.println("Hello Agent:  " + args);
  }
}

To run the agent, set the value of Premain-Class in the META-INF/MANIFEST.MF file as the agent path, and then pack the agent into a JAR package. Alternatively, use IDEA to export the agent as a JAR package.

java
echo 'Premain-Class: org.xunche.agent.HelloAgent' > manifest.mf
javac org/xunche/agent/HelloAgent.java
javac org/xunche/app/HelloWorld.java
jar cvmf manifest.mf hello-agent.jar org/

Next, compile and run the test code. To simplify this process, I have placed the compiled class under the same directory as the agent JAR package.

java
java -javaagent:hello-agent.jar=xunche org/xunche/app/HelloWorld

As the output shows, the premain function in the agent runs prior to the main function.

java
Hello Agent: xunche
Hello World

2.2 A Complicated Example

Now that we have learned about agents through the previous example, let's take a look at the practical usage of an agent. The following example uses the agent to monitor functions. The general idea is as follows:

The agent inserts timestamp recording code at the execution entry and exit of a non-JDK function by using ASM, and after the execution of the function, the time cost can be calculated from the timestamps.

First, let's look at the test code. The main method calls the sayHi function, which outputs "hi, xunche" and sleeps for a random period of time.

java
package org.xunche.app;
public class HelloXunChe {
    public static void main(String[] args) throws InterruptedException {
        HelloXunChe helloXunChe = new HelloXunChe();
        helloXunChe.sayHi();
    }
    public void sayHi() throws InterruptedException {
        System.out.println("hi, xunche");
        sleep();
    }
    public void sleep() throws InterruptedException {
        Thread.sleep((long) (Math.random() * 200));
    }
}

The agent then uses ASM to insert the code to calculate the execution time cost of a function. When JVM loads a class, the code is inserted into each function of the class. See the code as follows, and JDK's built-in ASM is used here. Also, use ASM's official class libraries for the same purpose.

java
package org.xunche.agent;
import jdk.internal.org.objectweb.asm.*;
import jdk.internal.org.objectweb.asm.commons.AdviceAdapter;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
public class TimeAgent {
    public static void premain(String args, Instrumentation instrumentation) {
        instrumentation.addTransformer(new TimeClassFileTransformer());
    }
    private static class TimeClassFileTransformer implements ClassFileTransformer {
        @Override
        public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) {
            if (className.startsWith("java") || className.startsWith("jdk") || className.startsWith("javax") || className.startsWith("sun") || className.startsWith("com/sun")|| className.startsWith("org/xunche/agent")) {
                //return null或者执行异常会执行原来的字节码
                return null;
            }
            System.out.println("loaded class: " + className);
            ClassReader reader = new ClassReader(classfileBuffer);
            ClassWriter writer = new ClassWriter(reader, ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
            reader.accept(new TimeClassVisitor(writer), ClassReader.EXPAND_FRAMES);
            return writer.toByteArray();
        }
    }
    public static class TimeClassVisitor extends ClassVisitor {
        public TimeClassVisitor(ClassVisitor classVisitor) {
            super(Opcodes.ASM5, classVisitor);
        }
        @Override
        public MethodVisitor visitMethod(int methodAccess, String methodName, String methodDesc, String signature, String[] exceptions) {
            MethodVisitor methodVisitor = cv.visitMethod(methodAccess, methodName, methodDesc, signature, exceptions);
            return new TimeAdviceAdapter(Opcodes.ASM5, methodVisitor, methodAccess, methodName, methodDesc);
        }
    }
    public static class TimeAdviceAdapter extends AdviceAdapter {
        private String methodName;
        protected TimeAdviceAdapter(int api, MethodVisitor methodVisitor, int methodAccess, String methodName, String methodDesc) {
            super(api, methodVisitor, methodAccess, methodName, methodDesc);
            this.methodName = methodName;
        }
        @Override
        protected void onMethodEnter() {
            //在方法入口处植入
            if ("<init>".equals(methodName)|| "<clinit>".equals(methodName)) {
                return;
            }
            mv.visitTypeInsn(NEW, "java/lang/StringBuilder");
            mv.visitInsn(DUP);
            mv.visitMethodInsn(INVOKESPECIAL, "java/lang/StringBuilder", "<init>", "()V", false);
            mv.visitVarInsn(ALOAD, 0);
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/Object", "getClass", "()Ljava/lang/Class;", false);
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/Class", "getName", "()Ljava/lang/String;", false);
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/StringBuilder", "append", "(Ljava/lang/String;)Ljava/lang/StringBuilder;", false);
            mv.visitLdcInsn(".");
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/StringBuilder", "append", "(Ljava/lang/String;)Ljava/lang/StringBuilder;", false);
            mv.visitLdcInsn(methodName);
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/StringBuilder", "append", "(Ljava/lang/String;)Ljava/lang/StringBuilder;", false);
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/StringBuilder", "toString", "()Ljava/lang/String;", false);
            mv.visitMethodInsn(INVOKESTATIC, "org/xunche/agent/TimeHolder", "start", "(Ljava/lang/String;)V", false);
        }
        @Override
        protected void onMethodExit(int i) {
            //在方法出口植入
            if ("<init>".equals(methodName) || "<clinit>".equals(methodName)) {
                return;
            }
            mv.visitTypeInsn(NEW, "java/lang/StringBuilder");
            mv.visitInsn(DUP);
            mv.visitMethodInsn(INVOKESPECIAL, "java/lang/StringBuilder", "<init>", "()V", false);
            mv.visitVarInsn(ALOAD, 0);
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/Object", "getClass", "()Ljava/lang/Class;", false);
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/Class", "getName", "()Ljava/lang/String;", false);
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/StringBuilder", "append", "(Ljava/lang/String;)Ljava/lang/StringBuilder;", false);
            mv.visitLdcInsn(".");
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/StringBuilder", "append", "(Ljava/lang/String;)Ljava/lang/StringBuilder;", false);
            mv.visitLdcInsn(methodName);
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/StringBuilder", "append", "(Ljava/lang/String;)Ljava/lang/StringBuilder;", false);
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/StringBuilder", "toString", "()Ljava/lang/String;", false);
            mv.visitVarInsn(ASTORE, 1);
            mv.visitFieldInsn(GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;");
            mv.visitTypeInsn(NEW, "java/lang/StringBuilder");
            mv.visitInsn(DUP);
            mv.visitMethodInsn(INVOKESPECIAL, "java/lang/StringBuilder", "<init>", "()V", false);
            mv.visitVarInsn(ALOAD, 1);
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/StringBuilder", "append", "(Ljava/lang/String;)Ljava/lang/StringBuilder;", false);
            mv.visitLdcInsn(": ");
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/StringBuilder", "append", "(Ljava/lang/String;)Ljava/lang/StringBuilder;", false);
            mv.visitVarInsn(ALOAD, 1);
            mv.visitMethodInsn(INVOKESTATIC, "org/xunche/agent/TimeHolder", "cost", "(Ljava/lang/String;)J", false);
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/StringBuilder", "append", "(J)Ljava/lang/StringBuilder;", false);
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/lang/StringBuilder", "toString", "()Ljava/lang/String;", false);
            mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/PrintStream", "println", "(Ljava/lang/String;)V", false);
        }
    }
}

The preceding code is a bit long, and you can skip the ASM section. We register a transformer using instrumentation.addTransformer. In the transformer, we have overridden the transform function. The classfileBuffer input parameter in the function is the original bytecode, and the return value is the actual bytecode to be loaded. The onMethodEnter function calls TimeHolder's start function while passing in the current function's name.

The onMethodExit calls TimeHolder's cost function while passing in the current function's name, and outputs the return value of the cost function.

The following shows the code for TimeHolder.

java
package org.xunche.agent;
import java.util.HashMap;
import java.util.Map;
public class TimeHolder {
    private static Map<String, Long> timeCache = new HashMap<>();
    public static void start(String method) {
        timeCache.put(method, System.currentTimeMillis());
    }
    public static long cost(String method) {
        return System.currentTimeMillis() - timeCache.get(method);
    }
}

The agent code is now finished. The ASM part is not the focus in this article, and you may learn more about it in a specific article about ASM subsequent editions. After running the agent, let's look at the current code in the class. Compared to the original test code, each function has been appended with the monitoring code.

java
package org.xunche.app;
import org.xunche.agent.TimeHolder;
public class HelloXunChe {
    public HelloXunChe() {
    }
    public static void main(String[] args) throws InterruptedException {
        TimeHolder.start(args.getClass().getName() + "." + "main");
        HelloXunChe helloXunChe = new HelloXunChe();
        helloXunChe.sayHi();
        HelloXunChe helloXunChe = args.getClass().getName() + "." + "main";
        System.out.println(helloXunChe + ": " + TimeHolder.cost(helloXunChe));
    }
    public void sayHi() throws InterruptedException {
        TimeHolder.start(this.getClass().getName() + "." + "sayHi");
        System.out.println("hi, xunche");
        this.sleep();
        String var1 = this.getClass().getName() + "." + "sayHi";
        System.out.println(var1 + ": " + TimeHolder.cost(var1));
    }
    public void sleep() throws InterruptedException {
        TimeHolder.start(this.getClass().getName() + "." + "sleep");
        Thread.sleep((long)(Math.random() * 200.0D));
        String var1 = this.getClass().getName() + "." + "sleep";
        System.out.println(var1 + ": " + TimeHolder.cost(var1));
    }
}

3) Agentmain

As shown in the previous example, the premain method modifies the bytecode before the startup of an application to achieve desired functions. Different from the premain method, the agentmain method intercepts the loading of a class during the runtime of a Java application through the attach API provided by JDK. The agentmain method is explained in the following example.

3.1 Practice Example

The goal of this example is to realize a tool to remotely collect function calling information from a running Java process. You may be thinking that it sounds like BTrace. In fact, BTrace is realized in the same way. Due to limited time, the code in this example is rough, and the point is to grasp the idea behind it rather than the details of the code.

The implementation idea is as follows:

The agent modifies the bytecode of the functions of a specified class to collect the input parameters and return values of the functions. Then, the agent sends the collected results to the server through a socket.
The server accesses the running Java process through the attach API and loads the agent.
The server loads the agent while specifying the target class and functions.
The server opens a port to receive requests from the target process.

Let's look at the test code first. Every 100 milliseconds, it runs the sayHi function and sleeps for a random period of time.

java
package org.xunche.app;
public class HelloTraceAgent {
    public static void main(String[] args) throws InterruptedException {
        HelloTraceAgent helloTraceAgent = new HelloTraceAgent();
        while (true) {
            helloTraceAgent.sayHi("xunche");
            Thread.sleep(100);
        }
    }
    public String sayHi(String name) throws InterruptedException {
        sleep();
        String hi = "hi, " + name + ", " + System.currentTimeMillis();
        return hi;
    }
    public void sleep() throws InterruptedException {
        Thread.sleep((long) (Math.random() * 200));
    }
}

Next, let's look at the code for the agent. Similar to the previous agent that monitors the time cost of a function, this agent inserts the code at the exit of a function to collect the input parameters and return values of the function by using ASM. Then, the agent sends the collected information to the server through a socket.

java
package org.xunche.agent;
import jdk.internal.org.objectweb.asm.*;
import jdk.internal.org.objectweb.asm.commons.AdviceAdapter;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.security.ProtectionDomain;
public class TraceAgent {
    public static void agentmain(String args, Instrumentation instrumentation) throws ClassNotFoundException, UnmodifiableClassException {
        if (args == null) {
            return;
        }
        int index = args.lastIndexOf(".");
        if (index != -1) {
            String className = args.substring(0, index);
            String methodName = args.substring(index + 1);
            //目标代码已经加载，需要重新触发加载流程，才会通过注册的转换器进行转换
            instrumentation.addTransformer(new TraceClassFileTransformer(className.replace(".", "/"), methodName), true);
            instrumentation.retransformClasses(Class.forName(className));
        }
    }
    public static class TraceClassFileTransformer implements ClassFileTransformer {
        private String traceClassName;
        private String traceMethodName;
        public TraceClassFileTransformer(String traceClassName, String traceMethodName) {
            this.traceClassName = traceClassName;
            this.traceMethodName = traceMethodName;
        }
        @Override
        public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) {
            //过滤掉Jdk、agent、非指定类的方法
            if (className.startsWith("java") || className.startsWith("jdk") || className.startsWith("javax") || className.startsWith("sun")
                    || className.startsWith("com/sun") || className.startsWith("org/xunche/agent") || !className.equals(traceClassName)) {
                //return null会执行原来的字节码
                return null;
            }
            ClassReader reader = new ClassReader(classfileBuffer);
            ClassWriter writer = new ClassWriter(reader, ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
            reader.accept(new TraceVisitor(className, traceMethodName, writer), ClassReader.EXPAND_FRAMES);
            return writer.toByteArray();
        }
    }
    public static class TraceVisitor extends ClassVisitor {
        private String className;
        private String traceMethodName;
        public TraceVisitor(String className, String traceMethodName, ClassVisitor classVisitor) {
            super(Opcodes.ASM5, classVisitor);
            this.className = className;
            this.traceMethodName = traceMethodName;
        }
        @Override
        public MethodVisitor visitMethod(int methodAccess, String methodName, String methodDesc, String signature, String[] exceptions) {
            MethodVisitor methodVisitor = cv.visitMethod(methodAccess, methodName, methodDesc, signature, exceptions);
            if (traceMethodName.equals(methodName)) {
                return new TraceAdviceAdapter(className, methodVisitor, methodAccess, methodName, methodDesc);
            }
            return methodVisitor;
        }
    }
    private static class TraceAdviceAdapter extends AdviceAdapter {
        private final String className;
        private final String methodName;
        private final Type[] methodArgs;
        private final String[] parameterNames;
        private final int[] lvtSlotIndex;
        protected TraceAdviceAdapter(String className, MethodVisitor methodVisitor, int methodAccess, String methodName, String methodDesc) {
            super(Opcodes.ASM5, methodVisitor, methodAccess, methodName, methodDesc);
            this.className = className;
            this.methodName = methodName;
            this.methodArgs = Type.getArgumentTypes(methodDesc);
            this.parameterNames = new String[this.methodArgs.length];
            this.lvtSlotIndex = computeLvtSlotIndices(isStatic(methodAccess), this.methodArgs);
        }
        @Override
        public void visitLocalVariable(String name, String description, String signature, Label start, Label end, int index) {
            for (int i = 0; i < this.lvtSlotIndex.length; ++i) {
                if (this.lvtSlotIndex[i] == index) {
                    this.parameterNames[i] = name;
                }
            }
        }
        @Override
        protected void onMethodExit(int opcode) {
            //排除构造方法和静态代码块
            if ("<init>".equals(methodName) || "<clinit>".equals(methodName)) {
                return;
            }
            if (opcode == RETURN) {
                push((Type) null);
            } else if (opcode == LRETURN || opcode == DRETURN) {
                dup2();
                box(Type.getReturnType(methodDesc));
            } else {
                dup();
                box(Type.getReturnType(methodDesc));
            }
            Type objectType = Type.getObjectType("java/lang/Object");
            push(lvtSlotIndex.length);
            newArray(objectType);
            for (int j = 0; j < lvtSlotIndex.length; j++) {
                int index = lvtSlotIndex[j];
                Type type = methodArgs[j];
                dup();
                push(j);
                mv.visitVarInsn(ALOAD, index);
                box(type);
                arrayStore(objectType);
            }
            visitLdcInsn(className.replace("/", "."));
            visitLdcInsn(methodName);
            mv.visitMethodInsn(INVOKESTATIC, "org/xunche/agent/Sender", "send", "(Ljava/lang/Object;[Ljava/lang/Object;Ljava/lang/String;Ljava/lang/String;)V", false);
        }
        private static int[] computeLvtSlotIndices(boolean isStatic, Type[] paramTypes) {
            int[] lvtIndex = new int[paramTypes.length];
            int nextIndex = isStatic ? 0 : 1;
            for (int i = 0; i < paramTypes.length; ++i) {
                lvtIndex[i] = nextIndex;
                if (isWideType(paramTypes[i])) {
                    nextIndex += 2;
                } else {
                    ++nextIndex;
                }
            }
            return lvtIndex;
        }
        private static boolean isWideType(Type aType) {
            return aType == Type.LONG_TYPE || aType == Type.DOUBLE_TYPE;
        }
        private static boolean isStatic(int access) {
            return (access & 8) > 0;
        }
    }
}

The preceding code shows the code for the agent. The onMethodExit function obtains the request parameters and response parameters and calls Sender's send function. In this example, the idea about accessing local variable tables comes from the LocalVariableTableParameterNameDiscoverer of Spring.

Next, let's look at the code for Sender.

java
public class Sender {
    private static final int SERVER_PORT = 9876;
    public static void send(Object response, Object[] request, String className, String methodName) {
        Message message = new Message(response, request, className, methodName);
        try {
            Socket socket = new Socket("localhost", SERVER_PORT);
            socket.getOutputStream().write(message.toString().getBytes());
            socket.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
    private static class Message {
        private Object response;
        private Object[] request;
        private String className;
        private String methodName;
        public Message(Object response, Object[] request, String className, String methodName) {
            this.response = response;
            this.request = request;
            this.className = className;
            this.methodName = methodName;
        }
        @Override
        public String toString() {
            return "Message{" +
                    "response=" + response +
                    ", request=" + Arrays.toString(request) +
                    ", className='" + className + '\'' +
                    ", methodName='" + methodName + '\'' +
                    '}';
        }
    }
}

Considering that it's easy to understand the Sender code. Next, let's look at the code for the server. The server opens a port for listening, accepts requests, and loads the agent through the attach API.

java
package org.xunche.app;
import com.sun.tools.attach.AgentInitializationException;
import com.sun.tools.attach.AgentLoadException;
import com.sun.tools.attach.AttachNotSupportedException;
import com.sun.tools.attach.VirtualMachine;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.net.ServerSocket;
import java.net.Socket;
public class TraceAgentMain {
    private static final int SERVER_PORT = 9876;
    public static void main(String[] args) throws IOException, AttachNotSupportedException, AgentLoadException, AgentInitializationException {
        new Server().start();
        //attach的进程
        VirtualMachine vm = VirtualMachine.attach("85241");
        //加载agent并指明需要采集信息的类和方法
        vm.loadAgent("trace-agent.jar", "org.xunche.app.HelloTraceAgent.sayHi");
        vm.detach();
    }
    private static class Server implements Runnable {
        @Override
        public void run() {
            try {
                ServerSocket serverSocket = new ServerSocket(SERVER_PORT);
                while (true) {
                    Socket socket = serverSocket.accept();
                    InputStream input = socket.getInputStream();
                    BufferedReader reader = new BufferedReader(new InputStreamReader(input));
                    System.out.println("receive message:" + reader.readLine());
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        public void start() {
            Thread thread = new Thread(this);
            thread.start();
        }
    }
}

After running the preceding program, note that the server has received the request and response information from org.xunche.app.HelloTraceAgent.sayHi.

java
receive message:Message{response=hi, xunche, 1581599464436, request=[xunche], className='org.xunche.app.HelloTraceAgent', methodName='sayHi'}

4) Summary

This article described the basic usage of Java agents using premain and agentmain methods and implemented a tool to collect the calling information of a function from a running process. In fact, the functions of agents are far more extensive than what is described here. Agents are used in many objects, such as BTrace, Arms, and Springloaded.

The views expressed herein are for reference only and don't necessarily represent the official views of Alibaba Cloud.

Community

How to Use Java Agents with ASM Bytecode Framework?

1) Overview

2) Premain

2.1 A Simple Example

2.2 A Complicated Example

3) Agentmain

3.1 Practice Example

4) Summary

Read previous post:

Read next post:

Alibaba Clouder

You may also like

Comments

Alibaba Clouder

Related Products

Alibaba Cloud Service Mesh

Super App Solution for Telcos

Data Lake Formation

Microservices Engine (MSE)