Mybatis中@Intercepts、StatementHandler、ResultSetHandler用法记录，主要用于实现数据加解密

公司的海外业务中，有涉及数据隐私保护的要求。因此对于一些敏感数据，比如用户的姓名、地址、电话等，需要经过一层加密操作，才能把数据存到库中。因此，需要写一个方法实现在插入钱对数据进行加密。
方法如下：

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.xxxx.annotation.encrypt.EncryptEnabled;
import com.xxxx.annotation.encrypt.EncryptField;
import com.xxxx.annotation.encrypt.EncryptJSONField;
import com.xxxx.annotation.encrypt.EncryptJSONFieldKey;
import com.xxxx.exception.BusinessException;
import com.xxxx.response.LepusResultCode;
import com.xxxx.service.LocalBladeService;
import com.xxxx.util.JsonUtils;
import com.xxxx.util.MybatisPluginUtils;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.apache.ibatis.session.Configuration;

import java.lang.reflect.Field;
import java.sql.Connection;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Properties;

@Intercepts({
        @Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})
})
@Slf4j
@RequiredArgsConstructor
public class EncryptWriteInterceptor implements Interceptor {

    private final LocalBladeService localBladeService;
    private static final String MAPPED_STATEMENT = "delegate.mappedStatement";
    private static final String BOUND_SQL = "delegate.boundSql";

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        StatementHandler statementHandler = MybatisPluginUtils.realTarget(invocation.getTarget());
        MetaObject metaObject = SystemMetaObject.forObject(statementHandler);
        MappedStatement mappedStatement = (MappedStatement) metaObject.getValue(MAPPED_STATEMENT);
        SqlCommandType commandType = mappedStatement.getSqlCommandType();

        BoundSql boundSql = (BoundSql) metaObject.getValue(BOUND_SQL);
        Object params = boundSql.getParameterObject();
        if (params instanceof Map) {
            return invocation.proceed();
        }

        EncryptEnabled encryptEnabled = Objects.isNull(params)
                ? null : params.getClass().getAnnotation(EncryptEnabled.class);
        if (Objects.nonNull(encryptEnabled) && encryptEnabled.value()) {
            handleParameters(mappedStatement.getConfiguration(), boundSql, params, commandType);
        }

        return invocation.proceed();
    }

    private void handleParameters(Configuration config, BoundSql boundSql, Object param, SqlCommandType commandType) {
        if (!isWriteCommand(commandType)) {
            return;
        }

        Map<String, Object> newValues = Maps.newHashMap();
        List<Field> needEncryptFields = Lists.newArrayList();
        List<String> encryptFieldValues = Lists.newArrayList();
        MetaObject metaObject = config.newMetaObject(param);

        for (Field field : param.getClass().getDeclaredFields()) {
            Object value = metaObject.getValue(field.getName());
            Object newValue = value;
            if (!(value instanceof CharSequence)) {
                continue;
            }

            if (isEncryptField(field, value)) {
                needEncryptFields.add(field);
                encryptFieldValues.add((String) value);
            }

            if (isEncryptJsonField(field, value)) {
                newValue = handleEncryptJsonField(field, newValue);
                newValues.put(field.getName(), newValue);
            }

        }

        // 批量加密，批量赋值
        if (CollectionUtils.isNotEmpty(needEncryptFields)) {
            List<String> secrets = localBladeService.encryptData(encryptFieldValues);
            if (CollectionUtils.isEmpty(secrets) || secrets.size() != needEncryptFields.size()) {
                return;
            }
            for (int i = 0; i < needEncryptFields.size(); i++) {
                Field field = needEncryptFields.get(i);
                newValues.put(field.getName(), secrets.get(i));
            }
        }

        for (Map.Entry<String, Object> entry : newValues.entrySet()) {
            boundSql.setAdditionalParameter(entry.getKey(), entry.getValue());
        }

    }

    private boolean isEncryptField(Field field, Object val) {
        EncryptField encryptField = field.getAnnotation(EncryptField.class);
        return Objects.nonNull(encryptField) && Objects.nonNull(val);
    }

    private boolean isEncryptJsonField(Field field, Object val) {
        EncryptJSONField encryptJSONField = field.getAnnotation(EncryptJSONField.class);
        return Objects.nonNull(encryptJSONField) && Objects.nonNull(val);
    }

    private boolean isWriteCommand(SqlCommandType commandType) {
        return commandType == SqlCommandType.UPDATE || commandType == SqlCommandType.INSERT;
    }

    private Object handleEncryptJsonField(Field field, Object value) {
        EncryptJSONField encryptJSONField = field.getAnnotation(EncryptJSONField.class);
        Object newValue = value;
        if (encryptJSONField != null && value != null) {
            newValue = processJsonField(newValue, encryptJSONField);
        }
        return newValue;
    }

    private Object processJsonField(Object val, EncryptJSONField encryptJSONField) {
        try {
            Map<String, Object> map = JsonUtils.parseToMap(val.toString());
            EncryptJSONFieldKey[] keys = encryptJSONField.keys();
            List<String> needEncryptJsonKeys = Lists.newArrayList(), encryptJsonValues = Lists.newArrayList();
            for (EncryptJSONFieldKey jsonFieldKey : keys) {
                String key = jsonFieldKey.key();
                Object oldData = map.get(key);
                if (Objects.nonNull(oldData)) {
                    needEncryptJsonKeys.add(key);
                    encryptJsonValues.add(oldData.toString());
                }
            }

            // 批量加密，批量赋值
            if (CollectionUtils.isNotEmpty(needEncryptJsonKeys)) {
                List<String> secrets = localBladeService.encryptData(encryptJsonValues);
                if (CollectionUtils.isEmpty(secrets) || secrets.size() != needEncryptJsonKeys.size()) {
                    log.error("批量加密 json 字段失败， {}", JsonUtils.toJson(needEncryptJsonKeys));
                    throw new BusinessException(LepusResultCode.FAIL_BATCH_ENCRYPT);
                }
                for (int i = 0; i < needEncryptJsonKeys.size(); i++) {
                    String key = needEncryptJsonKeys.get(i);
                    map.put(key, secrets.get(i));
                }
            }

            return JsonUtils.toJson(map);
        } catch (Throwable t) {
            log.error("加密 json 字段失败， cause: {}", t.getMessage(), t);
            throw new BusinessException(LepusResultCode.FAIL_BATCH_ENCRYPT);
        }
    }

    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {
        // do nothing
    }
}

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;

import com.xxxx.annotation.encrypt.EncryptEnabled;
import com.xxxx.annotation.encrypt.EncryptField;
import com.xxxx.annotation.encrypt.EncryptJSONField;
import com.xxxx.annotation.encrypt.EncryptJSONFieldKey;
import com.xxxx.exception.BusinessException;
import com.xxxx.response.LepusResultCode;
import com.xxxx.service.LocalBladeService;
import com.xxxx.util.JsonUtils;
import com.xxxx.util.MDCUtils;
import com.xxxx.util.MybatisPluginUtils;

import org.apache.commons.collections4.MapUtils;
import org.apache.ibatis.executor.resultset.ResultSetHandler;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.ResultMap;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.springframework.util.CollectionUtils;

import java.lang.reflect.Field;
import java.sql.Statement;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Properties;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicBoolean;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;

import static com.xxxx.util.threadPool.ThreadPoolManager.ENCRYPTED_EXECUTOR;
@Intercepts({
        @Signature(type = ResultSetHandler.class, method = "handleResultSets", args = {Statement.class})
})
@Slf4j
@RequiredArgsConstructor
public class DecryptReadInterceptor implements Interceptor {

    private final LocalBladeService localBladeService;
    private static final String MAPPED_STATEMENT = "mappedStatement";

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        @SuppressWarnings("unchecked") final List<Object> results = (List<Object>) invocation.proceed();
        if (CollectionUtils.isEmpty(results)) {
            return results;
        }

        final ResultSetHandler statementHandler = MybatisPluginUtils.realTarget(invocation.getTarget());
        final MetaObject metaObject = SystemMetaObject.forObject(statementHandler);
        final MappedStatement mappedStatement = (MappedStatement) metaObject.getValue(MAPPED_STATEMENT);
        final ResultMap resultMap = mappedStatement.getResultMaps().isEmpty() ? null : mappedStatement.getResultMaps().get(0);

        Object result0 = results.get(0);
        // 确保 DTO 类有注解
        EncryptEnabled encryptEnabled = result0.getClass().getAnnotation(EncryptEnabled.class);
        if (encryptEnabled == null || !encryptEnabled.value()) {
            return results;
        }

        // 确保类中包含注解字段
        Map<String, EncryptField> encryptFieldMap = consEncryptedMapByResultMap(resultMap);
        Map<String, EncryptJSONField> encryptJSONFieldMap = consEncryptedJsonMapByResultMap(resultMap);
        if (encryptFieldMap.isEmpty() && encryptJSONFieldMap.isEmpty()) {
            return results;
        }

        // 依次解密字段
//        for (Object result : results) {
//            final MetaObject resultMetaObj = mappedStatement.getConfiguration().newMetaObject(result);
//            handleStringValues(resultMetaObj, Lists.newArrayList(encryptFieldMap.keySet()));
//
//            for (Map.Entry<String, EncryptJSONField> entry : encryptJSONFieldMap.entrySet()) {
//                handleJsonVal(resultMetaObj, entry.getKey(), entry.getValue());
//            }
//        }

        // 对于一个object中，对于EncryptField和EncryptJSONField使用多线程进行加解密
        final CountDownLatch downLatch = new CountDownLatch(2 * results.size());
        AtomicBoolean interrupt = new AtomicBoolean(false);
        for (Object result : results) {
            final MetaObject resultMetaObj = mappedStatement.getConfiguration().newMetaObject(result);
            ENCRYPTED_EXECUTOR.submit(() -> asynEncryptField(resultMetaObj, encryptFieldMap, downLatch, interrupt));
            ENCRYPTED_EXECUTOR.submit(() -> asynEncryptJSONField(resultMetaObj, encryptJSONFieldMap, downLatch, interrupt));
        }
        downLatch.await();
        if (interrupt.get()) {
            log.error("fail to batch decrypt");
            throw new BusinessException(LepusResultCode.FAIL_BATCH_ENCRYPT);
        }
        return results;
    }

    public void asynEncryptField(MetaObject resultMetaObj, Map<String, EncryptField> encryptFieldMap, CountDownLatch downLatch, AtomicBoolean interrupt) {
        String rid = MDCUtils.requestId();
        try {
            MDCUtils.initRequestId(rid);
            handleStringValues(resultMetaObj, Lists.newArrayList(encryptFieldMap.keySet()));
        } catch (RuntimeException e) {
            log.error("encryptField fail : {}", e.getMessage());
            interrupt.set(true);
        } finally {
            MDCUtils.clearRequestId();
            downLatch.countDown();
        }
    }

    public void asynEncryptJSONField(MetaObject resultMetaObj, Map<String, EncryptJSONField> encryptJSONFieldMap, CountDownLatch downLatch, AtomicBoolean interrupt) {
        String rid = MDCUtils.requestId();
        try {
            MDCUtils.initRequestId(rid);
            for (Map.Entry<String, EncryptJSONField> entry : encryptJSONFieldMap.entrySet()) {
                handleJsonVal(resultMetaObj, entry.getKey(), entry.getValue());
            }
        } catch (RuntimeException e) {
            log.error("encryptJSONField fail : {}", e.getMessage());
            interrupt.set(true);
        } finally {
            MDCUtils.clearRequestId();
            downLatch.countDown();
        }
    }

    private Map<String, EncryptField> consEncryptedMapByResultMap(ResultMap resultMap) {
        if (resultMap == null) {
            return Maps.newHashMap();
        }

        return obtainEncryptedMapByType(resultMap.getType());
    }

    private Map<String, EncryptField> obtainEncryptedMapByType(Class<?> clazz) {
        Map<String, EncryptField> encryptFieldMap = Maps.newTreeMap();

        for (Field field : clazz.getDeclaredFields()) {
            EncryptField encryptField = field.getAnnotation(EncryptField.class);
            if (Objects.nonNull(encryptField)) {
                encryptFieldMap.put(field.getName(), encryptField);
            }
        }

        return encryptFieldMap;
    }

    private Map<String, EncryptJSONField> consEncryptedJsonMapByResultMap(ResultMap resultMap) {
        if (resultMap == null) {
            return Maps.newHashMap();
        }

        return obtainEncryptedJsonByType(resultMap.getType());
    }

    private Map<String, EncryptJSONField> obtainEncryptedJsonByType(Class<?> clazz) {
        Map<String, EncryptJSONField> encryptedJSONFieldMap = Maps.newTreeMap();

        for (Field field : clazz.getDeclaredFields()) {
            EncryptJSONField encryptJSONField = field.getAnnotation(EncryptJSONField.class);
            if (Objects.nonNull(encryptJSONField)) {
                encryptedJSONFieldMap.put(field.getName(), encryptJSONField);
            }
        }

        return encryptedJSONFieldMap;
    }

    /**
     * 批量解密「已加密」字段，并重新给对象赋值
     */
    private void handleStringValues(MetaObject resultMetaObj, List<String> keys) {
        if (CollectionUtils.isEmpty(keys)) {
            return;
        }

        List<String> needDecryptStringValues = Lists.newArrayList();
        keys.forEach(key -> {
            Object valueObj = resultMetaObj.getValue(key);
            if(valueObj == null) return;
            if(valueObj instanceof String) {
                needDecryptStringValues.add((String)valueObj);
            }
        });
        if(CollectionUtils.isEmpty(needDecryptStringValues)) return;
        List<String> secrets = localBladeService.decryptData(needDecryptStringValues);
        for (int i = 0; i < keys.size(); i++) {
            resultMetaObj.setValue(keys.get(i), secrets.get(i));
        }
    }

    private void handleJsonVal(MetaObject resultMetaObj, String property, EncryptJSONField encryptJSONField) {
        String value = (String) resultMetaObj.getValue(property);
        Map<String, Object> objectMap = JsonUtils.parseToMap(value);
        if (MapUtils.isEmpty(objectMap)) {
            return;
        }

        //log.info("Needed decrypt data: " + JsonUtils.toJson(objectMap));
        List<String> encryptedKeys = Lists.newArrayList(), encryptedValues = Lists.newArrayList();
        try {
            EncryptJSONFieldKey[] keys = encryptJSONField.keys();
            for (EncryptJSONFieldKey jsonFieldKey : keys) {
                String key = jsonFieldKey.key();
                if (objectMap.containsKey(key)) {
                    encryptedKeys.add(key);
                    encryptedValues.add((String) objectMap.get(key));
                }
            }

            List<String> secrets = localBladeService.decryptData(encryptedValues);
            for (int i = 0; i < encryptedKeys.size(); i++) {
                objectMap.put(encryptedKeys.get(i), secrets.get(i));
            }
        } catch (Throwable t) {
            log.error("json 字段解密失败，cause: {}", t.getMessage(), t);
            throw new BusinessException(LepusResultCode.FAIL_BATCH_DECRYPT, t.getMessage());
        }

        resultMetaObj.setValue(property, JsonUtils.toJson(objectMap));
    }

    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {
        // do nothing
    }
}

import java.lang.annotation.Documented;
import java.lang.annotation.ElementType;
import java.lang.annotation.Inherited;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

/**
 * 标记在DTO类上，用于说明是否支持加解密
 */
@Target({ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Inherited
@Documented
public @interface EncryptEnabled {
    /**
     * 是否开启加解密和脱敏模式
     *
     * @return SensitiveEncryptEnabled
     */
    boolean value() default true;
}

import java.lang.annotation.Documented;
import java.lang.annotation.ElementType;
import java.lang.annotation.Inherited;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

/**
 * 标记了注解的字段会在写请求时对数据进行加密，在读请求时进行解密
 */
@Target({ElementType.FIELD, ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
@Inherited
@Documented
public @interface EncryptField {
}

import java.lang.annotation.Documented;
import java.lang.annotation.ElementType;
import java.lang.annotation.Inherited;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

/**
 * 对json内的key_value进行加解密
 */
@Target({ElementType.FIELD, ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
@Inherited
@Documented
public @interface EncryptJSONField {
    /**
     * 需要加解密的字段的数组
     */
    EncryptJSONFieldKey[] keys();
}

import java.lang.annotation.Documented;
import java.lang.annotation.ElementType;
import java.lang.annotation.Inherited;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

/**
 * json字段中需要脱敏的key字段
 *
 * @author chenhaiyang
 */
@Target({ElementType.FIELD, ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
@Inherited
@Documented
public @interface EncryptJSONFieldKey {
    /**
     * json 中的 key name
     */
    String key();
}

import com.xxxx.dao.interceptors.DecryptReadInterceptor;
import com.xxxx.dao.interceptors.EncryptWriteInterceptor;
import com.xxxx.service.LocalBladeService;
import org.mybatis.spring.boot.autoconfigure.ConfigurationCustomizer;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import lombok.RequiredArgsConstructor;

/**
 * 加解密插件配置
 *
 **/
@Configuration
@RequiredArgsConstructor
public class EncryptPluginConfig {

    private final LocalBladeService localBladeService;

    @Bean
    ConfigurationCustomizer configurationCustomizer() {

        EncryptWriteInterceptor encryptWriteInterceptor = new EncryptWriteInterceptor(localBladeService);
        DecryptReadInterceptor decryptReadInterceptor = new DecryptReadInterceptor(localBladeService);

        return configuration -> {
            configuration.addInterceptor(decryptReadInterceptor);
            configuration.addInterceptor(encryptWriteInterceptor);
        };
    }


}

第一个是实现数据加密的类
第二个是实现数据解密的类
后面四个是加密相关的注解
最后一个是加解密插件配置

其中第一个和第二个类中用到了
注解：@Intercepts，StatementHandler，ResultSetHandler类。下面探究一个这三个的用法。

一：@Intercepts

package org.apache.ibatis.plugin;

import java.lang.annotation.Documented;
import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

@Documented
@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.TYPE})
public @interface Intercepts {
    Signature[] value();
}

MyBatis自定义插件机制分析
盘点 MyBatis : 自定义插件的使用及 PageHelper

二: StatementHandler

mybatis四大神器之StatementHandler

三: ResultSetHandler

mybatis四大神器之ResultSetHandler