新增数据权限处理器

Author: qmdx

mybatis-plus-annotation/src/main/java/com/baomidou/mybatisplus/annotation/InterceptorIgnore.java

@@ -40,4 +40,9 @@
      * 垃圾SQL拦截 {@link com.baomidou.mybatisplus.extension.plugins.inner.IllegalSQLInnerInterceptor}
      */
     String illegalSql() default "";
+
+    /**
+     * 数据权限 {@link com.baomidou.mybatisplus.extension.plugins.inner.DataPermissionInterceptor}
+     */
+    String dataPermission() default "";
 }

mybatis-plus-core/src/main/java/com/baomidou/mybatisplus/core/plugins/InterceptorIgnoreHelper.java

@@ -74,6 +74,10 @@ public static boolean willIgnoreIllegalSql(String id) {
         return willIgnore(id, i -> i.getIllegalSql() != null && i.getIllegalSql());
     }
 
+    public static boolean willIgnoreDataPermission(String id) {
+        return willIgnore(id, i -> i.getDataPermission() != null && i.getDataPermission());
+    }
+
     public static boolean willIgnore(String id, Function<InterceptorIgnoreCache, Boolean> function) {
         InterceptorIgnoreCache cache = INTERCEPTOR_IGNORE_CACHE.get(id);
         if (cache != null) {
@@ -130,5 +134,6 @@ public static class InterceptorIgnoreCache {
         private Boolean dynamicTableName;
         private Boolean blockAttack;
         private Boolean illegalSql;
+        private Boolean dataPermission;
     }
 }

mybatis-plus-extension/src/main/java/com/baomidou/mybatisplus/extension/parser/JsqlParserSupport.java

@@ -14,6 +14,8 @@
 import org.apache.ibatis.logging.LogFactory;
 
 /**
+ * https://github.com/JSQLParser/JSqlParser
+ *
  * @author miemie
  * @since 2020-06-22
  */

mybatis-plus-extension/src/main/java/com/baomidou/mybatisplus/extension/plugins/handler/DataPermissionHandler.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright (c) 2011-2020, baomidou (jobob@qq.com).
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package com.baomidou.mybatisplus.extension.plugins.handler;
+
+import net.sf.jsqlparser.expression.Expression;
+
+/**
+ * 数据权限处理器
+ *
+ * @author hubin
+ * @since 3.4.1 +
+ */
+public interface DataPermissionHandler {
+
+    /**
+     * 获取数据权限 SQL 片段
+     *
+     * @param where             待执行 SQL Where 条件表达式
+     * @param mappedStatementId Mybatis MappedStatement Id 根据该参数可以判断具体执行方法
+     * @return JSqlParser 条件表达式
+     */
+    Expression getSqlSegment(Expression where, String mappedStatementId);
+}

mybatis-plus-extension/src/main/java/com/baomidou/mybatisplus/extension/plugins/handler/TableNameHandler.java

@@ -16,6 +16,8 @@
 package com.baomidou.mybatisplus.extension.plugins.handler;
 
 /**
+ * 动态表名处理器
+ *
  * @author miemie
  * @since 3.4.0
  */

mybatis-plus-extension/src/main/java/com/baomidou/mybatisplus/extension/plugins/inner/DataPermissionInterceptor.java

@@ -0,0 +1,64 @@
+/*
+ * Copyright (c) 2011-2020, baomidou (jobob@qq.com).
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package com.baomidou.mybatisplus.extension.plugins.inner;
+
+import com.baomidou.mybatisplus.core.plugins.InterceptorIgnoreHelper;
+import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
+import com.baomidou.mybatisplus.extension.parser.JsqlParserSupport;
+import com.baomidou.mybatisplus.extension.plugins.handler.DataPermissionHandler;
+import lombok.*;
+import net.sf.jsqlparser.expression.Expression;
+import net.sf.jsqlparser.statement.select.PlainSelect;
+import net.sf.jsqlparser.statement.select.Select;
+import org.apache.ibatis.executor.Executor;
+import org.apache.ibatis.mapping.BoundSql;
+import org.apache.ibatis.mapping.MappedStatement;
+import org.apache.ibatis.session.ResultHandler;
+import org.apache.ibatis.session.RowBounds;
+
+import java.sql.SQLException;
+
+/**
+ * 数据权限处理器
+ *
+ * @author hubin
+ * @since 3.4.1 +
+ */
+@Data
+@NoArgsConstructor
+@AllArgsConstructor
+@ToString(callSuper = true)
+@EqualsAndHashCode(callSuper = true)
+@SuppressWarnings({"rawtypes"})
+public class DataPermissionInterceptor extends JsqlParserSupport implements InnerInterceptor {
+    private DataPermissionHandler dataPermissionHandler;
+
+    @Override
+    public void beforeQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) throws SQLException {
+        if (InterceptorIgnoreHelper.willIgnoreDataPermission(ms.getId())) return;
+        PluginUtils.MPBoundSql mpBs = PluginUtils.mpBoundSql(boundSql);
+        mpBs.sql(parserSingle(mpBs.sql(), ms.getId()));
+    }
+
+    @Override
+    protected void processSelect(Select select, int index, String sql, Object obj) {
+        PlainSelect plainSelect = (PlainSelect) select.getSelectBody();
+        Expression sqlSegment = dataPermissionHandler.getSqlSegment(plainSelect.getWhere(), (String) obj);
+        if (null != sqlSegment) {
+            plainSelect.setWhere(sqlSegment);
+        }
+    }
+}

mybatis-plus-extension/src/test/java/com/baomidou/mybatisplus/extension/plugins/inner/DataPermissionInterceptorTest.java

@@ -0,0 +1,90 @@
+package com.baomidou.mybatisplus.extension.plugins.inner;
+
+import com.baomidou.mybatisplus.extension.plugins.handler.DataPermissionHandler;
+import net.sf.jsqlparser.JSQLParserException;
+import net.sf.jsqlparser.expression.Expression;
+import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
+import net.sf.jsqlparser.expression.operators.conditional.OrExpression;
+import net.sf.jsqlparser.parser.CCJSqlParserUtil;
+import org.junit.jupiter.api.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * 数据权限拦截器测试
+ *
+ * @author hubin
+ * @since 3.4.1 +
+ */
+public class DataPermissionInterceptorTest {
+    private static String TEST_1 = "com.baomidou.userMapper.selectByUsername";
+    private static String TEST_2 = "com.baomidou.userMapper.selectById";
+    private static String TEST_3 = "com.baomidou.roleMapper.selectByCompanyId";
+    private static String TEST_4 = "com.baomidou.roleMapper.selectById";
+
+    /**
+     * 这里可以理解为数据库配置的数据权限规则 SQL
+     */
+    private static Map<String, String> sqlSegmentMap = new HashMap<String, String>() {
+        {
+            put(TEST_1, "username='123' or userId IN (1,2,3)");
+            put(TEST_2, "u.state=1 and u.amount > 1000");
+            put(TEST_3, "companyId in (1,2,3)");
+            put(TEST_4, "username like 'abc%'");
+        }
+    };
+
+    private static DataPermissionInterceptor interceptor = new DataPermissionInterceptor(new DataPermissionHandler() {
+
+        @Override
+        public Expression getSqlSegment(Expression where, String mappedStatementId) {
+            try {
+                String sqlSegment = sqlSegmentMap.get(mappedStatementId);
+                Expression sqlSegmentExpression = CCJSqlParserUtil.parseCondExpression(sqlSegment);
+                if (null != where) {
+                    System.out.println("原 where = " + where.toString());
+                    if (mappedStatementId.equals(TEST_4)) {
+                        // 这里测试返回 OR 条件
+                        return new OrExpression(where, sqlSegmentExpression);
+                    }
+                    return new AndExpression(where, sqlSegmentExpression);
+                }
+                return sqlSegmentExpression;
+            } catch (JSQLParserException e) {
+                e.printStackTrace();
+            }
+            return null;
+        }
+    });
+
+    @Test
+    void test1() {
+        assertSql(TEST_1, "select * from sys_user",
+            "SELECT * FROM sys_user WHERE username = '123' OR userId IN (1, 2, 3)");
+    }
+
+    @Test
+    void test2() {
+        assertSql(TEST_2, "select u.username from sys_user u join sys_user_role r on u.id=r.user_id where r.role_id=3",
+            "SELECT u.username FROM sys_user u JOIN sys_user_role r ON u.id = r.user_id WHERE r.role_id = 3 AND u.state = 1 AND u.amount > 1000");
+    }
+
+    @Test
+    void test3() {
+        assertSql(TEST_3, "select * from sys_role where company_id=6",
+            "SELECT * FROM sys_role WHERE company_id = 6 AND companyId IN (1, 2, 3)");
+    }
+
+    @Test
+    void test4() {
+        assertSql(TEST_4, "select * from sys_role where id=3",
+            "SELECT * FROM sys_role WHERE id = 3 OR username LIKE 'abc%'");
+    }
+
+    void assertSql(String mappedStatementId, String sql, String targetSql) {
+        assertThat(interceptor.parserSingle(sql, mappedStatementId)).isEqualTo(targetSql);
+    }
+}