ソースを参照

调整 QueryWrapper 需要主动开启检查 SQL 注入过滤

hubin 2 年 前
コミット
2b3cf4c331

+ 24 - 0
mybatis-plus-core/src/main/java/com/baomidou/mybatisplus/core/conditions/query/QueryWrapper.java

@@ -18,10 +18,12 @@ package com.baomidou.mybatisplus.core.conditions.query;
 import com.baomidou.mybatisplus.core.conditions.AbstractWrapper;
 import com.baomidou.mybatisplus.core.conditions.SharedString;
 import com.baomidou.mybatisplus.core.conditions.segments.MergeSegments;
+import com.baomidou.mybatisplus.core.exceptions.MybatisPlusException;
 import com.baomidou.mybatisplus.core.metadata.TableFieldInfo;
 import com.baomidou.mybatisplus.core.metadata.TableInfoHelper;
 import com.baomidou.mybatisplus.core.toolkit.CollectionUtils;
 import com.baomidou.mybatisplus.core.toolkit.StringPool;
+import com.baomidou.mybatisplus.core.toolkit.sql.SqlInjectionUtils;
 
 import java.util.List;
 import java.util.Map;
@@ -77,6 +79,28 @@ public class QueryWrapper<T> extends AbstractWrapper<T, String, QueryWrapper<T>>
         this.sqlFirst = sqlFirst;
     }
 
+
+    /**
+     * 检查 SQL 注入过滤
+     */
+    private boolean checkSqlInjection;
+
+    /**
+     * 开启检查 SQL 注入
+     */
+    public QueryWrapper<T> checkSqlInjection() {
+        this.checkSqlInjection = true;
+        return this;
+    }
+
+    @Override
+    protected String columnToString(String column) {
+        if (checkSqlInjection && SqlInjectionUtils.check(column)) {
+            throw new MybatisPlusException("Discovering SQL injection column: " + column);
+        }
+        return column;
+    }
+
     @Override
     public QueryWrapper<T> select(boolean condition, List<String> columns) {
         if (condition && CollectionUtils.isNotEmpty(columns)) {

+ 11 - 0
mybatis-plus-core/src/test/java/com/baomidou/mybatisplus/core/conditions/QueryWrapperTest.java

@@ -2,8 +2,10 @@ package com.baomidou.mybatisplus.core.conditions;
 
 import com.baomidou.mybatisplus.core.MybatisConfiguration;
 import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
+import com.baomidou.mybatisplus.core.exceptions.MybatisPlusException;
 import com.baomidou.mybatisplus.core.metadata.TableInfoHelper;
 import org.apache.ibatis.builder.MapperBuilderAssistant;
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
 import java.time.LocalDate;
@@ -165,6 +167,15 @@ class QueryWrapperTest extends BaseWrapperTest {
         logParams(wrapper);
     }
 
+    @Test
+    void testCheckSqlInjection() {
+        QueryWrapper qw = new QueryWrapper<Entity>().checkSqlInjection().eq("a", "b");
+        Assertions.assertEquals("WHERE (a = #{ew.paramNameValuePairs.MPGENVAL1})", qw.getCustomSqlSegment());
+
+        qw.orderByAsc("select 1 from xxx");
+        Assertions.assertThrows(MybatisPlusException.class, () -> qw.getCustomSqlSegment());
+    }
+
     private List<Object> getList() {
         List<Object> list = new ArrayList<>();
         for (int i = 0; i < 2; i++) {