Implement a MyBatis encryption and decryption plugin from scratch

Author: vivo Internet Server Team - Li Gang

This article introduces the process of using the MyBatis plugin to implement database field encryption and decryption.

1. Background of demand

For security and compliance reasons, companies need to encrypt some fields stored in plaintext in the database to prevent unauthorized access and personal information leakage.

Since the project has stopped iterating and the cost of transformation is too high, we chose the MyBatis plug-in to implement database encryption and decryption, ensuring that the specified fields can be encrypted when writing data to the database, and the specified fields can be decrypted when reading data.

2. Analysis of ideas

2.1 System Architecture

1. Add a ciphertext field to each field that needs to be encrypted (intrusive to the business), modify the database, mapper.xml and DO objects, and close the encryption and decryption of the plaintext/ciphertext field through a plug-in.
2. The custom Executor encrypts the plaintext fields of operations such as SELECT/UPDATE/INSERT/DELETE and sets them to the ciphertext fields.
3. The custom plug-in ResultSetHandler is responsible for decrypting the query result, decrypting the ciphertext field of SELECT and other operations and setting it to the plaintext field.

2.2 System flow

1. A new decryption process control switch is added to control whether to write only the original field/double write/write only the encrypted field when writing, and whether to read the original field or the encrypted field when reading.
2. A new historical data encryption task is added, which encrypts historical data in batches and writes it to the encrypted field.
3. For security reasons, there will be some verification/compensation tasks in the process, which will not be repeated here.

3. Program formulation

3.1 Introduction to MyBatis plugin

MyBatis reserves the org.apache.ibatis.plugin.Interceptor interface. By implementing this interface, we can intercept the execution process of MyBatis. The definition of the interface is as follows:

 public interface Interceptor {
  Object intercept(Invocation invocation) throws Throwable;
  Object plugin(Object target);
  void setProperties(Properties properties);
}

There are three methods:

• [intercept] : The specific process of plug-in execution, the incoming Invocation is the encapsulation of the proxy method by MyBatis.
• [plugin] : Use the current Interceptor to create a proxy. The usual implementation is Plugin.wrap(target, this), and jdk is used to create a dynamic proxy object in the wrap method.
• [setProperties] : Refer to the code below, you can set the parameters when configuring the plugin in the MyBatis configuration file, and call the Properties.getProperty("param1") method in the setProperties function to get the configured value.

 <plugins>
    <plugin interceptor="com.xx.xx.xxxInterceptor">
        <property name="param1" value="value1"/>
    </plugin>
</plugins>

Before implementing the intercept function to intercept the execution process of MyBatis, we need to use the @Intercepts annotation to specify the interception method.

 @Intercepts({ @Signature(type = Executor.class, method = "update", args = { MappedStatement.class, Object.class }),
        @Signature(type = Executor.class, method = "query", args = { MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class }) })

Referring to the above code, we can specify the classes and methods that need to be intercepted. Of course, we cannot intercept arbitrary objects. The classes that can be intercepted by MyBatis are the following four.

1. Executor
2. StatementHandler
3. ParameterHandler
4. ResultSetHandler

Returning to the requirements of database encryption, we need to select the classes that can be used to implement entry encryption and exit decryption from the above four classes. Before introducing these four classes, you need to have a certain understanding of the execution process of MyBatis.

3.2 Spring-MyBatis execution process

(1) Spring creates sqlSessionFactory through sqlSessionFactoryBean. When using sqlSessionFactoryBean, we usually specify configLocation and mapperLocations to tell sqlSessionFactoryBean where to read the configuration file and where to read the mapper file.

(2) After getting the location of the configuration file and mapper file, call XmlConfigBuilder.parse() and XmlMapperBuilder.parse() respectively to create Configuration and MappedStatement. As the name suggests, the Configuration class stores all the configurations of MyBatis, while the MappedStatement class stores the The encapsulation of a SQL statement, MappedStatement is stored in the Configuration object in the form of a map, and the key is the full path of the corresponding method.

(3) Spring scans all Mapper interfaces through ClassPathMapperScanner and creates BeanDefinition objects for them, but since they are essentially unimplemented interfaces, Spring will change the beanClass property of their BeanDefinition to MapperFactorybean.

(4) MapperFactoryBean also implements the FactoryBean interface. Spring will call the FactoryBean.getObject() method to obtain the bean when creating a bean. Finally, a proxy is created for the mapper interface through the newInstance method of mapperProxyFactory. The proxy is created by JDK, and the final proxy is generated. The object is MapperProxy.

(5) All interfaces that call mapper essentially call the MapperProxy.invoke method, and internally call various methods such as insert/update/delete of sqlSession.

 MapperMethod.java
public Object execute(SqlSession sqlSession, Object[] args) {
  Object result;
  if (SqlCommandType.INSERT == command.getType()) {
    Object param = method.convertArgsToSqlCommandParam(args);
    result = rowCountResult(sqlSession.insert(command.getName(), param));
  } else if (SqlCommandType.UPDATE == command.getType()) {
    Object param = method.convertArgsToSqlCommandParam(args);
    result = rowCountResult(sqlSession.update(command.getName(), param));
  } else if (SqlCommandType.DELETE == command.getType()) {
    Object param = method.convertArgsToSqlCommandParam(args);
    result = rowCountResult(sqlSession.delete(command.getName(), param));
  } else if (SqlCommandType.SELECT == command.getType()) {
    if (method.returnsVoid() && method.hasResultHandler()) {
      executeWithResultHandler(sqlSession, args);
      result = null;
    } else if (method.returnsMany()) {
      result = executeForMany(sqlSession, args);
    } else if (method.returnsMap()) {
      result = executeForMap(sqlSession, args);
    } else {
      Object param = method.convertArgsToSqlCommandParam(args);
      result = sqlSession.selectOne(command.getName(), param);
    }
  } else if (SqlCommandType.FLUSH == command.getType()) {
      result = sqlSession.flushStatements();
  } else {
    throw new BindingException("Unknown execution method for: " + command.getName());
  }
  if (result == null && method.getReturnType().isPrimitive() && !method.returnsVoid()) {
    throw new BindingException("Mapper method '" + command.getName()
        + " attempted to return null from a method with a primitive return type (" + method.getReturnType() + ").");
  }
  return result;
}

(6) SqlSession can be understood as a session, SqlSession will obtain the corresponding MappedStatement from Configuration and hand it over to Executor for execution.

 DefaultSqlSession.java
@Override
public <E> List<E> selectList(String statement, Object parameter, RowBounds rowBounds) {
  try {
    // 从configuration对象中使用被调用方法的全路径，获取对应的MappedStatement
    MappedStatement ms = configuration.getMappedStatement(statement);
    return executor.query(ms, wrapCollection(parameter), rowBounds, Executor.NO_RESULT_HANDLER);
  } catch (Exception e) {
    throw ExceptionFactory.wrapException("Error querying database.  Cause: " + e, e);
  } finally {
    ErrorContext.instance().reset();
  }
}

(7) Executor will first create StatementHandler, which can be understood as a statement execution.

(8) Then the Executor will obtain the connection. The specific method of obtaining the connection depends on the implementation of the Datasource. The connection can be obtained by means of connection pooling.

(9) Then call the StatementHandler.prepare method, which corresponds to the Connection.prepareStatement step in the JDBC execution process.

(10) The Executor then calls the parameterize method of the StatementHandler to set the parameters, corresponding to the StatementHandler.setXXX() of the JDBC execution process to set the parameters, and the ParameterHandler method will be created internally.

 SimpleExecutor.java
@Override
public <E> List<E> doQuery(MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) throws SQLException {
  Statement stmt = null;
  try {
    Configuration configuration = ms.getConfiguration();
    // 创建StatementHandler，对应第7步
    StatementHandler handler = configuration.newStatementHandler(wrapper, ms, parameter, rowBounds, resultHandler, boundSql);
    // 获取连接，再调用conncetion.prepareStatement创建prepareStatement，设置参数
    stmt = prepareStatement(handler, ms.getStatementLog());
    // 执行prepareStatement
    return handler.<E>query(stmt, resultHandler);
  } finally {
    closeStatement(stmt);
  }
}

(11) The returned result is processed by ResultSetHandler, the return value of JDBC is processed, and it is converted into a Java object.

3.3 When to create the MyBatis plugin

In the Configuration class, we can see four methods: newExecutor, newStatementHandler, newParameterHandler, and newResultSetHandler. The proxy class of the plug-in is created in these four methods. I take the creation of StatementHandler as an example:

 Configuration.java
public StatementHandler newStatementHandler(Executor executor, MappedStatement mappedStatement, Object parameterObject, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) {
  StatementHandler statementHandler = new RoutingStatementHandler(executor, mappedStatement, parameterObject, rowBounds, resultHandler, boundSql);
  // 使用责任链的形式创建代理
  statementHandler = (StatementHandler) interceptorChain.pluginAll(statementHandler);
  return statementHandler;
}
 
InterceptorChain.java
public Object pluginAll(Object target) {
  for (Interceptor interceptor : interceptors) {
    target = interceptor.plugin(target);
  }
  return target;
}

The interceptor.plugin corresponds to the method in the interceptor implemented by ourselves. The usual implementation is Plugin.wrap(target, this); . The method of creating a proxy internally is JDK.

3.4 MyBatis plugin can intercept class selection

Mybatis is essentially an encapsulation of the JDBC execution process. Combined with the above figure, we briefly summarize the functions of Mybatis, which can be delegated.

• [Executor] : The object that actually executes the SQL statement. When calling the method of sqlSession, it is essentially the method of calling the executor, and is also responsible for obtaining the connection and creating the StatementHandler.
• [StatementHandler] : Create and hold ParameterHandler and ResultSetHandler objects, operate JDBC statement and perform database operations.
• 【ParameterHandler】: Process the input parameters and set the parameters on the Java method to the executed statement.
• 【ResultSetHandler】 : Process the execution result of the SQL statement and convert the return value into a Java object.

For the encryption of the input parameters, we need to change the parameter value to the encrypted parameter before the ParameterHandler calls the prepareStatement.setXXX() method to set the parameter, so it seems that intercepting Executor/StatementHandler/ParameterHandler can be done.

But actually? Since we are not encrypting the original field, but adding an encrypted field, what problems will this bring? Please see the dynamic SQL with encrypted fields in the mapper.xml file below:

 <select id="selectUserList" resultMap="BaseResultMap" parameterType="com.xxx.internet.demo.entity.UserInfo">
        SELECT
        *
        FROM
        `t_user_info`
        <where>
            <if test="phone != null">
                `phone` = #{phone}
            </if>
<!--            明文字段-->
            <if test="secret != null">
                AND `secret` = #{secret}
            </if>
<!--            加密后字段-->
            <if test="secretCiper != null">
                AND `secret_ciper` = #{secretCiper}
            </if>
            <if test="name">
                AND `name` = #{name}
            </if>
        </where>
        ORDER BY `update_time` DESC
    </select>

It can be seen that this statement has a dynamic tag, so it must not be directly handed over to JDBC to create prepareStatement, it needs to be parsed into static SQL first, and this step is done before the Executor calls StatementHandler.parameterize(), which is executed by MappedStatementHandler. The getBoundSql(Object parameterObject) function parses the dynamic tag and generates a static SQL statement. The parameterObject here can be regarded as a Map for the time being, and the key values are the parameter name and parameter value.

Then let's take a look at the problems with using StatementHandler and ParameterHandler for parameter encryption. When MappedStatementHandler.getBoundSql is executed, the encrypted parameters are not written in parameterObject. When judging the label, it must be no, and the final generated static SQL must not be Contains encrypted fields. No matter how we deal with parameterObject in StatementHandler and ParameterHandler, the encryption of input parameters cannot be realized.

Therefore, in the encryption of input parameters, we can only choose to intercept the update and query methods of Executor.

What about decryption of the return value? Referring to the flowchart, we can intercept ResultSetHandler and Executor, which is true. In terms of processing return values, the two are equivalent. The return value of ResultSetHandler.handleResultSet() is directly passed to Executor, and then by Executor is transparently transmitted to SqlSession, so you can choose one of the two.

4. Implementation of the plan

After knowing the object to be intercepted, you can start to implement the encryption and decryption plug-in. First define a method dimension annotation.

 /**
 * 通过注解来表明，我们需要对那个字段进行加密
 */
@Target({ ElementType.METHOD })
@Retention(RetentionPolicy.RUNTIME)
@Inherited
@Documented
public @interface TEncrypt {
    /**
     * 加密时从srcKey到destKey
     * @return
     */
    String[] srcKey() default {};
 
    /**
     * 解密时从destKey到srcKey
     * @return
     */
    String[] destKey() default {};
}

Put this annotation on the DAO layer method that needs to be encrypted and decrypted.

 UserMapper.java
public interface UserMapper {
    @TEncrypt(srcKey = {"secret"}, destKey = {"secretCiper"})
    List<UserInfo> selectUserList(UserInfo userInfo);
    }

Modify the xxxMapper.xml file

 <mapper namespace="com.xxx.internet.demo.mapper.UserMapper">
    <resultMap id="BaseResultMap" type="com.xxx.internet.demo.entity.UserInfo">
        <id column="id" jdbcType="BIGINT" property="id" />
        <id column="phone" jdbcType="VARCHAR" property="phone"/>
        <id column="secret" jdbcType="VARCHAR" property="secret"/>
<!--        加密后映射-->
        <id column="secret_ciper" jdbcType="VARCHAR" property="secretCiper"/>
        <id column="name" jdbcType="VARCHAR" property="name" />
    </resultMap>
 
    <select id="selectUserList" resultMap="BaseResultMap" parameterType="com.xxx.internet.demo.entity.UserInfo">
        SELECT
        *
        FROM
        `t_user_info`
        <where>
            <if test="phone != null">
                `phone` = #{phone}
            </if>
<!--            明文字段-->
            <if test="secret != null">
                AND `secret` = #{secret}
            </if>
<!--            加密后字段-->
            <if test="secretCiper != null">
                AND `secret_ciper` = #{secretCiper}
            </if>
            <if test="name">
                AND `name` = #{name}
            </if>
        </where>
        ORDER BY `update_time` DESCv
    </select>
</mapper>

After making the above modifications, we can write the encryption plug-in

 @Intercepts({ @Signature(type = Executor.class, method = "update", args = { MappedStatement.class, Object.class }),
        @Signature(type = Executor.class, method = "query", args = { MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class }) })
public class ExecutorEncryptInterceptor implements Interceptor {
    private static final ObjectFactory        DEFAULT_OBJECT_FACTORY         = new DefaultObjectFactory();
 
    private static final ObjectWrapperFactory DEFAULT_OBJECT_WRAPPER_FACTORY = new DefaultObjectWrapperFactory();
 
    private static final ReflectorFactory     REFLECTOR_FACTORY              = new DefaultReflectorFactory();
 
    private static final List<String>         COLLECTION_NAME  = Arrays.asList("list");
 
    private static final String               COUNT_SUFFIX                   = "_COUNT";
 
    @Override
    public Object intercept(Invocation invocation) throws Throwable {
 
        // 获取拦截器拦截的设置参数对象DefaultParameterHandler
        final Object[] args = invocation.getArgs();
        MappedStatement mappedStatement = (MappedStatement) args[0];
        Object parameterObject = args[1];
 
        // id字段对应执行的SQL的方法的全路径，包含类名和方法名
        String id = mappedStatement.getId();
        String className = id.substring(0, id.lastIndexOf("."));
        String methodName = id.substring(id.lastIndexOf(".") + 1);
 
        // 分页插件会生成一个count语句，这个语句的参数也要做处理
        if (methodName.endsWith(COUNT_SUFFIX)) {
            methodName = methodName.substring(0, methodName.lastIndexOf(COUNT_SUFFIX));
        }
 
        // 动态加载类并获取类中的方法
        final Method[] methods = Class.forName(className).getMethods();
 
        // 遍历类的所有方法并找到此次调用的方法
        for (Method method : methods) {
            if (method.getName().equalsIgnoreCase(methodName) && method.isAnnotationPresent(TEncrypt.class)) {
 
                // 获取方法上的注解以及注解对应的参数
                TEncrypt paramAnnotation = method.getAnnotation(TEncrypt.class);
 
                // 支持加密的操作，这里只修改参数
                if (parameterObject instanceof Map) {
                    List<String> paramAnnotations = findParams(method);
                    parameterMapHandler((Map) parameterObject, paramAnnotation, mappedStatement.getSqlCommandType(), paramAnnotations);
                } else {
                    encryptParam(parameterObject, paramAnnotation, mappedStatement.getSqlCommandType());
                }
            }
        }
 
        return invocation.proceed();
    }
}

The main process of encryption is as follows:

1. Determine whether the method called this time is annotated with @TEncrypt.
2. Get the annotation and the parameters configured on the annotation.
3. Traverse the parameterObject to find the fields that need to be encrypted.
4. Call the encryption method to get the encrypted value.
5. Write encrypted fields and values to parameterObject.

The difficulty is mainly in the parsing of parameterObject. At the Executor level, parameterObject is no longer a simple Object[], but an object created by the MapperMethod.convertArgsToSqlCommandParam(Object[] args) method. Since this object needs to be processed, We must first know the process of its creation.

Referring to the creation process of the parameterObject in the figure above, the processing of the parameterObject by the encryption plug-in is essentially a reverse process. If it is a list, we traverse each value in the list, and if it is a map, we traverse each value in the map.

After obtaining the Object to be processed, traverse each property in the Object to determine whether it is in the srcKeys parameter of the @TEncrypt annotation. If so, encrypt it and set it to the Object.

The logic of the decryption plug-in is basically the same as that of the encryption plug-in, and will not be repeated here.

V. Challenges

5.1 The paging plugin automatically generates the count statement

Many places in the business code use com.github.pagehelper for physical paging. Refer to the demo below. When using PageRowBounds, the pagehelper plugin will help us get the total number of eligible data and set it to the total property of the rowBounds object.

 PageRowBounds rowBounds = new PageRowBounds(0, 10);
List<User> list = userMapper.selectIf(1, rowBounds);
long total = rowBounds.getTotal();

Then the problem comes. On the surface, we only executed the statement userMapper.selectIf(1, rowBounds), and pagehelper is a physical paging realized by rewriting SQL to increase limit and offset. During the execution of the entire statement, there is no data from the database. Read out all the data that meets the conditions, so how does pagehelper get the total number of data?

The answer is that pagehelper will execute one more count statement. Let's not talk about the principle of executing an additional count statement, let's first see what problems can be caused by adding a count statement.

Referring to the previous selectUserList interface, assuming that we want to select data whose secret is a certain value, the final execution after processing by the encryption plugin is roughly the following statement "select * from t\_user\_info where secret\_ciper = ? order by update\_time limit ?, ?".

However, since pagehelper will execute another statement, and since the statement does not have the @TEncrypt annotation, it will not be intercepted by the encryption plugin. The final executed count statement is similar to this: "select count(*) from t\_user \_info where secret = ? order by update_time".

It can be clearly seen that the first statement uses secret_ciper as the query condition, and the count statement uses secret as the query condition, which will cause the final amount of data to be inconsistent with the actual amount of data.

Therefore, we have specially treated the count statement in the code of the encryption plug-in. Since the id of the mappedStatement corresponding to the count statement added by the pagehelper is fixed at the end of "\_COUNT", and this id is the full path of the method in the corresponding mapper, For example, the id of the original statement is "com.xxx.internet.demo.entity.UserInfo.selectUserList", then the id of the count statement is "com.xxx.internet.demo.entity.UserInfo.selectUserList\_COUNT", remove " _COUNT", we can then judge whether there are annotations on the corresponding method.

6. Summary

This article introduces the exploration process of using the MyBatis plug-in to implement database field encryption and decryption. There are many details that need to be paid attention to in the actual development process, and my understanding of MyBatis has deepened throughout the process. In general, this solution is relatively lightweight, and although it is intrusive to business code, it can control the impact to a minimum.