Remove explicit setEntityExpansion calls (#394)

When fixing XXE, users may find it helpful to also remove explicit turning on off entity expansion.
pixee · Jun 21, 2024 · ec403a7 · ec403a7
1 parent adb461a
commit ec403a7
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 9 deletions.
diff --git a/core-codemods/src/test/resources/sonar-xxe-s2755/Test.java.after b/core-codemods/src/test/resources/sonar-xxe-s2755/Test.java.after
@@ -72,7 +72,6 @@ public class XXEVuln {
         DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
         dbf.setFeature("http://xml.org/sax/features/external-general-entities", false);
         dbf.setFeature("http://xml.org/sax/features/external-parameter-entities", false);
-        dbf.setExpandEntityReferences(true);
         DocumentBuilder db = dbf.newDocumentBuilder();
         return db.parse(new InputSource(new StringReader(xml)));
     }

diff --git a/...java/io/codemodder/remediation/xxe/DocumentBuilderFactoryAndSAXParserAtCreationFixer.java b/...java/io/codemodder/remediation/xxe/DocumentBuilderFactoryAndSAXParserAtCreationFixer.java
@@ -44,6 +44,6 @@ public XXEFixAttempt tryFix(final int line, final Integer column, CompilationUni
 
     Statement statement = variableDeclarationStmtRef.get();
     return addFeatureDisablingStatements(
-        cu, newFactoryVariable.getNameAsExpression(), statement, false);
+        newFactoryVariable.getNameAsExpression(), statement, false);
   }
 }
diff --git a/...-base/src/main/java/io/codemodder/remediation/xxe/DocumentBuilderFactoryAtParseFixer.java b/...-base/src/main/java/io/codemodder/remediation/xxe/DocumentBuilderFactoryAtParseFixer.java
@@ -104,7 +104,7 @@ public XXEFixAttempt tryFix(final int line, final Integer column, CompilationUni
             "DocumentBuilder was initialized with a factory call without a statement");
       }
       return XMLFeatures.addFeatureDisablingStatements(
-          cu, factoryNameExpr, newDocumentBuilderStatement.get(), true);
+          factoryNameExpr, newDocumentBuilderStatement.get(), true);
     } else if (parserAssignmentNode instanceof Parameter) {
       return new XXEFixAttempt(true, false, "DocumentBuilder came from outside the method scope");
     }

diff --git a/framework/codemodder-base/src/main/java/io/codemodder/remediation/xxe/XMLFeatures.java b/framework/codemodder-base/src/main/java/io/codemodder/remediation/xxe/XMLFeatures.java
@@ -1,6 +1,6 @@
 package io.codemodder.remediation.xxe;
 
-import com.github.javaparser.ast.CompilationUnit;
+import com.github.javaparser.ast.Node;
 import com.github.javaparser.ast.NodeList;
 import com.github.javaparser.ast.expr.BooleanLiteralExpr;
 import com.github.javaparser.ast.expr.MethodCallExpr;
@@ -39,10 +39,7 @@ static MethodCallExpr createGeneralEntityDisablingCall(final NameExpr nameExpr)
   }
 
   static XXEFixAttempt addFeatureDisablingStatements(
-      final CompilationUnit cu,
-      final NameExpr variable,
-      final Statement statementToInjectAround,
-      final boolean before) {
+      final NameExpr variable, final Statement statementToInjectAround, final boolean before) {
     Optional<BlockStmt> block = ASTs.findBlockStatementFrom(statementToInjectAround);
     if (block.isEmpty()) {
       return new XXEFixAttempt(true, false, "No block statement found for newFactory() call");
@@ -62,6 +59,22 @@ static XXEFixAttempt addFeatureDisablingStatements(
       index++;
     }
     existingStatements.addAll(index, fixStatements);
+
+    // search for any dbf.setExpandEntityReferences(true) calls and remove them
+    blockStmt.findAll(MethodCallExpr.class).stream()
+        .filter(
+            m ->
+                "setExpandEntityReferences".equals(m.getNameAsString())
+                    && m.getScope().isPresent()
+                    && m.getScope().get().equals(variable)
+                    && m.getArguments().size() == 1
+                    && m.getArguments().get(0).isBooleanLiteralExpr()
+                    && m.getArguments().get(0).asBooleanLiteralExpr().getValue())
+        .map(Node::getParentNode)
+        .filter(Optional::isPresent)
+        .map(Optional::get)
+        .forEach(Node::remove);
+
     return new XXEFixAttempt(true, true, null);
   }
 }
diff --git a/...rk/codemodder-base/src/main/java/io/codemodder/remediation/xxe/XMLReaderAtParseFixer.java b/...rk/codemodder-base/src/main/java/io/codemodder/remediation/xxe/XMLReaderAtParseFixer.java
@@ -72,6 +72,6 @@ public XXEFixAttempt tryFix(final int line, final Integer column, final Compilat
       return new XXEFixAttempt(true, false, "No statement found for parse() call");
     }
     return XMLFeatures.addFeatureDisablingStatements(
-        cu, parser.asNameExpr(), parseStatement.get(), true);
+        parser.asNameExpr(), parseStatement.get(), true);
   }
 }