FiAttach传入两个参数,一个是agent.jar的路径,一个是存放希望运行时进行替换的类文件的文件夹路径。

程序自动检测当前的Java应用,将agent.jar附着到虚拟机进程,并将文件夹下的类文件动态替换进去(用新的类替换虚拟机中原来加载的类)。

import java.io.IOException;
import java.util.List;
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 com.sun.tools.attach.VirtualMachineDescriptor;

public class FiAttach {
  public static void main(String[] args) {
    List<VirtualMachineDescriptor> vmdList = VirtualMachine.list();
    if (args.length < 2) {
      System.out.println("Error! Run Command: java com.taobao.fi.FiAttach agentJarPath agentArgs");
      return;
    }
    
    String agentJarPath = args[0];
    String agentArgs = args[1];
    
    System.out.println("agentJarPath: " + agentJarPath);
    System.out.println("agentArgs: " + agentArgs);
    
    for (VirtualMachineDescriptor vmd : vmdList) {
      // 注意,目前只支持jboss和tomcat,否则判断会失效!
      // vmd.displayName(): org.jboss.Main -b 0.0.0.0 -Djboss.server.home.dir=/home/admin/deploy/.default -Djboss.server.home.url=file:/home/admin/deploy/.default
      if (vmd.displayName().startsWith("org.jboss.Main") || vmd.displayName().startsWith("org.apache.catalina.startup.Bootstrap")) {
        try {
          VirtualMachine vm = VirtualMachine.attach(vmd);
          vm.loadAgent(agentJarPath, agentArgs);
          vm.detach();
        } catch (AttachNotSupportedException e) {
          e.printStackTrace();
        } catch (IOException e) {
          e.printStackTrace();
        } catch (AgentLoadException e) {
          e.printStackTrace();
        } catch (AgentInitializationException e) {
          e.printStackTrace();
        }
      }
    }
  }
}
程序编译时,需要依赖JDK_HOME/lib/tools.jar

下面看agent.jar的实现,AgentMain.java:
import java.util.Set;
import java.util.HashSet;
import java.io.File;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;

public class AgentMain {

  private static Set<String> fiClsFileNames = new HashSet<String>();

  private static Transformer transformer = new Transformer();

  // 标识是否之前做过故障注入
  private static boolean hasFi = false;

  private static void updateClsFileNames(String fiClassFolderPath) {
    fiClsFileNames.clear();
    File fiClassFolderFile = new File(fiClassFolderPath);
    if (!fiClassFolderFile.isDirectory()) {
      return;
    }

    File[] fiClassFiles = fiClassFolderFile.listFiles();
    for (File fiClassFile : fiClassFiles) {
      fiClsFileNames.add(fiClassFile.getName());
    }
  }

  // 判断是否是已经进行过故障注入的类 或者是 将要进行故障注入的类
  private static boolean isPrevFiCls(String clsName) {
    String clsFileName = clsName + ".class";
    return fiClsFileNames.contains(clsFileName);
  }

  // 判断是否是将要进行故障注入的类(注意:在这之前,需要调用updateCurrClsFileNames())
  private static boolean isWillingFiCls(String clsName) {
    String clsFileName = clsName + ".class";
    return fiClsFileNames.contains(clsFileName);
  }

  public static void agentmain(String agentArgs, Instrumentation inst)
      throws ClassNotFoundException, UnmodifiableClassException,
      InterruptedException {

    System.out.println("AgentMain::agentmain!!");

    synchronized (AgentMain.class) {

      String fiClsFolderPath = agentArgs;

      if (hasFi) {
        inst.removeTransformer(transformer);

        Class[] classes = inst.getAllLoadedClasses();
        for (Class cls : classes) {
          System.out.println("AgentMain::agentmain, recover class: "
              + cls.getName());
          if (isPrevFiCls(cls.getName())) {
            // 触发已加载的类 还原对类的更改
            inst.retransformClasses(cls);
          }
        }
      }

      updateClsFileNames(fiClsFolderPath);
      
      transformer.setFiClsFolderPath(fiClsFolderPath);
      // 这里应该不存在线程安全隐患,因为attach动作总是人为触发的
      transformer.setFiClsFileNames(fiClsFileNames);

      // 添加转换器
      inst.addTransformer(transformer, true);

      // 更改当前已加载的类
      Class[] classes = inst.getAllLoadedClasses();
      for (Class cls : classes) {
        if (isWillingFiCls(cls.getName())) {
          System.out
              .println("AgentMain::agentmain, transform class: "
                  + cls.getName());
          inst.retransformClasses(cls);
        }
      }

      hasFi = true;
    }
  }
}

Transformer.java:
import java.util.Set;
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.security.ProtectionDomain;

public class Transformer implements ClassFileTransformer {

  private String fiClsFolderPath;

  private Set<String> fiClsFileNames = null;

  public String getFiClsFolderPath() {
    return fiClsFolderPath;
  }

  public void setFiClsFolderPath(String fiClsFolderPath) {
    this.fiClsFolderPath = fiClsFolderPath;
  }

  public Set<String> getFiClsFileNames() {
    return fiClsFileNames;
  }

  public void setFiClsFileNames(Set<String> fiClsFileNames) {
    this.fiClsFileNames = fiClsFileNames;
  }

  private boolean isFiCls(String clsName) {
    String clsFileName = clsName + ".class";
    return fiClsFileNames.contains(clsFileName);
  }

  public static byte[] getBytesFromFile(String fileName) {
    System.out.println("[Transformer]: getBytesFromFile: " + fileName);
    try {
      // precondition
      File file = new File(fileName);
      InputStream is = new FileInputStream(file);
      long length = file.length();
      byte[] bytes = new byte[(int) length];

      // Read in the bytes
      int offset = 0;
      int numRead = 0;
      while (offset < bytes.length
          && (numRead = is.read(bytes, offset, bytes.length - offset)) >= 0) {
        offset += numRead;
      }

      if (offset < bytes.length) {
        throw new IOException("Could not completely read file "
            + file.getName());
      }
      is.close();
      return bytes;
    } catch (Exception e) {
      System.out.println("error occurs in _ClassTransformer!"
          + e.getClass().getName());
      return null;
    }
  }

  @Override
  public byte[] transform(ClassLoader loader, String className,
      Class<?> classBeingRedefined, ProtectionDomain protectionDomain,
      byte[] classfileBuffer) throws IllegalClassFormatException {
    System.out.println("transform: " + className);

    // 如果不是将要进行故障注入的类,直接返回null,意即不做任何的转换处理
    if (!isFiCls(className.replace("/", "."))) {
      return null;
    }

    return getBytesFromFile(fiClsFolderPath + File.separator
        + className.replace("/", ".") + ".class");
  }
}

类文件名的命名格式,举例:com.taobao.A.class这样。

这里,如果转换后的类(更改后的类)需要依赖某个类(记为类B),可以将这个类B的源码放置到agent工程,随着agent.jar打包进去。虚拟机在加载agent.jar后,也会将该类装载进去。
这样,转换后的类也可以访问到类B。

注意,为了打成agent,需要在源码目录下新建META-INF文件夹
文件夹内新建文件MANIFEST.MF,内容如下:
Manifest-Version: 1.0
Agent-Class: com.taobao.fi.AgentMain
Can-Redefine-Classes: false
Can-Retransform-Classes: false
Boot-Class-Path: fiagent.jar

特别注意,此文件是空格敏感的。每一行不容许有多余的空格。否则,打包出来的agent.jar,虚拟机会不认的。

利用eclipse的导出jar包时,记得要选择使用该工程源码目录下的MANIFEST.MF文件。

如果你想还原成原来的类,只需要将类文件夹下的类删除,然后,重新执行FiAttach即可。

本文是本人实作了Java故障注入测试工具后的总结,供业界同仁参考。题外话,像btrace也是基于此原理。

此文完。