Mybatis动态替换表名-拦截器

需求

最近接到一个安全漏洞整改，说是存在Sql注入问题。一看代码，有人写了这样的Sql：

select * from ${tablename} where id=#{id};

老大说要改掉，不然代码检测一直过不去，不能跟其它系统对接。

代码

拦截器

package com.xhfeng.studentmananger.config;

import org.apache.ibatis.binding.MapperMethod;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.reflection.DefaultReflectorFactory;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.factory.DefaultObjectFactory;
import org.apache.ibatis.reflection.wrapper.DefaultObjectWrapperFactory;
import org.springframework.stereotype.Component;

import java.sql.Connection;
import java.util.Objects;
import java.util.Properties;

/**
 * Mybatis预执行替换表
 * @author jike18
 * @Date 2022-09-19 16:13
 */
@Component
@Intercepts({@Signature(type= StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})})
public class MyBatisSwapTableInterceptor implements Interceptor {

    private final static String SOURCE_TABLE = "wait_swap_table";

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        StatementHandler statementHandler = (StatementHandler) invocation.getTarget();
        MetaObject metaObject = MetaObject.forObject(statementHandler, new DefaultObjectFactory(),
                new DefaultObjectWrapperFactory(), new DefaultReflectorFactory());
        BoundSql boundSql = (BoundSql)metaObject.getValue("delegate.boundSql");
        // 获取执行的SQL
        String sql = boundSql.getSql();
        // 获取执行的参数
        Object parameterObject = boundSql.getParameterObject();
        // 需要拦截的参数是map类型，{tableName="",id=""...}，因此可以转Map的继续，其他的放行
        if(parameterObject instanceof MapperMethod.ParamMap){
            // 判断是否存在需要替换的表，不存在放行
            if (sql.contains(SOURCE_TABLE)) {
                // 替换的表名，从参数中获取
                String tableName = String.valueOf(((MapperMethod.ParamMap) parameterObject).get("tableName"));
                if(Objects.nonNull(tableName)){
                    sql = sql.replaceAll(SOURCE_TABLE, tableName);
                    // 替换执行的的SQL.
                    metaObject.setValue("delegate.boundSql.sql", sql);
                }
            }
        }
        return invocation.proceed();
    }

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

    @Override
    public void setProperties(Properties properties) {
    }
}

修改后的sql

select * from wait_swap_table where id=#{id};