【Java】Agent初探

Java Agent是什么

在官方文档上的描述中文摘录如下：

“提供允许 Java 代理检测运行在JVM 上的程序的服务。检测机制是对方法的字节码进行修改。”

翻译过来的意思就是通过JVM提供的钩子（类加载），允许开发者对目标类字节码修改。
给我们带来的好处就是虚拟机级别的AOP，无侵入字节码增强。

Java Agent怎么用

按照文档的操作指示，使用Java Agent需要以下几个步骤:

准备Agent.jar 。代理Jar文件的MANIFEST.MF必须包含 Premain-Class 属性。代理Jar必须实现premain 方法
public static void premain(String agentArgs, Instrumentation inst);

应用端使用命令行启动JavaAgent
-javaagent:jarpath=[= options]

Agent端准备premain()：

package com.example.agent;
import java.lang.instrument.Instrumentation;
public class TinyAgent {
    public static void premain(String agentArgs, Instrumentation instrument) {
        Class[] allLoadedClasses = instrument.getAllLoadedClasses();
        System.out.println("loadedClasses length:" + allLoadedClasses.length);
        //instrumentation.addTransformer(new TinyTransformer(agentArgs));
        System.out.println("TinyAgent premain method");
    }
}

配置清单：MANIFEST.MF

    <build>
        <finalName>agent</finalName>
        <plugins>
            <plugin>
                <groupId>org.apache.maven.plugins</groupId>
                <artifactId>maven-jar-plugin</artifactId>
                <version>3.2.0</version>
                <configuration>
                    <archive>
                       <manifestEntries>
                           <Premain-Class>
                               com.example.agent.TinyAgent
                           </Premain-Class>
                       </manifestEntries>
                    </archive>
                </configuration>
            </plugin>
        </plugins>
    </build>

应用端main()准备，其中Hello.say()简单的打印hello字符串

package com.example.MiniApm;
import com.example.Hello;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;

@SpringBootApplication
public class MiniApmApplication {
    public static void main(String[] args) {
        SpringApplication.run(MiniApmApplication.class, args);
        Hello h = new Hello();
        h.say();
    }
}

命令行启动:
-javaagent:E:\workspace\MiniApm\agent\target\agent.jar

输出结果：

loadedClasses length:488
TinyAgent premain method

  .   ____          _            __ _ _
 /\\ / ___'_ __ _ _(_)_ __  __ _ \ \ \ \
( ( )\___ | '_ | '_| | '_ \/ _` | \ \ \ \
 \\/  ___)| |_)| | | | | || (_| |  ) ) ) )
  '  |____| .__|_| |_|_| |_\__, | / / / /
 =========|_|==============|___/=/_/_/_/
 :: Spring Boot ::        (v2.2.6.RELEASE)
 ............................................
2020-05-05 22:52:36.624  INFO 10800 --- [  restartedMain] o.s.b.w.embedded.tomcat.TomcatWebServer  : Tomcat started on port(s): 8080 (http) with context path ''
2020-05-05 22:52:36.639  INFO 10800 --- [  restartedMain] com.example.MiniApm.MiniApmApplication   : Started MiniApmApplication in 2.64 seconds (JVM running for 3.152)
-------------hello

我们看到控制台输出两行文字

loadedClasses length:488 (1)
TinyAgent premain method (2)

这代表我们的Java Agent机制生效。由于我们启动的是Spring boot应用，第一行加载的类才488个，从数量上看也不匹配，而且debug观察之后根本没有org.springframework的类。

问题：spring相关类哪去了，Agent机制到底生效没有？

带着这个问题去OpenJdk上寻找答案，需要提前下载好源码

Java Agent底层原理

在分析OpenJdk源码之前先看下premain()，顾名思义在main方法之前执行。

public static void premain(String agentArgs, Instrumentation instrument)

接口Instrumentation是Jvm传递过来的，它的实现类IDE点进去就能看到:

package sun.instrument;
public class InstrumentationImpl implements Instrumentation {
    // jvm invoke premain()
    private void loadClassAndCallPremain(String var1, String var2)  {
        this.loadClassAndStartAgent(var1, "premain", var2);
    }
    static {
        // step 1: 加载动态链接库instrument.dll
        System.loadLibrary("instrument");
    }
}

静态代码块先执行，系统在加载动态链接库D:\Java\jdk1.8.0_202\jre\bin\instrument.dll

动态库被加载完成之后会寻找Agent入口方法Agent_OnLoad()，截取部分代码如下：openjdk\jdk\src\share\instrument\InvocationAdapter.c

/*
 *  看下面的注释就能大概了解这个方法是程序的入口
 *  This will be called once for every -javaagent on the command line.
 *  Each call to Agent_OnLoad will create its own agent and agent data.
 *
 *  The argument tail string provided to Agent_OnLoad will be of form
 *  <jarfile>[=<options>]. The tail string is split into the jarfile and
 *  options components. The jarfile manifest is parsed and the value of the
 *  Premain-Class attribute will become the agent's premain class. The jar
 *  file is then added to the system class path, and if the Boot-Class-Path
 *  attribute is present then all relative URLs in the value are processed
 *  to create boot class path segments to append to the boot class path.
 */
JNIEXPORT jint JNICALL
Agent_OnLoad(JavaVM *vm, char *tail, void * reserved) {
    JPLISInitializationError initerror  = JPLIS_INIT_ERROR_NONE;
    jint                     result     = JNI_OK;
    JPLISAgent *             agent      = NULL;

    // 创建JPLISAgent结构体 在JPLISAgent.c中
    initerror = createNewJPLISAgent(vm, &agent);
    premainClass = getAttribute(attributes, "Premain-Class");
    /*
     * Add to the jarfile
     */
    appendClassPath(agent, jarfile);
    bootClassPath = getAttribute(attributes, "Boot-Class-Path");
        switch (initerror) {
    case JPLIS_INIT_ERROR_NONE:
      result = JNI_OK;
      break;
    case JPLIS_INIT_ERROR_CANNOT_CREATE_NATIVE_AGENT:
      result = JNI_ERR;
      fprintf(stderr, "java.lang.instrument/-javaagent: cannot create native agent.\n");
      break;
    case JPLIS_INIT_ERROR_FAILURE:
      result = JNI_ERR;
      fprintf(stderr, "java.lang.instrument/-javaagent: initialization of native agent failed.\n");
      break;
    ......
}

这里面能看到主要就是创建JPLISAgent结构体以及读取MANIFEST.MF配置信息

struct _JPLISAgent {
    JavaVM *          mJVM;                   
    JPLISEnvironment  mNormalEnvironment;     
    JPLISEnvironment  mRetransformEnvironment;
    ...
    jobject           mInstrumentationImpl; //sun/instrument/InstrumentationImpl
    jmethodID         mPremainCaller;         // loadClassAndCallPremain
    jmethodID         mTransform;             // transform
    char const *      mAgentClassName;        /* agent class name */
    char const *      mOptionsString;         /* -javaagent options */
};

createNewJPLISAgent的内部核心逻辑就在initializeJPLISAgent，开启VMInit事件注册回调函数eventHandlerVMInit

JPLISInitializationError
initializeJPLISAgent(   JPLISAgent *    agent,
                        JavaVM *        vm,
                        jvmtiEnv *      jvmtienv) {
    ......
    /* now turn on the VMInit event */
    if ( jvmtierror == JVMTI_ERROR_NONE ) {
        jvmtiEventCallbacks callbacks;
        memset(&callbacks, 0, sizeof(callbacks));
        // 注册VMInit事件回调 在InvocationAdapter.c
        callbacks.VMInit = &eventHandlerVMInit;

        jvmtierror = (*jvmtienv)->SetEventCallbacks( jvmtienv,
                                                     &callbacks,
                                                     sizeof(callbacks));
        check_phase_ret_blob(jvmtierror, JPLIS_INIT_ERROR_FAILURE);
        jplis_assert(jvmtierror == JVMTI_ERROR_NONE);
    }
    ....
}

Jvm事件回调有很多本文相关的主要有两个回调VMInit和ClassFileLoadHook，枚举如下：

typedef struct {
                              /*   50 : VM Initialization Event */
    jvmtiEventVMInit VMInit;
                              /*   51 : VM Death Event */
    jvmtiEventVMDeath VMDeath;
    ....
                              /*   54 : Class File Load Hook */
    jvmtiEventClassFileLoadHook ClassFileLoadHook;
                              /*   55 : Class Load */
    jvmtiEventClassLoad ClassLoad;
                              /*   56 : Class Prepare */
    jvmtiEventClassPrepare ClassPrepare;
                              /*   57 : VM Start Event */
    jvmtiEventVMStart VMStart;
    
} jvmtiEventCallbacks;

注册完VMInit事件注册回调函数eventHandlerVMInit就可以等待调度执行了

/*
 *  JVMTI callback support
 *
 *  We have two "stages" of callback support.
 *  At OnLoad time, we install a VMInit handler.
 *  When the VMInit handler runs, we remove the VMInit handler and install a
 *  ClassFileLoadHook handler.
 */

void JNICALL
eventHandlerVMInit( jvmtiEnv *      jvmtienv,
                    JNIEnv *        jnienv,
                    jthread         thread) {
    JPLISEnvironment * environment  = NULL;
    jboolean           success      = JNI_FALSE;

    environment = getJPLISEnvironment(jvmtienv);

    /* process the premain calls on the all the JPL agents */
    if ( environment != NULL ) {
        jthrowable outstandingException = preserveThrowable(jnienv);
        // java进程启动 in JPLISAgent.c
        success = processJavaStart( environment->mAgent,
                                    jnienv);
        restoreThrowable(jnienv, outstandingException);
    }

    /* if we fail to start cleanly, bring down the JVM */
    if ( !success ) {
        abortJVM(jnienv, JPLIS_ERRORMESSAGE_CANNOTSTART);
    }
}

执行函数体方法processJavaStart()：

创建实现类完善JPLISAgent剩余的field，之前初始化有些字段为null
开启ClassFileLoadHook事件，注册回调函数eventHandlerClassFileLoadHook
执行sun.instrument.InstrumentationImpl.loadClassAndCallPremain()
等待ClassFileLoadHook事件回调执行

/*
 * If this call fails, the JVM launch will ultimately be aborted,
 * so we don't have to be super-careful to clean up in partial failure
 * cases.
 */
jboolean
processJavaStart(   JPLISAgent *    agent,
                    JNIEnv *        jnienv) {
    jboolean    result;
    result = initializeFallbackError(jnienv);
    jplis_assert(result);

    /*
     *  Now make the InstrumentationImpl instance.
     */
    if ( result ) {
        // 创建实现类 完善JPLISAgent 剩余的field
        result = createInstrumentationImpl(jnienv, agent);
        jplis_assert(result);
    }

    /*
     *  Then turn off the VMInit handler and turn on the ClassFileLoadHook.
     *  This way it is on before anyone registers a transformer.
     */
    if ( result ) {
        // 开启ClassFileLoadHook事件,注册回调eventHandlerClassFileLoadHook
        result = setLivePhaseEventHandlers(agent);
        jplis_assert(result);
    }

    /*
     *  Load the Java agent, and call the premain.
     */
    if ( result ) {
        result = startJavaAgent(agent, jnienv,
                                agent->mAgentClassName, agent->mOptionsString,
                                agent->mPremainCaller);//loadClassAndCallPremain
    }

    /*
     * Finally surrender all of the tracking data that we don't need any more.
     * If something is wrong, skip it, we will be aborting the JVM anyway.
     */
    if ( result ) {
        deallocateCommandLineData(agent);
    }

    return result;
}

jvm事件调度开始执行eventHandlerClassFileLoadHook：

// 在字节码加载之前拦截 委托给初始化的sun/instrument/InstrumentationImpl.transform()
void JNICALL
eventHandlerClassFileLoadHook(  jvmtiEnv *              jvmtienv,
                                JNIEnv *                jnienv,
                                jclass                  class_being_redefined,
                                jobject                 loader,
                                const char*             name,
                                jobject                 protectionDomain,
                                jint                    class_data_len,
                                const unsigned char*    class_data,
                                jint*                   new_class_data_len,
                                unsigned char**         new_class_data) {
    JPLISEnvironment * environment  = NULL;

    environment = getJPLISEnvironment(jvmtienv);

    if ( environment != NULL ) {
        jthrowable outstandingException = preserveThrowable(jnienv);
        // 委托给java的transform() in JPLISAgent.c
        transformClassFile( environment->mAgent,
                            jnienv,
                            loader,
                            name,
                            class_being_redefined,
                            protectionDomain,
                            class_data_len,
                            class_data,
                            new_class_data_len,
                            new_class_data,
                            environment->mIsRetransformer);
        restoreThrowable(jnienv, outstandingException);
    }
}

委托java的transform()执行命令

void
transformClassFile(             JPLISAgent *            agent,
                                JNIEnv *                jnienv,
                                jobject                 loaderObject,
                                const char*             name,
                                jclass                  classBeingRedefined,
                                jobject                 protectionDomain,
                                jint                    class_data_len,
                                const unsigned char*    class_data,
                                jint*                   new_class_data_len,
                                unsigned char**         new_class_data,
                                jboolean                is_retransformer) {
            // 开始调用transform()
            transformedBufferObject = (*jnienv)->CallObjectMethod(
                                                jnienv,
                                                agent->mInstrumentationImpl,
                                                agent->mTransform,// transform
                                                loaderObject,
                                                classNameStringObject,
                                                classBeingRedefined,
                                                protectionDomain,
                                                classFileBufferObject,
                                                is_retransformer);
            errorOutstanding = checkForAndClearThrowable(jnienv);
            jplis_assert_msg(!errorOutstanding, "transform method call failed");
    ...
}

看到这里，上面的问题就有答案了
“spring相关类哪去了，Agent机制到底生效没有？”

我们想要增强的相关类在InstrumentationImpl.transform()里面由jvm类加载事件钩子传递进来，下面来证实一下。

Java Agent Transform

最上面的例子把注释打开，调用方法追加到TransformerManager中，所以上面的c代码必然有从管理器中获取transformer的逻辑，这里就不细说了。

package com.example.agent;
import java.lang.instrument.Instrumentation;
public class TinyAgent {
    public static void premain(String agentArgs, Instrumentation instrument) {
        Class[] allLoadedClasses = instrument.getAllLoadedClasses();
        System.out.println("loadedClasses length:" + allLoadedClasses.length);
        // 添加到TransformerManager
        instrumentation.addTransformer(new TinyTransformer(agentArgs));
        System.out.println("TinyAgent premain method");
    }
}

TinyTransformer这个类比较复杂打算后面单独写一篇文章细说，这里只是做个简单的实现：

package com.example.agent;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.net.URL;
import java.security.ProtectionDomain;

public class TinyTransformer implements ClassFileTransformer {
    private final String args;

    public TinyTransformer(String args) {
        this.args = args;
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        if (className.startsWith("com/example/Hello")) {
            String hello = "com/example/Hello.class";
            String agent = "E:/workspace/MiniApm/agent/target/classes/com/example/Hello.class";
            URL resource = loader.getResource(hello);
            System.out.println("resource in apm   :" + resource);
            System.out.println("resource in agent :" + agent);

            try {
                File file = new File(agent);
                InputStream in = new FileInputStream(file);
                byte[] bytes = new byte[in.available()];
                in.read(bytes);
                in.close();
                return bytes;
            } catch (IOException e) {
                e.printStackTrace();
            }
        }

        return classfileBuffer;
    }
}

这个实现逻辑是简单的把com/example/Hello.class重写成新的代理类实现，之后返回新的class数组，实现代码的增强。

输出结果如下：

before say()
-------------agent hello
after say()

原始结果：

-------------hello

小结：

因为实现的比较简单，就是class文件的简单替换。真实项目中这种替换是无法实施的，市面上有许多开源的字节码增强工具，比如bytebuddy，具体下一篇文章准备写一些实战。

增强方法timer
增强Caller/Runner
增强某些开源软件
插件化命令行Options