1. fix green screen when use XXE fix

2. support more XML Parser to detect XXE vuln
RebortBoss · Oct 15, 2020 · d42c3c9 · d42c3c9
1 parent b4dd612
commit d42c3c9
Show file tree

Hide file tree

Showing 17 changed files with 646 additions and 161 deletions.
diff --git a/src/main/java/com/immomo/momosec/lang/java/rule/momosecurity/XxeInspector.java b/src/main/java/com/immomo/momosec/lang/java/rule/momosecurity/XxeInspector.java
diff --git a/src/test/java/com/immomo/momosec/lang/java/MomoJavaCodeInsightFixtureTestCase.java b/src/test/java/com/immomo/momosec/lang/java/MomoJavaCodeInsightFixtureTestCase.java
@@ -46,7 +46,7 @@ protected void doTest(@NotNull InspectionProfileEntry inspection, @NotNull Strin
         myFixture.testHighlightingAllFiles(true, false, false, names);
     }
 
-    protected void testQuickFixInClassInitializer(String classInner, String expectAddStat, LocalQuickFix quickFix) {
+    protected void testQuickFixInClassInitializer(String classInner, String[] expectAddStats, LocalQuickFix quickFix) {
         Project project = myFixture.getProject();
         PsiClass aClass = JavaPsiFacade.getElementFactory(project).createClassFromText(classInner, null);
 
@@ -65,17 +65,24 @@ protected void testQuickFixInClassInitializer(String classInner, String expectAd
         if (firstStat instanceof PsiTryStatement) {
             PsiTryStatement tryStatement = (PsiTryStatement)initializer.getBody().getStatements()[0];
             assert tryStatement.getTryBlock() != null;
-            assert tryStatement.getTryBlock().getStatements().length == 1;
             expStat = (PsiExpressionStatement)tryStatement.getTryBlock().getStatements()[0];
         } else {
             expStat = (PsiExpressionStatement)firstStat;
         }
         assert expStat.getExpression() instanceof PsiMethodCallExpression;
 
-        Assert.assertEquals(expectAddStat, expStat.getExpression().getText());
+        String[] actual = new String[expectAddStats.length];
+        PsiElement currStat = expStat;
+        for (int i=0; i<expectAddStats.length && currStat != null; i++) {
+            actual[i] = ((PsiExpressionStatement)currStat).getExpression().getText();
+            do {
+                currStat = currStat.getNextSibling();
+            } while(i < expectAddStats.length - 1 && currStat != null && !(currStat instanceof PsiStatement));
+        }
+        Assert.assertArrayEquals(expectAddStats, actual);
     }
 
-    protected void testQuickFixEntityInLocalVariable(String methodInner, String expectAddStat, LocalQuickFix quickFix) {
+    protected void testQuickFixEntityInLocalVariable(String methodInner, String[] expectAddStats, LocalQuickFix quickFix) {
         Project project = myFixture.getProject();
         PsiMethod method = JavaPsiFacade.getElementFactory(project).createMethodFromText(
                 "public foo () {\n" +
@@ -98,15 +105,22 @@ protected void testQuickFixEntityInLocalVariable(String methodInner, String expe
                 new MockProblemDescriptor(rExp, "", ProblemHighlightType.GENERIC_ERROR_OR_WARNING);
         quickFix.applyFix(project, descriptor);
 
-        assert method.getBody().getStatements().length == 2;
         assert method.getBody().getStatements()[1] instanceof PsiExpressionStatement;
         PsiExpressionStatement expStat = (PsiExpressionStatement)method.getBody().getStatements()[1];
         assert expStat.getExpression() instanceof PsiMethodCallExpression;
 
-        Assert.assertEquals(expectAddStat, expStat.getExpression().getText());
+        String[] actual = new String[expectAddStats.length];
+        PsiElement currStat = expStat;
+        for (int i=0; i<expectAddStats.length && currStat != null; i++) {
+            actual[i] = ((PsiExpressionStatement)currStat).getExpression().getText();
+            do {
+                currStat = currStat.getNextSibling();
+            } while(i < expectAddStats.length - 1 && currStat != null && !(currStat instanceof PsiStatement));
+        }
+        Assert.assertArrayEquals(expectAddStats, actual);
     }
 
-    protected void testQuickFixEntityInMethodAssignment(String classInner, String expectAddStat, LocalQuickFix quickFix) {
+    protected void testQuickFixEntityInMethodAssignment(String classInner, String[] expectAddStats, LocalQuickFix quickFix) {
         Project project = myFixture.getProject();
         PsiClass aClass = JavaPsiFacade.getElementFactory(project).createClassFromText(
                 classInner, null);
@@ -130,11 +144,18 @@ protected void testQuickFixEntityInMethodAssignment(String classInner, String ex
                 new MockProblemDescriptor(rExp, "", ProblemHighlightType.GENERIC_ERROR_OR_WARNING);
         quickFix.applyFix(project, descriptor);
 
-        assert method.getBody().getStatements().length == 2;
         assert method.getBody().getStatements()[1] instanceof PsiExpressionStatement;
         PsiExpressionStatement expStat = (PsiExpressionStatement)method.getBody().getStatements()[1];
         assert expStat.getExpression() instanceof PsiMethodCallExpression;
 
-        Assert.assertEquals(expectAddStat, expStat.getExpression().getText());
+        String[] actual = new String[expectAddStats.length];
+        PsiElement currStat = expStat;
+        for (int i=0; i<expectAddStats.length && currStat != null; i++) {
+            actual[i] = ((PsiExpressionStatement)currStat).getExpression().getText();
+            do {
+                currStat = currStat.getNextSibling();
+            } while(i < expectAddStats.length - 1 && currStat != null && !(currStat instanceof PsiStatement));
+        }
+        Assert.assertArrayEquals(expectAddStats, actual);
     }
 }
diff --git a/src/test/java/com/immomo/momosec/lang/java/rule/momosecurity/XStreamUnserializeTest.java b/src/test/java/com/immomo/momosec/lang/java/rule/momosecurity/XStreamUnserializeTest.java
@@ -29,7 +29,7 @@ public void testIfFindAllVulns() {
     public void testLocalVariableFix() {
         testQuickFixEntityInLocalVariable(
                 "XStream x = new XStream(new DomDriver());",
-                "XStream.setupDefaultSecurity(x)",
+                new String[]{"XStream.setupDefaultSecurity(x)"},
                 new XStreamUnserialize.XStreamUnserializeQuickFix()
         );
     }
@@ -40,15 +40,15 @@ public void testMethodAssignmentFix() {
                 "public void foo() {\n" +
                 "  x = new XStream(new DomDriver());\n" +
                 "}",
-                "XStream.setupDefaultSecurity(x)",
+                new String[]{"XStream.setupDefaultSecurity(x)"},
                 new XStreamUnserialize.XStreamUnserializeQuickFix()
         );
     }
 
     public void testClassFieldFix() {
         testQuickFixInClassInitializer(
                 "static XStream x = new XStream(new DomDriver());",
-                "XStream.setupDefaultSecurity(x)",
+                new String[]{"XStream.setupDefaultSecurity(x)"},
                 new XStreamUnserialize.XStreamUnserializeQuickFix()
         );
     }