Fix more XXE cases

core-codemods/src/main/java/io/codemodder/codemods/codeql/CodeQLXXECodemod.java

@@ -7,7 +7,7 @@
 import io.codemodder.providers.sarif.codeql.ProvidedCodeQLScan;
 import io.codemodder.remediation.GenericRemediationMetadata;
 import io.codemodder.remediation.Remediator;
-import io.codemodder.remediation.xxe.XXEIntermediateXMLStreamReaderRemediator;
+import io.codemodder.remediation.xxe.XMLStreamReaderIntermediateRemediator;
 import java.util.Optional;
 import javax.inject.Inject;
 
@@ -24,7 +24,7 @@ public final class CodeQLXXECodemod extends CodeQLRemediationCodemod {
   @Inject
   public CodeQLXXECodemod(@ProvidedCodeQLScan(ruleId = "java/xxe") final RuleSarif sarif) {
     super(GenericRemediationMetadata.XXE.reporter(), sarif);
-    this.remediator = new XXEIntermediateXMLStreamReaderRemediator<>();
+    this.remediator = new XMLStreamReaderIntermediateRemediator<>();
   }
 
   @Override

framework/codemodder-base/src/main/java/io/codemodder/remediation/xxe/TransformerFactoryAtCreationFixStrategy.java

@@ -21,7 +21,7 @@
 import java.util.Optional;
 
 /**
- * Fix strategy for XXE vulnerabilities anchored to the TransformerParser newInstance() calls. Finds
+ * Fix strategy for XXE vulnerabilities anchored to the TransformerParser.newInstance() calls. Finds
  * the parser's declaration and add statements disabling external entities and features.
  */
 final class TransformerFactoryAtCreationFixStrategy extends MatchAndFixStrategy {
@@ -70,7 +70,7 @@ public SuccessOrReason fix(final CompilationUnit cu, final Node node) {
   }
 
   /**
-   * Matches against TransformerFactory.newInstance() calls
+   * Matches against XMLInputFactory.newInstance() calls.
    *
    * @param node
    * @return

framework/codemodder-base/src/main/java/io/codemodder/remediation/xxe/XMLInputFactoryAtNewInstanceFixStrategy.java

@@ -0,0 +1,58 @@
+package io.codemodder.remediation.xxe;
+
+import com.github.javaparser.ast.CompilationUnit;
+import com.github.javaparser.ast.Node;
+import com.github.javaparser.ast.body.VariableDeclarator;
+import com.github.javaparser.ast.expr.Expression;
+import com.github.javaparser.ast.expr.MethodCallExpr;
+import com.github.javaparser.ast.stmt.Statement;
+import io.codemodder.ast.ASTs;
+import io.codemodder.remediation.SuccessOrReason;
+import java.util.List;
+import java.util.Optional;
+
+final class XMLInputFactoryAtNewInstanceFixStrategy
+    extends io.codemodder.remediation.MatchAndFixStrategy {
+
+  /**
+   * Matches (XmlInputFactory | var xif) y = XMLInputFactory.newInstance(), where the node is the
+   * newDocumentBuilder call.
+   *
+   * @param node the node to match
+   * @return true if this is an assignment to an XMLInputFactory
+   */
+  @Override
+  public boolean match(final Node node) {
+    Optional<MethodCallExpr> method =
+        ASTs.isInitializedToType(node, "newInstance", List.of("var", "XMLInputFactory"));
+    if (method.isEmpty()) {
+      return false;
+    }
+    MethodCallExpr newInstance = method.get();
+    Optional<Expression> scope = newInstance.getScope();
+    if (scope.isEmpty()) {
+      return false;
+    }
+    return scope.get().toString().equals("XMLInputFactory");
+  }
+
+  @Override
+  public SuccessOrReason fix(final CompilationUnit cu, final Node node) {
+    MethodCallExpr newInstance = (MethodCallExpr) node;
+    Optional<VariableDeclarator> initExpr = ASTs.isInitExpr(newInstance);
+    if (initExpr.isEmpty()) {
+      return SuccessOrReason.reason("Not an initialization expression");
+    }
+    VariableDeclarator xmlInputFactoryVariable = initExpr.get();
+    Optional<Statement> variableDeclarationStmtRef =
+        xmlInputFactoryVariable.findAncestor(Statement.class);
+
+    if (variableDeclarationStmtRef.isEmpty()) {
+      return SuccessOrReason.reason("Not assigned as part of statement");
+    }
+
+    Statement stmt = variableDeclarationStmtRef.get();
+    return XMLFixBuilder.addXMLInputFactoryDisablingStatement(
+        xmlInputFactoryVariable.getNameAsExpression(), stmt, false);
+  }
+}

framework/codemodder-base/src/main/java/io/codemodder/remediation/xxe/XXEIntermediateXMLStreamReaderFixStrategy.java → framework/codemodder-base/src/main/java/io/codemodder/remediation/xxe/XMLStreamReaderIntermediateFixStrategy.java

@@ -15,7 +15,7 @@
  * Fix strategy for XXE vulnerabilities anchored to arguments of a XMLStreamReader calls. Finds the
  * parser's declaration and add statements disabling external entities and features.
  */
-final class XXEIntermediateXMLStreamReaderFixStrategy implements RemediationStrategy {
+final class XMLStreamReaderIntermediateFixStrategy implements RemediationStrategy {
 
   @Override
   public SuccessOrReason fix(final CompilationUnit cu, final Node node) {
@@ -27,7 +27,7 @@ public SuccessOrReason fix(final CompilationUnit cu, final Node node) {
     if (maybeCall.isEmpty()) {
       return SuccessOrReason.reason("Not a method call");
     }
-    // get the xmlstreamreader scope variable
+    // get the XMLStreamReader scope variable
     MethodCallExpr createXMLStreamReaderCall = maybeCall.get();
     Expression xmlStreamReaderScope = createXMLStreamReaderCall.getScope().get();
 

framework/codemodder-base/src/main/java/io/codemodder/remediation/xxe/XXEIntermediateXMLStreamReaderRemediator.java → framework/codemodder-base/src/main/java/io/codemodder/remediation/xxe/XMLStreamReaderIntermediateRemediator.java

@@ -13,11 +13,12 @@
 import java.util.Optional;
 import java.util.function.Function;
 
-public class XXEIntermediateXMLStreamReaderRemediator<T> implements Remediator<T> {
+/** Remediator for XXE vulnerabilities anchored to the XMLStreamReader calls. */
+public final class XMLStreamReaderIntermediateRemediator<T> implements Remediator<T> {
 
   private final SearcherStrategyRemediator<T> searchStrategyRemediator;
 
-  public XXEIntermediateXMLStreamReaderRemediator() {
+  public XMLStreamReaderIntermediateRemediator() {
     this.searchStrategyRemediator =
         new SearcherStrategyRemediator.Builder<T>()
             .withSearcherStrategyPair(
@@ -30,20 +31,20 @@ public XXEIntermediateXMLStreamReaderRemediator() {
                                 .filter(Node::hasScope)
                                 .isPresent())
                     .build(),
-                new XXEIntermediateXMLStreamReaderFixStrategy())
+                new XMLStreamReaderIntermediateFixStrategy())
             .build();
   }
 
   @Override
   public CodemodFileScanningResult remediateAll(
-      CompilationUnit cu,
-      String path,
-      DetectorRule detectorRule,
-      Collection<T> findingsForPath,
-      Function<T, String> findingIdExtractor,
-      Function<T, Integer> findingStartLineExtractor,
-      Function<T, Optional<Integer>> findingEndLineExtractor,
-      Function<T, Optional<Integer>> findingColumnExtractor) {
+      final CompilationUnit cu,
+      final String path,
+      final DetectorRule detectorRule,
+      final Collection<T> findingsForPath,
+      final Function<T, String> findingIdExtractor,
+      final Function<T, Integer> findingStartLineExtractor,
+      final Function<T, Optional<Integer>> findingEndLineExtractor,
+      final Function<T, Optional<Integer>> findingColumnExtractor) {
     return searchStrategyRemediator.remediateAll(
         cu,
         path,

framework/codemodder-base/src/main/java/io/codemodder/remediation/xxe/XXERemediator.java

@@ -29,6 +29,8 @@ public XXERemediator(final NodePositionMatcher matcher) {
             .withMatchAndFixStrategyAndNodeMatcher(
                 new TransformerFactoryAtCreationFixStrategy(), matcher)
             .withMatchAndFixStrategyAndNodeMatcher(new XMLReaderAtParseFixStrategy(), matcher)
+            .withMatchAndFixStrategyAndNodeMatcher(
+                new XMLInputFactoryAtNewInstanceFixStrategy(), matcher)
             .build();
   }
 

framework/codemodder-base/src/test/java/io/codemodder/remediation/xxe/DefaultXXEIntermediateXMLStreamReaderRemediatorTest.java → framework/codemodder-base/src/test/java/io/codemodder/remediation/xxe/DefaultXMLStreamReaderIntermediateRemediatorTest.java

@@ -14,14 +14,14 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
-final class DefaultXXEIntermediateXMLStreamReaderRemediatorTest {
+final class DefaultXMLStreamReaderIntermediateRemediatorTest {
 
-  private XXEIntermediateXMLStreamReaderRemediator<Object> remediator;
+  private XMLStreamReaderIntermediateRemediator<Object> remediator;
   private DetectorRule rule;
 
   @BeforeEach
   void setup() {
-    remediator = new XXEIntermediateXMLStreamReaderRemediator<>();
+    remediator = new XMLStreamReaderIntermediateRemediator<>();
     rule = new DetectorRule("xxe", "XXE Fixed At XMLStreamReader", null);
   }