Bugfix: issue#2 XML foreach Error

RebortBoss · Oct 15, 2020 · 058b25b · 058b25b
1 parent d42c3c9
commit 058b25b
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 31 deletions.
diff --git a/src/main/java/com/immomo/momosec/lang/xml/rule/momosecurity/MybatisXmlSQLi.java b/src/main/java/com/immomo/momosec/lang/xml/rule/momosecurity/MybatisXmlSQLi.java
@@ -22,10 +22,7 @@
 import com.intellij.lang.ASTFactory;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.*;
-import com.intellij.psi.xml.XmlDocument;
-import com.intellij.psi.xml.XmlElementType;
-import com.intellij.psi.xml.XmlTag;
-import com.intellij.psi.xml.XmlText;
+import com.intellij.psi.xml.*;
 import com.intellij.xml.util.XmlUtil;
 import org.apache.commons.lang3.StringUtils;
 import org.jetbrains.annotations.Nls;
@@ -167,9 +164,13 @@ private void fixXmlText(XmlText xmlText, int offset) {
                             suffixXmlText = textElements[0];
                         }
 
-                        // 4. 先修正尾部文本再追加
-                        fixXmlText(suffixXmlText, 0);
+                        // 4. 先追加再修正尾部文本
                         parent.add(suffixXmlText);
+
+                        XmlTagChild[] xmlTagChildren = parent.getValue().getChildren();
+                        if (xmlTagChildren[xmlTagChildren.length - 1] instanceof XmlText) {
+                            fixXmlText((XmlText)xmlTagChildren[xmlTagChildren.length - 1], 0);
+                        }
                     } else {
                         fixXmlText(xmlText, m.end());
                     }

diff --git a/src/test/java/com/immomo/momosec/lang/xml/rule/momosecurity/MybatisXmlSQLiTest.java b/src/test/java/com/immomo/momosec/lang/xml/rule/momosecurity/MybatisXmlSQLiTest.java
@@ -20,6 +20,7 @@
 import com.intellij.codeInspection.ProblemHighlightType;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.XmlElementFactory;
+import com.intellij.psi.impl.source.xml.XmlTagImpl;
 import com.intellij.psi.xml.XmlTag;
 import com.intellij.psi.xml.XmlText;
 import com.intellij.testFramework.MockProblemDescriptor;
@@ -62,41 +63,41 @@ public void testNotHandleTableVar() {
     }
 
     public void testWhereInQuickFix() {
-        Project project = myFixture.getProject();
-        XmlText text = XmlElementFactory.getInstance(project).createDisplayText(
-                "where id in ${ids}"
-        );
-
-        MockProblemDescriptor descriptor = new MockProblemDescriptor(text, "", ProblemHighlightType.GENERIC_ERROR_OR_WARNING);
-        MybatisXmlSQLi.MybatisXmlSQLiQuickFix quickFix = new MybatisXmlSQLi.MybatisXmlSQLiQuickFix();
-        quickFix.applyFix(project, descriptor);
+        commonWhereInReplaceTest(
+                "where id in ${ids}",
 
-        assert text.getNextSibling() instanceof XmlTag;
-        XmlTag foreach = (XmlTag)text.getNextSibling();
-        Assert.assertEquals(
+                "where id in \n" +
                 "<foreach collection=\"ids\" item=\"idsItem\" open=\"(\" separator=\",\" close=\")\">\n" +
                 "#{idsItem}\n" +
-                "</foreach>",
-                foreach.getText());
+                "</foreach>\n"
+        );
     }
 
-    public void testWhereInWithQuoteQuickFix() {
-        Project project = myFixture.getProject();
-        XmlText text = XmlElementFactory.getInstance(project).createDisplayText(
-                "where id in (${ids})"
+    public void testDoubleWhereInQuickFix() {
+        commonWhereInReplaceTest(
+                "and (createdBy in ${userNameList} or projectId IN ${id})",
+
+                "and (createdBy in \n" +
+                "<foreach collection=\"userNameList\" item=\"userNameListItem\" open=\"(\" separator=\",\" close=\")\">\n" +
+                "#{userNameListItem}\n" +
+                "</foreach>\n" +
+                " or projectId IN \n" +
+                "<foreach collection=\"id\" item=\"idItem\" open=\"(\" separator=\",\" close=\")\">\n" +
+                "#{idItem}\n" +
+                "</foreach>\n" +
+                ")"
         );
 
-        MockProblemDescriptor descriptor = new MockProblemDescriptor(text, "", ProblemHighlightType.GENERIC_ERROR_OR_WARNING);
-        MybatisXmlSQLi.MybatisXmlSQLiQuickFix quickFix = new MybatisXmlSQLi.MybatisXmlSQLiQuickFix();
-        quickFix.applyFix(project, descriptor);
+    }
 
-        assert text.getNextSibling() instanceof XmlTag;
-        XmlTag foreach = (XmlTag)text.getNextSibling();
-        Assert.assertEquals(
+    public void testWhereInWithQuoteQuickFix() {
+        commonWhereInReplaceTest(
+                "where id in (${ids})",
+                "where id in \n" +
                 "<foreach collection=\"ids\" item=\"idsItem\" open=\"(\" separator=\",\" close=\")\">\n" +
                 "#{idsItem}\n" +
-                "</foreach>",
-                foreach.getText());
+                "</foreach>\n"
+        );
     }
 
     public void testLikeQuickFix() {
@@ -129,4 +130,15 @@ public void commonHashReplaceTest(String origin, String expect) {
 
         Assert.assertEquals(expect, text.getText());
     }
+
+    public void commonWhereInReplaceTest(String origin, String expect) {
+        Project project = myFixture.getProject();
+        XmlText text = XmlElementFactory.getInstance(project).createDisplayText(origin);
+
+        MockProblemDescriptor descriptor = new MockProblemDescriptor(text, "", ProblemHighlightType.GENERIC_ERROR_OR_WARNING);
+        MybatisXmlSQLi.MybatisXmlSQLiQuickFix quickFix = new MybatisXmlSQLi.MybatisXmlSQLiQuickFix();
+        quickFix.applyFix(project, descriptor);
+
+        Assert.assertEquals(expect, ((XmlTagImpl) text.getParent()).getValue().getText());
+    }
 }