miemie 4 سال پیش
والد
کامیت
77fe13a94b

+ 22 - 2
mybatis-plus-extension/src/main/java/com/baomidou/mybatisplus/extension/plugins/inner/BlockAttackInnerInterceptor.java

@@ -19,6 +19,9 @@ import com.baomidou.mybatisplus.core.plugins.InterceptorIgnoreHelper;
 import com.baomidou.mybatisplus.core.toolkit.Assert;
 import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
 import com.baomidou.mybatisplus.extension.parser.JsqlParserSupport;
+import net.sf.jsqlparser.expression.Expression;
+import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
+import net.sf.jsqlparser.expression.operators.relational.NotEqualsTo;
 import net.sf.jsqlparser.statement.delete.Delete;
 import net.sf.jsqlparser.statement.update.Update;
 import org.apache.ibatis.executor.statement.StatementHandler;
@@ -50,11 +53,28 @@ public class BlockAttackInnerInterceptor extends JsqlParserSupport implements In
 
     @Override
     protected void processDelete(Delete delete, int index, Object obj) {
-        Assert.notNull(delete.getWhere(), "Prohibition of full table deletion");
+        this.checkWhere(delete.getWhere(), "Prohibition of full table deletion");
     }
 
     @Override
     protected void processUpdate(Update update, int index, Object obj) {
-        Assert.notNull(update.getWhere(), "Prohibition of table update operation");
+        this.checkWhere(update.getWhere(), "Prohibition of table update operation");
+    }
+
+    protected void checkWhere(Expression where, String ex) {
+        Assert.notNull(where, ex);
+        if (where instanceof EqualsTo) {
+            // example: 1=1
+            EqualsTo equalsTo = (EqualsTo) where;
+            Expression leftExpression = equalsTo.getLeftExpression();
+            Expression rightExpression = equalsTo.getRightExpression();
+            Assert.isFalse(leftExpression.toString().equals(rightExpression.toString()), ex);
+        } else if (where instanceof NotEqualsTo) {
+            // example: 1 != 2
+            NotEqualsTo notEqualsTo = (NotEqualsTo) where;
+            Expression leftExpression = notEqualsTo.getLeftExpression();
+            Expression rightExpression = notEqualsTo.getRightExpression();
+            Assert.isTrue(leftExpression.toString().equals(rightExpression.toString()), ex);
+        }
     }
 }

+ 21 - 4
mybatis-plus-extension/src/test/java/com/baomidou/mybatisplus/extension/plugins/inner/BlockAttackInnerInterceptorTest.java

@@ -14,14 +14,31 @@ class BlockAttackInnerInterceptorTest {
     private final BlockAttackInnerInterceptor interceptor = new BlockAttackInnerInterceptor();
 
     @Test
-    void test() {
+    void update() {
+        check("update user set name = null", "null where");
+        check("update user set name = null where 1=1", "1=1");
+        check("update user set name = null where 1<>2", "1<>2");
+        check("update user set name = null where 1!=2", "1!=2");
+//        check("update user set name = null where 1=1 and 2=2", "1=1 and 2=2");
+    }
+
+    @Test
+    void delete() {
+        check("delete from user", "null where");
+        check("delete from user where 1=1", "1=1");
+        check("delete from user where 1<>2", "1<>2");
+        check("delete from user where 1!=2", "1!=2");
+//        check("delete from user where 1=1 and 2=2", "1=1 and 2=2");
+    }
+
+    void check(String sql, String as) {
         Exception e = null;
         try {
-            interceptor.parserSingle("update set name = null", null);
+            interceptor.parserSingle(sql, null);
         } catch (Exception x) {
             e = x;
         }
-        assertThat(e).isNotNull();
-        assertThat(e).isInstanceOf(MybatisPlusException.class);
+        assertThat(e).as(as).isNotNull();
+        assertThat(e).as(as).isInstanceOf(MybatisPlusException.class);
     }
 }