springboot mybatis中通过添加拦截器实现数据权限过滤

1.自定义注解。在需要的拦截的方法上添加注解

package com.bimstudy.annotation;


import com.bimstudy.base.Constant;

import java.lang.annotation.*;

/**
 * 数据域，用于标注哪个角色的数据域
 */
@Target( { ElementType.METHOD } )
@Retention( RetentionPolicy.RUNTIME )
@Documented
public @interface DataScope {
    Constant.RolesRelation value() default Constant.RolesRelation.Train;
    String scopeName() default "create_user_id";
    /**
     * 是否需要SQL过滤
     */
    boolean needSqlFilter() default true;
}

@DataScope
List<User> findUsers();

自定义数据权限上下文信息

package com.bimstudy.annotation;

import lombok.Builder;
import lombok.Data;

/**
 * 数据域上下文
 */
public class DataScopeRequestContext {

    private final static ThreadLocal<Payload> requestHolder = new ThreadLocal<>();

    public static void add(Payload payload){
        requestHolder.set(payload);
    }
    public static void remove(){
        requestHolder.remove();
    }
    public static Payload get(){
        return requestHolder.get();
    }

    @Data
    @Builder
    public static class Payload{
        /**
         * 用户id
         */
        private Long userId;
        /**
         * 租户id
         */
        private Long tenantId;
        /**
         * 应用id
         */
        private Long appId;
        /**
         * 仅为目标角色
         */
        private Boolean isOnlyTargetRole;
        /**
         * 是否需要过滤
         */
        private Boolean needSqlFilter;
        /**
         * 字段域
         */
        private String scopeName;
    }
}

2.定义spring aop拦截器，在拦截器中写入上下文内容

package com.bimstudy.aspect;

import com.bimstudy.annotation.DataScope;
import com.bimstudy.annotation.DataScopeRequestContext;
import com.bimstudy.service.IUserService;
import com.bimstudy.util.ComUtil;
import com.bimstudy.util.JWTUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.shiro.SecurityUtils;
import org.apache.shiro.subject.Subject;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;

/**
 * 数据域拦截处理
 */
@Slf4j
@Aspect
@Component
@Order
public class DataScopeAop {

    @Autowired
    private IUserService userService;

    @Around("@annotation(dataScope)")
    public Object around(ProceedingJoinPoint jp, DataScope dataScope) throws Throwable {
        Subject subject = SecurityUtils.getSubject();
        Object principals = subject.getPrincipals();
        if (!ComUtil.isEmpty(principals)){
            String token = principals.toString();
            Long userNo = Long.valueOf(JWTUtil.getUserNo(token));
            Long loginFromAppId = Long.valueOf(JWTUtil.getLoginFromAppId(token));
            Long tenantId = Long.valueOf(JWTUtil.getTenantId(token));
            boolean needSqlFilter = dataScope.needSqlFilter();
            boolean onlyTargetRole = userService.isOnlyTargetRole(userNo, tenantId, loginFromAppId,dataScope.value());
            log.info("是否仅为直播管理员:{},needSqlFilter:{}",onlyTargetRole,needSqlFilter);
            DataScopeRequestContext.add(DataScopeRequestContext.Payload.builder().userId(userNo).tenantId(tenantId).appId(loginFromAppId).isOnlyTargetRole(onlyTargetRole).needSqlFilter(needSqlFilter).scopeName(dataScope.scopeName()).build());
        }

        try{
            Object result = jp.proceed();
            return result;
        }catch (Exception e){
            log.error("DataScopeAop 执行业务逻辑异常，",e);
            throw e;
        }finally{
            DataScopeRequestContext.remove();
        }
    }

}

3.自定义mybatis sql拦截器

package com.bimstudy.interceptor;

import com.baomidou.mybatisplus.core.metadata.IPage;
import com.baomidou.mybatisplus.core.parser.ISqlParser;
import com.baomidou.mybatisplus.core.parser.SqlInfo;
import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import com.baomidou.mybatisplus.core.toolkit.StringPool;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.baomidou.mybatisplus.extension.plugins.PaginationInterceptor;
import com.baomidou.mybatisplus.extension.toolkit.SqlParserUtils;
import com.bimstudy.annotation.DataScopeRequestContext;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.ParameterMapping;
import org.apache.ibatis.mapping.SqlCommandType;
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.stereotype.Component;

import java.sql.Connection;
import java.util.*;

/**
 * 数据权限拦截器
 */
@Slf4j
@Component
@Intercepts({@Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})})
public class DataScopeInterceptor extends PaginationInterceptor {
    /**
     * COUNT SQL 解析
     */
    private ISqlParser sqlParser;
    /**
     * 溢出总页数，设置第一页
     */
    private boolean overflow = false;
    /**
     * 方言类型
     */
    private String dialectType;
    /**
     * 方言实现类
     */
    private String dialectClazz;
    /**
     * 定位结尾关键词
     */
    private static final List<String> keyWords = Arrays.asList("group by","order by","limit");

    @Override
    public Object intercept(Invocation invocation) throws Throwable {

        StatementHandler statementHandler = PluginUtils.realTarget(invocation.getTarget());
        MetaObject metaObject = SystemMetaObject.forObject(statementHandler);

        // SQL 解析
        this.sqlParser(metaObject);

        // 先判断是不是SELECT操作
        MappedStatement mappedStatement = (MappedStatement) metaObject.getValue("delegate.mappedStatement");
        if (!SqlCommandType.SELECT.equals(mappedStatement.getSqlCommandType())) {
            return invocation.proceed();
        }

        // 针对定义了rowBounds，做为mapper接口方法的参数
        BoundSql boundSql = (BoundSql) metaObject.getValue("delegate.boundSql");
        Object paramObj = boundSql.getParameterObject();
        List<ParameterMapping> mappings = new ArrayList<>(boundSql.getParameterMappings());

        // 判断参数里是否有page对象
        IPage page = null;
        if (paramObj instanceof IPage) {
            page = (IPage) paramObj;
        } else if (paramObj instanceof Map) {
            for (Object arg : ((Map) paramObj).values()) {
                if (arg instanceof IPage) {
                    page = (IPage) arg;
                    break;
                }
            }
        }

        String originalSql = boundSql.getSql();
        //去除分号
        originalSql = originalSql.replace(";"," ");

        DataScopeRequestContext.Payload payload = DataScopeRequestContext.get();
        if (payload != null) {

            Long userId = payload.getUserId();
            if( payload.getNeedSqlFilter() && userId!=null){

                //不需要分页的场合，如果 size 小于 0 返回结果集
                if (null == page || page.getSize() < 0) {
                    String buildSql = concatExtCondition(originalSql, payload.getScopeName(), userId.toString(),false);
                    metaObject.setValue("delegate.boundSql.sql", buildSql);
                    metaObject.setValue("delegate.boundSql.parameterMappings", mappings);
                    return invocation.proceed();
                }

                String buildSql = concatExtCondition(originalSql, payload.getScopeName(), userId.toString(),true);
                metaObject.setValue("delegate.boundSql.sql", buildSql);
                metaObject.setValue("delegate.boundSql.parameterMappings", mappings);
                //重新计算总页数，总条数信息
                if (page.isSearchCount()) {
                    Connection connection = (Connection) invocation.getArgs()[0];
                    SqlInfo sqlInfo = SqlParserUtils.getOptimizeCountSql(page.optimizeCountSql(), sqlParser, buildSql);
                    String countSql = sqlInfo.getSql().substring(0, sqlInfo.getSql().lastIndexOf("LIMIT"));
                    Predicate mybatis_plus = mapped -> mapped.getProperty().startsWith("mybatis_plus_");
                    List limitParameter =
                            boundSql.getParameterMappings().stream().filter(mybatis_plus).collect(Collectors.toList());
                    boundSql.getParameterMappings().removeIf(mybatis_plus);
                    this.queryTotal(overflow, countSql, mappedStatement, boundSql, page, connection);
                    boundSql.getParameterMappings().addAll(limitParameter);
                    if (page.getTotal() <= 0) {
                        return invocation.proceed();
                    }
                }


            }
            DataScopeRequestContext.remove();
        }
        return invocation.proceed();
    }

    /**
     * 查询SQL拼接 create_user_id 过滤条件
     * @param originalSql
     * @param field
     * @param value
     * @return
     */
    public static String concatExtCondition(String originalSql, String field, String value,boolean pageEnable) {
        if (StringUtils.isNotEmpty(originalSql)) {
            int whereIndex = originalSql.toLowerCase().indexOf("where");
            String conditionString;
            if (whereIndex <0) {
                conditionString = " where " + field + " " + StringPool.EQUALS + " " + value + " ";
            }else {
                conditionString = " and " + field + " " + StringPool.EQUALS + " " + value + "  ";
            }
            return insertConditionString(originalSql, conditionString);
        }
        return originalSql;
    }

    private static String insertConditionString(String originalSql, String conditionString) {
        StringBuilder buildSql = new StringBuilder(originalSql);
        int keyWordIndex = getKeyWordIndex(originalSql);
        if (keyWordIndex > 0) {
            buildSql.insert(keyWordIndex, conditionString);
        } else {
            buildSql.append(conditionString);
        }
        return buildSql.toString();
    }

    private static int getKeyWordIndex(String originalSql) {
        int keyWordIndex=-1;
        for (int i = 0; i < keyWords.size(); i++) {
            String keyWord = keyWords.get(i);
            keyWordIndex = originalSql.toLowerCase().indexOf(keyWord);
            if (keyWordIndex >0) {
                break;
            }
        }
        return keyWordIndex;
    }

    @Override
    public Object plugin(Object target) {
        DataScopeRequestContext.Payload payload = DataScopeRequestContext.get();
        if (target instanceof StatementHandler && payload !=null && payload.getNeedSqlFilter() && payload.getIsOnlyTargetRole()) {
            return Plugin.wrap(target, this);
        }
        return target;
    }

    @Override
    public void setProperties(Properties properties) {
        String dialectType = properties.getProperty("dialectType");
        String dialectClazz = properties.getProperty("dialectClazz");
        if (StringUtils.isNotEmpty(dialectType)) {
            this.dialectType = dialectType;
        }
        if (StringUtils.isNotEmpty(dialectClazz)) {
            this.dialectClazz = dialectClazz;
        }
    }



}

如果觉得文章对你有用，请随意赞赏

springboot mybatis中通过添加拦截器实现数据权限过滤

http://bmap.xyz//archives/1705026835582

作者

清风飒影

发布于

2024-01-12

更新于

2024-01-12

许可协议

CC BY 4.0