Java Agent的使用和原理

package com.study.agent;

import com.study.transformer.PrintOutTransformer;
import java.lang.instrument.Instrumentation;

public class MyFirstAgent {

    public static void premain(String agentArgs, Instrumentation inst) {
        System.out.println("启动时加载自定义探针...");
        inst.addTransformer(new PrintOutTransformer(), true);
    }
}

package com.study.transformer;

import javassist.*;
import javassist.bytecode.AccessFlag;
import javassist.bytecode.AnnotationsAttribute;
import javassist.bytecode.annotation.Annotation;
import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.Objects;


public class PrintOutTransformer implements ClassFileTransformer {
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classFileBuffer) throws IllegalClassFormatException {
        //将斜杠替换为点，使之符合Java的全限定名规则
        className = className.replaceAll("/", ".");
        try {
            ClassPool pool = ClassPool.getDefault();
            CtClass cc = pool.get(className);
            CtMethod[] personFlies = cc.getDeclaredMethods();

            AnnotationsAttribute annotationsAttribute = (AnnotationsAttribute) cc.getClassFile().getAttribute(AnnotationsAttribute.visibleTag);
            Annotation annotation = annotationsAttribute.getAnnotation("org.springframework.web.bind.annotation.RestController");
            if(Objects.nonNull(annotation)){
                //借助JavaAssist工具，进行字节码插桩
                for (CtMethod personFly : personFlies) {
                    if (!Objects.equals(personFly.getModifiers(), AccessFlag.PUBLIC)){
                        //如果不是public方法，则跳过
                        continue;
                    }
                    //定义局部变量
                    personFly.addLocalVariable("_start", CtClass.longType);
                    personFly.addLocalVariable("_end", CtClass.longType);
                    //在目标方法前后，插入代码
                    personFly.insertBefore("_start = System.currentTimeMillis();");
                    personFly.insertAfter("_end = System.currentTimeMillis();System.out.println(\"接口耗时：\" + (_end-_start));", true);
                }
                cc.detach();
                return  cc.toBytecode();
            }
            return null;
        } catch (NotFoundException | CannotCompileException | IOException e) {
            throw new RuntimeException(e);
        }
    }
}

package com.study.client.web;

import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import java.util.Random;
import java.util.concurrent.TimeUnit;

@RestController
@RequestMapping("/")
public class IndexController {


    @GetMapping("/hello")
    public String test(){
        System.out.println("假装在执行业务逻辑...");
        try {
            TimeUnit.MILLISECONDS.sleep(new Random().nextInt(2000));
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
        nested();
        System.out.println("假装执行业务逻辑结束");
        return "Hello";
    }

    public static void nested(){
        System.out.println("执行嵌套的内部逻辑");
    }
}

package com.study.agent;

import com.study.transformer.PrintOutTransformer;
import java.lang.instrument.Instrumentation;

public class MyFirstAgent {

    public static void agentmain(String agentArgs, Instrumentation inst) {
        System.out.println("动态加载自定义探针...");
        inst.addTransformer(new PrintOutTransformer(), true);
        Class[] allLoadedClasses = inst.getAllLoadedClasses();
        for (Class clazz : allLoadedClasses) {
            Annotation[] declaredAnnotations = null;
            try {
                declaredAnnotations = clazz.getDeclaredAnnotations();
                if (isAnnotationPresent("org.springframework.web.bind.annotation.RestController", declaredAnnotations)) {
                    //重新转换已经加载的目标类，使其生效
                    inst.retransformClasses(clazz);
                }
            } catch (Exception e) {
                //do nothing
            }
        }
    }

    private static boolean isAnnotationPresent(String annotation, Annotation[] annotationArr) {
        if(Objects.isNull(annotationArr)) {
            return false;
        }
        for (Annotation ann : annotationArr) {
            return annotation.equals(ann.annotationType().getName());
        }
        return false;
    }
}

package com.study.attach;

import com.sun.tools.attach.*;
import java.io.IOException;
import java.util.List;

public class AgentAttachDemo {

    public static void main(String[] args) throws IOException, AttachNotSupportedException, AgentLoadException, AgentInitializationException, InterruptedException {
        List<VirtualMachineDescriptor> vm = VirtualMachine.list();
        for (VirtualMachineDescriptor vmd : vm) {
            if ("com.study.client.JavaAgentDemoClientApplication".equals(vmd.displayName())){
                System.out.println("注入探针");
                VirtualMachine virtualMachine = VirtualMachine.attach(vmd.id());
                virtualMachine.loadAgent("/Users/navy/codes/gitee/study/javaagent-demo/target/javaagent-demo-1.0-SNAPSHOT.jar");
                virtualMachine.detach();
            }
        }
    }
}