掘金后端 ( ) • 2024-04-25 18:16

在上篇文章中由于字数限制，没有将本文内容加进去，所以这里单独写一篇关于Java Agnet实际应用的文章。

如果你看完了上篇文章：关于Java Agent的使用、工作原理、及hotspot源码解析，那么此篇的应用文章就相当轻松了当然你需要使用过 mybatis plus 这个框架（因为我们本文是对这个框架的代码进行插桩）不过我想干Java的这个（mybatis plus）应该是基本功，对这个框架就不多介绍了。

1、我们的目的：（写时加密，读时解密）

因为我这里使用baomidou(mybatis plus)开发的，所以直接找到这个方法（注意我可不是凭想象来的而是经过了对增删改查这几个方法的debug 最终发现baomidou的mybatis plus最终都会走com.baomidou.mybatisplus.core.override类的execute方法，如下：）所以说需要插桩的地方我们就找到了。接下来开干！

2、编写Java Agent

2.1、编写premain方法并实现ClassFileTransformer的transform类，以便类加载时回调到transform 从而对指定的类中的方法进行增强即插桩！


package com.xzll.agent.config;

import com.xzll.agent.config.advice.MysqlFieldEncryptAndDecryptAdvice;
import javassist.*;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;

/**
 * @Author: hzz
 * @Date: 2023/3/3 09:15:21
 * @Description: MYSQL加解密 agent
 */
public class MysqlFieldCryptByExecuteBodyAgent {

   /**
    * 实现了带Instrumentation参数的premain()方法。
    */
   public static void premain(String args, Instrumentation inst) throws Exception {
      //调用addTransformer()方法对启动时所有的类(应用层)进行拦截
      inst.addTransformer(new MysqlReadWriteTransformer(), true);
   }

   static class MysqlReadWriteTransformer implements ClassFileTransformer {
      /**
       * 如果事先知道哪些类需要修改，最简单的修改类方式如下：
       * <p>
       * 1、通过调用ClassPool.get()方法获取一个CtClass对象
       * 2、修改它
       * 3、调用CtClass对象的writeFile()或toBytecode()方法获取修改后的类文件
       **/
      @Override
      public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
      //拦截指定要插桩 & 增强的类
         if ("com/baomidou/mybatisplus/core/override/MybatisMapperMethod".equals(className)) {
            CtClass clazz = null;
            System.out.println("对MybatisMapperMethod执行插桩实现读解密，写加密。");
            try {
               // 从ClassPool获得CtClass对象 (ClassPool对象是CtClass对象的容器，CtClass对象是类文件的抽象表示)
               final ClassPool classPool = ClassPool.getDefault();
               //这一步必不可少 和类加载器有关系，且maven中要配置 addClasspath=true
               //不加的话插桩时找不到MysqlFieldEncryptAndDecryptAdvice这个类
               classPool.insertClassPath(new ClassClassPath(MysqlFieldEncryptAndDecryptAdvice.class));
               clazz = classPool.get("com.baomidou.mybatisplus.core.override.MybatisMapperMethod");
               CtMethod getTime = clazz.getDeclaredMethod("execute");
               String body = "{\n" +
                              "return com.xzll.agent.config.advice.MysqlFieldEncryptAndDecryptAdvice.executeAgent($0,$1,$2);\n" +
                           "}\n";
               getTime.setBody(body);
               //通过CtClass的toBytecode(); 方法来获取 被修改后的字节码
               return clazz.toBytecode();
            } catch (Exception ex) {
               ex.printStackTrace();
            } finally {
               if (null != clazz) {
                  //调用CtClass对象的detach()方法后，对应class的其他方法将不能被调用。但是，你能够通过ClassPool的get()方法，
                  //重新创建一个代表对应类的CtClass对象。如果调用ClassPool的get()方法， ClassPool将重新读取一个类文件，并且重新创建一个CtClass对象，并通过get()方法返回
                  //如下所说：
                  //detach的意思是将内存中曾经被javassist加载过的Date对象移除，如果下次有需要在内存中找不到会重新走javassist加载
                  clazz.detach();
               }
               System.out.println("对sqlSession插桩完成");
            }
         }
         return classfileBuffer;
      }
   }
}

2.2、对指定方法插桩与增强

package com.xzll.agent.config.advice;


import cn.hutool.core.annotation.AnnotationUtil;
import cn.hutool.core.util.ReflectUtil;
import com.baomidou.mybatisplus.core.metadata.IPage;
import com.baomidou.mybatisplus.core.override.MybatisMapperMethod;
import com.xzll.agent.config.util.AESUtils;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.ibatis.binding.BindingException;
import org.apache.ibatis.binding.MapperMethod;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.session.SqlSession;

import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.Optional;

/**
 * @Author: 黄壮壮
 * @Date: 2024/4/20 10:37:11
 * @Description: 
 * 
 * 对baomidou 的 com.baomidou.mybatisplus.core.override.MybatisMapperMethod 类中的 execute方法进行增强，
 * 写时拦截到需要加密的字段（DO类上带有注解 @SensitiveData且字段上带有注解@EncryptTransaction的实体类中的字段 ） 进行aes加密，读时拦截需要解密（被@EncryptTransaction修饰）的字段用aes工具解密
 */
public class MysqlFieldEncryptAndDecryptAdvice {


   /**
    * 解密字段（如果类和其中的字段 存在被敏感注解修饰的话）
    *
    * @param args
    */
   private static <T> T decryptRead(T resultObject) {
      try {
         if (Objects.nonNull(resultObject)) {
            if (resultObject instanceof ArrayList) {
               List resultList = (List) resultObject;
               if (!CollectionUtils.isEmpty(resultList) && existSensitiveData(resultList.get(0))) {
                  for (Object result : resultList) {
                     decrypt(result);
                  }
               }
            } else {
               if (existSensitiveData(resultObject)) {
                  decrypt(resultObject);
               }
            }
         }
      } catch (Exception exception) {
         exception.printStackTrace();
      }
      return resultObject;
   }


   public static <T> T decrypt(T result) {
      //取出resultType的类
      Class<?> resultClass = result.getClass();
      Field[] declaredFields = resultClass.getDeclaredFields();
      for (Field field : declaredFields) {
         //取出所有被DecryptTransaction注解的字段 将其解密
         if (Objects.nonNull(field.getAnnotation(DecryptTransaction.class))) {
            field.setAccessible(true);
            try {
               Object object = field.get(result);
               String value = (String) object;
               if (StringUtils.isNotBlank(value)) {
                  //对注解的字段进行逐一解密
                  try {
                     value = AESUtils.decrypt(value);
                  } catch (Exception e) {
                  }
                  field.set(result, value);
               }
            } catch (Exception e) {
            }
         }
      }
      return result;
   }


   public static <T> T encrypt(T result) {
      //取出resultType的类
      Class<?> resultClass = result.getClass();
      Field[] declaredFields = resultClass.getDeclaredFields();
      for (Field field : declaredFields) {
         //取出所有被DecryptTransaction注解的字段 将其解密
         if (Objects.nonNull(field.getAnnotation(EncryptTransaction.class))) {
            field.setAccessible(true);
            try {
               Object object = field.get(result);
               String value = (String) object;
               if (StringUtils.isNotBlank(value)) {
                  //对注解的字段进行逐一解密
                  try {
                     value = AESUtils.encrypt(value);
                  } catch (Exception e) {
                  }
                  field.set(result, value);
               }
            } catch (Exception e) {
            }
         }
      }
      return result;
   }


   /**
    * 加密写字段（如果类和其中的字段 存在被敏感注解修饰的话）
    *
    * @param args
    */
   public static void encryptWrite(Object[] args) {
      try {
         for (Object object : args) {
            if (object instanceof List) {
               List resultList = (List) object;
               if (!CollectionUtils.isEmpty(resultList) && existSensitiveData(resultList.get(0))) {
                  for (Object result : resultList) {
                     encrypt(result);
                  }
               }
            } else {
               if (existSensitiveData(object)) {
                  encrypt(object);
               }
            }
         }
      } catch (Exception exception) {
         exception.printStackTrace();
      }
   }

   /**
    * 将原 MybatisMapperMethod 类的 execute方法进行增强
    *
    * @param mapperMethod
    * @param sqlSession
    * @param args
    * @return
    */
   public static Object executeAgent(MybatisMapperMethod mapperMethod, SqlSession sqlSession, Object[] args) {
      Object result;
      Object param;
      /**
       * 由于源代码中，都是直接使用this来访问 成员变量/方法，这种方式在 MybatisMapperMethod 类肯定是可行的，
       * 但是不在此类时通过this无法访问，所以就有了下边的：使用反射 获取私有成员变量/执行成员方法 ）
       */

      //（使用反射获取 MybatisMapperMethod 类中的私有成员变量）
      MapperMethod.SqlCommand commandValue = (MapperMethod.SqlCommand) ReflectUtil.getFieldValue(mapperMethod, "command");
      MapperMethod.MethodSignature methodValue = (MapperMethod.MethodSignature) ReflectUtil.getFieldValue(mapperMethod, "method");
      switch (commandValue.getType()) {
         case INSERT:
            //原代码
            //param = this.method.convertArgsToSqlCommandParam(args);
            //result = this.rowCountResult(sqlSession.insert(commandValue.getName(), param));

            //改后代码
            MysqlFieldEncryptAndDecryptAdvice.encryptWrite(args);
            param = methodValue.convertArgsToSqlCommandParam(args);
            result = ReflectUtil.invoke(mapperMethod, "rowCountResult", sqlSession.insert(commandValue.getName(), param));
            break;
         case UPDATE:
            //param = this.method.convertArgsToSqlCommandParam(args);
            //result = this.rowCountResult(sqlSession.update(this.command.getName(), param));
            MysqlFieldEncryptAndDecryptAdvice.encryptWrite(args);
            param = methodValue.convertArgsToSqlCommandParam(args);
            result = ReflectUtil.invoke(mapperMethod, "rowCountResult", sqlSession.update(commandValue.getName(), param));
            break;
         case DELETE:
            //param = this.method.convertArgsToSqlCommandParam(args);
            //result = this.rowCountResult(sqlSession.delete(this.command.getName(), param));

            param = methodValue.convertArgsToSqlCommandParam(args);
            result = ReflectUtil.invoke(mapperMethod, "rowCountResult", sqlSession.delete(commandValue.getName(), param));
            break;
         case SELECT:
            if (methodValue.returnsVoid() && methodValue.hasResultHandler()) {
               //this.executeWithResultHandler(sqlSession, args);
               ReflectUtil.invoke(mapperMethod, "executeWithResultHandler", sqlSession, args);
               result = null;
            } else if (methodValue.returnsMany()) {
               //result = this.executeForMany(sqlSession, args);
               result = ReflectUtil.invoke(mapperMethod, "executeForMany", sqlSession, args);
            } else if (methodValue.returnsMap()) {
               //result = this.executeForMap(sqlSession, args);
               result = ReflectUtil.invoke(mapperMethod, "executeForMap", sqlSession, args);

            } else if (methodValue.returnsCursor()) {
               //result = this.executeForCursor(sqlSession, args);
               result = ReflectUtil.invoke(mapperMethod, "executeForCursor", sqlSession, args);
            } else if (IPage.class.isAssignableFrom(methodValue.getReturnType())) {
               //result = this.executeForIPage(sqlSession, args);
               result = ReflectUtil.invoke(mapperMethod, "executeForIPage", sqlSession, args);
            } else {
               param = methodValue.convertArgsToSqlCommandParam(args);
               result = sqlSession.selectOne(commandValue.getName(), param);
               if (methodValue.returnsOptional() && (result == null || !methodValue.getReturnType().equals(result.getClass()))) {
                  result = Optional.ofNullable(result);
               }
            }
            break;
         case FLUSH:
            result = sqlSession.flushStatements();
            break;
         default:
            throw new BindingException("Unknown execution method for: " + commandValue.getName());
      }

      if (result == null && methodValue.getReturnType().isPrimitive() && !methodValue.returnsVoid()) {
         throw new BindingException("Mapper method '" + commandValue.getName() + " attempted to return null from a method with a primitive return type (" + methodValue.getReturnType() + ").");
      } else {
         if (Objects.equals(commandValue.getType(), SqlCommandType.SELECT) && result != null) {
            result = MysqlFieldEncryptAndDecryptAdvice.decryptRead(result);
         }
         return result;
      }
   }


   /**
    * 是否存在敏感字段 true存在 false不存在
    *
    * @param object
    * @return
    */
   private static boolean existSensitiveData(Object object) {
      Class<?> objectClass = object.getClass();
      SensitiveData sensitiveData = AnnotationUtil.getAnnotation(objectClass, SensitiveData.class);
      return Objects.nonNull(sensitiveData);
   }
}

2.3、指定哪些字段需要加/解密（编写DO类，指定加/解密的字段）

以下是DO实体定义： 一个字段想被修饰首先此字段所在的类需要被 @SensitiveData注解修饰，然后再加上 @EncryptTransaction（加密）或 @DecryptTransaction（解密）注解，@SensitiveData存在的意义是防止无效的遍历，提前将不带此注解的排除，提高性能。