Spring Boot 3 整合 Mybatis-Plus 实现数据权限控制

@Aspect
@Component
public class DataScopeAspect {

	@Resource
	private SysDeptService sysDeptService;

	@Resource
	private SysUserRoleService sysUserRoleService;

	@Resource
	private SysRoleDeptService sysRoleDeptService;

	/**
	 * 配置织入点
	 */
	@Pointcut("@annotation(cn.harry.common.annotation.DataScope)")
	public void dataScopePointCut() {
	}

	@Before("dataScopePointCut()")
	public void doBefore(JoinPoint point) {
		handleDataScope(point);
	}

	protected void handleDataScope(final JoinPoint joinPoint) {
		Object params = joinPoint.getArgs()[0];

		if (params instanceof Map) {
			SysUser user = SecurityUtils.getSysUser();
			// 如果不是超级管理员，则进行数据过滤
			if (!SecurityUtils.isAdmin(user.getId())) {
				// 根据用户ID 获取数据权限标识 如果数据权限包含全部数据 直接返回
				List<Integer> dataScopes = sysUserRoleService
						.listDataScopesByUserId(user.getId());
				if (!CollectionUtils.isEmpty(dataScopes)) {
					if (dataScopes.contains(DataScopeEnums.ALL.getKey())) {
						return;
					}
				}
				Map map = (Map) params;
				map.put(CommonConstant.SQL_FILTER, this.getSQLFilter(user, joinPoint));
			}
			return;
		}
		throw new ApiException("数据权限接口，只能是Map类型参数，且不能为NULL");
	}

	/**
	 * 获取数据过滤的SQL
	 * @param user
	 * @param point
	 * @return
	 */
	private String getSQLFilter(SysUser user, JoinPoint point) {
		MethodSignature signature = (MethodSignature) point.getSignature();
		DataScope dataFilter = signature.getMethod().getAnnotation(DataScope.class);
		// 获取表的别名
		String tableAlias = dataFilter.tableAlias();
		if (StringUtils.isNotBlank(tableAlias)) {
			tableAlias += ".";
		}
		// 部门ID列表
		Set<Long> deptIdList = new HashSet<>();
		StringBuilder sqlFilter = new StringBuilder();

		// 1. 根据用户ID 获取角色列表
		List<Long> roleIdList = sysUserRoleService.listRoleIdByUserId(user.getId());
		if (!roleIdList.isEmpty()) {
			// 2.用户角色对应的部门ID列表
			List<Long> idList = sysRoleDeptService.queryDeptIdList(roleIdList);
			deptIdList.addAll(idList);
		}

		// 用户子部门ID列表
		if (dataFilter.subDept()) {
			List<Long> subDeptIdList = sysDeptService.getSubDeptIdList(user.getDeptId());
			deptIdList.addAll(subDeptIdList);
		}
		sqlFilter.append(" (");

		if (!deptIdList.isEmpty()) {
			sqlFilter.append(tableAlias).append(dataFilter.deptId()).append(" in(")
					.append(StringUtils.join(deptIdList, ",")).append(")");
		}

		// 没有设置本部门数据权限，也能查询本部门数据 user.getDeptId()
		if (dataFilter.user()) {
			if (!deptIdList.isEmpty()) {
				sqlFilter.append(" or ");
			}
			sqlFilter.append(tableAlias).append(dataFilter.deptId()).append("=")
					.append(user.getDeptId());
		}

		sqlFilter.append(")");

		if ("()".equals(sqlFilter.toString().trim())) {
			return null;
		}

		return sqlFilter.toString();
	}

}

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface DataScope {

	/** 表的别名 */
	String tableAlias() default "";

	/** true：没有本部门数据权限，也能查询本人数据 */
	boolean user() default true;

	/** true：拥有子部门数据权限 */
	boolean subDept() default false;

	/** 部门ID */
	String deptId() default "dept_id";

	/** 用户ID */
	String userId() default "user_id";

}

/**
 * 数据权限控制器
 *
 * @author harry
 * @公众号 Harry技术
 */
@Slf4j
public class MybatisPlusDataPermissionHandler implements DataPermissionHandler {

    @Override
    @SneakyThrows
    public Expression getSqlSegment(Expression where, String mappedStatementId) {

        log.info(" 数据权限控制器 ：{}", mappedStatementId);

        Class<?> clazz = Class.forName(mappedStatementId.substring(0, mappedStatementId.lastIndexOf(StringPool.DOT)));
        String methodName = mappedStatementId.substring(mappedStatementId.lastIndexOf(StringPool.DOT) + 1);
        Method[] methods = clazz.getDeclaredMethods();
        for (Method method : methods) {
            if (method.getName().equals(methodName)) {
                DataScope annotation = method.getAnnotation(DataScope.class);
                // 如果没有注解或者是超级管理员，直接返回
                if (annotation == null || SecurityUtils.isRoot()) {
                    return where;
                }
                return dataScopeFilter(annotation, where);
            }
        }
        return where;
    }

    @SneakyThrows
    private Expression dataScopeFilter(DataScope annotation, Expression where) {
        String deptAlias = annotation.deptAlias();
        String userAlias = annotation.userAlias();
        String deptIdColumnName = annotation.deptIdColumnName();
        String userIdColumnName = annotation.userIdColumnName();


        String deptColumnName = StrUtil.isNotBlank(deptAlias) ? (deptAlias + StringPool.DOT + deptIdColumnName) : deptIdColumnName;
        String userColumnName = StrUtil.isNotBlank(userAlias) ? (userAlias + StringPool.DOT + userIdColumnName) : userIdColumnName;

        // 获取当前用户的数据权限
        Integer dataScope = SecurityUtils.getDataScope();

        DataScopeEnums dataScopeEnum = IBaseEnum.getEnumByValue(dataScope, DataScopeEnums.class);

        Long deptId, userId;
        String appendSqlStr;
        switch (dataScopeEnum) {
            case ALL:
                return where;
            case DEPT:
                deptId = SecurityUtils.getDeptId();
                appendSqlStr = deptColumnName + StringPool.EQUALS + deptId;
                break;
            case SELF:
                userId = SecurityUtils.getUserId();
                appendSqlStr = userColumnName + StringPool.EQUALS + userId;
                break;
            // 默认部门及子部门数据权限
            default:
                deptId = SecurityUtils.getDeptId();
                appendSqlStr = deptColumnName + " IN ( SELECT id FROM sys_dept WHERE id = " + deptId + " OR FIND_IN_SET( " + deptId + " , ancestors ) )";
                break;
        }

        if (StrUtil.isBlank(appendSqlStr)) {
            return where;
        }

        Expression appendExpression = CCJSqlParserUtil.parseCondExpression(appendSqlStr);

        if (where == null) {
            return appendExpression;
        }

        return new AndExpression(where, appendExpression);
    }
}

@Override
@DataScope(userIdColumnName = "id")
List<SysUser> selectList(IPage<SysUser> page, @Param(Constants.WRAPPER) Wrapper<SysUser> queryWrapper);

/**
 * 数据权限枚举
 *
 * @author harry
 * @公众号 Harry技术
 */
@Getter
@AllArgsConstructor
public enum DataScopeEnums implements IBaseEnum<Integer> {
    /**
     * value 越小，数据权限范围越大
     */
    ALL(0, "所有数据权限"),
    DEPT_AND_SUB(1, "本部门及子部门数据"),
    DEPT(2, "本部门数据权限"),
    SELF(3, "本人数据权限"),
    ;
    private final Integer value;

    private final String label;

}