How to Use Java Agents with ASM Bytecode Framework?

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.

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
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");
}
}
java
package org.xunche.agent;
public class HelloAgent {
public static void premain(String args) {
System.out.println("Hello Agent: " + args);
}
}
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/
java
java -javaagent:hello-agent.jar=xunche org/xunche/app/HelloWorld
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:

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));
}
}
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);
}
}
}
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);
}
}
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 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.
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));
}
}
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;
}
}
}
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 + '\'' +
'}';
}
}
}
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();
}
}
}
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.

Original Source:

--

--

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store