Mybatis executes the sql core method to add, delete and change Executor.update, query ResultSetHandler.handleResultSets

1. The underlying core interface for mybatis to add, delete and modify is Executor.update. If we want to intercept it, we need to intercept this method

2. The underlying core interface for mybatis to execute queries is ResultSetHandler.handleResultSets. If we want to intercept it, we need to intercept this method

 

Define the interceptor for adding, deleting and modifying:

@Intercepts({
    @Signature(type = Executor.class, method = "update", args = {MappedStatement.class, Object.class}),
})
@Slf4j
public class IbatisUpdateInterceptor implements Interceptor {
    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        try {
            if (invocation.getTarget() instanceof Executor && invocation.getArgs().length == 2) {
                return invokeUpdate(invocation);
            }
        } finally {
            // Roll back to native
            MybatisContextHolder.rollback();
        }

        return invocation.proceed();
    }

    private Object invokeUpdate(Invocation invocation) throws Exception {
        Executor executor = (Executor) invocation.getTarget();
        // Get the first parameter
        MappedStatement ms = (MappedStatement) invocation.getArgs()[0];

        // Non insert/update, do not process
        if (ms.getSqlCommandType() != SqlCommandType.INSERT && ms.getSqlCommandType() != SqlCommandType.UPDATE) {
            return invocation.proceed();
        }

        Object paramObj = invocation.getArgs()[1];
        // No processing without parameters
        if (paramObj == null) {
            return invocation.proceed();
        }

        if (!needEncryptDecrypt(paramObj) && !isBatchUpdate(paramObj)) {
            return invocation.proceed();
        }

        // encryption
        if (paramObj instanceof Map) {
            Map map = (Map) paramObj;

            // Batch update
            if (map.containsKey("list")) {
                encryptList((Collection) map.get("list"));
            } else {
                for (Object updateObj : map.values()) {
                    if (needEncryptDecrypt(updateObj)) {
                        encryptSensitiveFields(updateObj);
                    }
                }
            }

        } else if(paramObj instanceof Collection) {
            encryptList((Collection) paramObj);
        } else {
            encryptSensitiveFields(paramObj);
        }

        return executor.update(ms, paramObj);
    }

    private static boolean needEncryptDecrypt(Object obj) {
        Class<?> clazz = obj.getClass();
        EncryptDecrypt annotation = AnnotationUtils.findAnnotation(clazz, EncryptDecrypt.class);

        return annotation != null && annotation.needEncryptDecrypt();
    }


    private void encryptList(Collection params) throws Exception {
        for (Object obj : params) {
            encryptSensitiveFields(obj);
        }
    }

    private boolean isBatchUpdate(Object paramObj) {
        if (!(paramObj instanceof Map)) {
            return false;
        }

        Map map = (Map) paramObj;
        if (map.containsKey("list")) {
            // For batch update
            return needEncryptDecrypt(map.get("list"));
        } else {
            for (Object updateObj : map.values()) {
                // For updateByExample and updateByExampleSelective
                if (needEncryptDecrypt(updateObj)) {
                    return true;
                }
            }
        }

        return false;
    }

    public static List<Field> getAllFields(Class clazz) {
        List<Field> allFields = Lists.newArrayList();
        while (clazz != null && clazz != Object.class) {
            Field[] tmpFields = clazz.getDeclaredFields();
            if (tmpFields != null) {
                allFields.addAll(Arrays.asList(tmpFields));
            }

            clazz = clazz.getSuperclass();
        }

        return allFields;
    }

    public static boolean isStatic(Field f){
        return (f.getModifiers() & Modifier.STATIC) == Modifier.STATIC;
    }

    private Object encryptSensitiveFields(Object paramObj) throws InvocationTargetException, IllegalAccessException {
        List<Field> allFields = getAllFields(paramObj.getClass());

        for (Field field : allFields) {
            Sensitive sensitive = field.getAnnotation(Sensitive.class);
            if (sensitive == null) {
                continue;
            }

            PropertyDescriptor ps = BeanUtils.getPropertyDescriptor(paramObj.getClass(), field.getName());
            if (ps.getReadMethod() == null || ps.getWriteMethod() == null) {
                continue;
            }

            Object value = ps.getReadMethod().invoke(paramObj);

            if (value != null) {
                ps.getWriteMethod().invoke(paramObj, encryptDbField(value));
                // register rollback
                MybatisContextHolder.registerRollbackEvent(() -> ps.getWriteMethod().invoke(paramObj, value));
            }
        }

        return paramObj;
    }

    private Object encryptDbField(Object value) {
        //encryption
        return null;
    }

    @Override
    public Object plugin(Object target) {
        return null;
    }

    @Override
    public void setProperties(Properties properties) {

    }
}

 

To define a query result Interceptor:

@Intercepts({
    @Signature(type = ResultSetHandler.class, method = "handleResultSets", args = Statement.class)
})
@Slf4j
public class IbatisResultInterceptor implements Interceptor {

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        Object result = invocation.proceed();
        // Decrypt
        if (result instanceof Collection) {
            Collection<Object> objList = (Collection) result;

            List<Object> decryptList = Lists.newArrayList();
            for (Object obj : objList) {
                decryptList.add(decrypt(obj));
            }
            return decryptList;
        } else {
            return decrypt(result);
        }
    }

    private Object decrypt(Object obj) throws InvocationTargetException, IllegalAccessException {
        List<Field> allFields = getAllFields(obj.getClass());

        for (Field field : allFields) {
            Sensitive sensitive = field.getAnnotation(Sensitive.class);
            // Non sensitive information
            if (sensitive == null) {
                continue;
            }

            PropertyDescriptor ps = BeanUtils.getPropertyDescriptor(obj.getClass(), field.getName());
            if (ps.getReadMethod() == null || ps.getWriteMethod() == null) {
                continue;
            }

            Object value = ps.getReadMethod().invoke(obj);

            if (value != null) {
                try {
                    ps.getWriteMethod().invoke(obj, decryptDbField((String) value));
                } catch (Exception e) {
                    log.error("decrypt error of field:{}", field.getName());
                    throw e;
                }
            }
        }

        return obj;
    }

    public List<Field> getAllFields(Class clazz) {
        List<Field> allFields = Lists.newArrayList();
        while (clazz != null && clazz != Object.class) {
            Field[] tmpFields = clazz.getDeclaredFields();
            if (tmpFields != null) {
                allFields.addAll(Arrays.asList(tmpFields));
            }

            clazz = clazz.getSuperclass();
        }

        return allFields;
    }

    /**
     * Decrypt field
     * @param value
     * @return
     */
    private Object decryptDbField(String value) {
        //Decryption, rsa,aes,des
        return null;
    }

    @Override
    public Object plugin(Object target) {
        return null;
    }

    @Override
    public void setProperties(Properties properties) {

    }
}

 

Register Interceptor:

@Bean
    @ConditionalOnProperty(value = "mybatis.encrypt.enabled", havingValue = "true")
    public IbatisResultInterceptor myIbatisResultInterceptor (
        SqlSessionFactory sqlSessionFactory) {
        IbatisResultInterceptor interceptor = new IbatisResultInterceptor();

        Properties properties = new Properties();
        interceptor.setProperties(properties);

        sqlSessionFactory.getConfiguration().addInterceptor(interceptor);

        return interceptor;
    }

 

172 original articles published, 65 praised, 330000 visitors+
Private letter follow

Tags: Mybatis

Posted on Thu, 16 Jan 2020 11:16:36 -0500 by Kestrad