chore: improve XXE attack prevention (#2008)

Reported by Aikido earlier today.

Author: triceo

benchmark/src/main/java/ai/timefold/solver/benchmark/impl/io/PlannerBenchmarkConfigIO.java

@@ -28,12 +28,12 @@ public PlannerBenchmarkConfig read(Reader reader) {
              * the solver element in benchmark configuration and thus the solver element's namespace needs to be overridden.
              */
             return genericJaxbIO.readOverridingNamespace(document,
-                    ElementNamespaceOverride.of(SolverConfig.XML_ELEMENT_NAME, SolverConfig.XML_NAMESPACE));
+                    new ElementNamespaceOverride(SolverConfig.XML_ELEMENT_NAME, SolverConfig.XML_NAMESPACE));
         } else if (rootElementNamespace == null || rootElementNamespace.isEmpty()) {
             // If not, add the missing namespace to maintain backward compatibility.
             return genericJaxbIO.readOverridingNamespace(document,
-                    ElementNamespaceOverride.of(PlannerBenchmarkConfig.XML_ELEMENT_NAME, PlannerBenchmarkConfig.XML_NAMESPACE),
-                    ElementNamespaceOverride.of(SolverConfig.XML_ELEMENT_NAME, SolverConfig.XML_NAMESPACE));
+                    new ElementNamespaceOverride(PlannerBenchmarkConfig.XML_ELEMENT_NAME, PlannerBenchmarkConfig.XML_NAMESPACE),
+                    new ElementNamespaceOverride(SolverConfig.XML_ELEMENT_NAME, SolverConfig.XML_NAMESPACE));
         } else { // If there is an unexpected namespace, fail fast.
             String errorMessage = String.format("The <%s/> element belongs to a different namespace (%s) than expected (%s).",
                     PlannerBenchmarkConfig.XML_ELEMENT_NAME, rootElementNamespace, PlannerBenchmarkConfig.XML_NAMESPACE);

benchmark/src/main/java/ai/timefold/solver/benchmark/impl/result/BenchmarkResultIO.java

@@ -92,7 +92,7 @@ protected PlannerBenchmarkResult readPlannerBenchmarkResult(File plannerBenchmar
 
     protected PlannerBenchmarkResult read(Reader reader) {
         return genericJaxbIO.readOverridingNamespace(reader,
-                ElementNamespaceOverride.of(SOLVER_CONFIG_XML_ELEMENT_NAME, SolverConfig.XML_NAMESPACE));
+                new ElementNamespaceOverride(SOLVER_CONFIG_XML_ELEMENT_NAME, SolverConfig.XML_NAMESPACE));
     }
 
     protected void write(PlannerBenchmarkResult plannerBenchmarkResult, Writer writer) {

benchmark/src/main/java/ai/timefold/solver/benchmark/impl/xsd/XsdAggregator.java

@@ -8,52 +8,52 @@
 import java.util.function.Predicate;
 import java.util.function.UnaryOperator;
 
-import javax.xml.XMLConstants;
 import javax.xml.parsers.DocumentBuilder;
 import javax.xml.parsers.DocumentBuilderFactory;
 import javax.xml.parsers.ParserConfigurationException;
 import javax.xml.transform.OutputKeys;
-import javax.xml.transform.Result;
 import javax.xml.transform.Transformer;
 import javax.xml.transform.TransformerConfigurationException;
 import javax.xml.transform.TransformerException;
-import javax.xml.transform.TransformerFactory;
 import javax.xml.transform.dom.DOMSource;
 import javax.xml.transform.stream.StreamResult;
 
 import ai.timefold.solver.core.config.solver.SolverConfig;
+import ai.timefold.solver.core.impl.io.jaxb.GenericJaxbIO;
 
+import org.jspecify.annotations.NullMarked;
 import org.w3c.dom.Attr;
 import org.w3c.dom.Document;
 import org.w3c.dom.Element;
 import org.w3c.dom.Node;
-import org.w3c.dom.NodeList;
 import org.xml.sax.SAXException;
 
 /**
  * This class merges solver.xsd and benchmark.xsd into a single XML Schema file that contains both Solver and Benchmark XML
  * types under a single namespace of the benchmark.xsd.
- *
+ * <p>
  * Both solver.xsd and benchmark.xsd declare its own namespace as they are supposed to be used for different purposes. As the
  * benchmark configuration contains solver configuration, the benchmark.xsd imports the solver.xsd. To avoid distributing
  * dependent schemas and using prefixes in users' XML configuration files, the types defined by solver.xsd are merged to
  * the benchmark.xsd under its namespace.
  */
+@NullMarked
 public final class XsdAggregator {
 
     private static final String TNS_PREFIX = "tns";
 
     public static void main(String[] args) {
         if (args.length != 3) {
-            String msg = "The XSD Aggregator expects 3 arguments:\n"
-                    + "1) a path to the solver XSD file. \n"
-                    + "2) a path to the benchmark XSD file. \n"
-                    + "3) a path to an output file where the merged benchmark XSD should be saved to.";
+            var msg = """
+                    The XSD Aggregator expects 3 arguments:
+                    1) a path to the solver XSD file.
+                    2) a path to the benchmark XSD file.
+                    3) a path to an output file where the merged benchmark XSD should be saved to.""";
             throw new IllegalArgumentException(msg);
         }
-        File solverXsd = checkFileExists(new File(args[0]));
-        File benchmarkXsd = checkFileExists(new File(args[1]));
-        File outputXsd = new File(args[2]);
+        var solverXsd = checkFileExists(new File(args[0]));
+        var benchmarkXsd = checkFileExists(new File(args[1]));
+        var outputXsd = new File(args[2]);
 
         if (!outputXsd.getParentFile().exists()) {
             outputXsd.getParentFile().mkdirs();
@@ -71,53 +71,46 @@ private static File checkFileExists(File file) {
     }
 
     private void mergeXmlSchemas(File solverSchemaFile, File benchmarkSchemaFile, File outputSchemaFile) {
-        DocumentBuilderFactory factory = DocumentBuilderFactory.newInstance();
-        Document solverSchema = parseXml(solverSchemaFile, factory);
-        Element solverRootElement = solverSchema.getDocumentElement();
-        Document benchmarkSchema = parseXml(benchmarkSchemaFile, factory);
+        var factory = GenericJaxbIO.createDocumentBuilderFactory();
+        var solverSchema = parseXml(solverSchemaFile, factory);
+        var solverRootElement = solverSchema.getDocumentElement();
+        var benchmarkSchema = parseXml(benchmarkSchemaFile, factory);
 
         removeReferencesToSolverConfig(benchmarkSchema, benchmarkSchemaFile);
 
         copySolverConfigTypes(benchmarkSchema, solverRootElement);
 
-        Transformer transformer = createTransformer();
-        DOMSource source = new DOMSource(benchmarkSchema);
-        Result result = new StreamResult(outputSchemaFile);
+        var source = new DOMSource(benchmarkSchema);
+        var result = new StreamResult(outputSchemaFile);
         try {
-            transformer.transform(source, result);
+            createTransformer().transform(source, result);
         } catch (TransformerException e) {
-            throw new IllegalArgumentException(
-                    "Failed to write the resulting XSD to a file (" + outputSchemaFile.getAbsolutePath() + ").", e);
+            throw new IllegalArgumentException("Failed to write the resulting XSD to a file (%s)."
+                    .formatted(outputSchemaFile.getAbsolutePath()), e);
         }
     }
 
     private Document parseXml(File xmlFile, DocumentBuilderFactory documentBuilderFactory) {
-        DocumentBuilder builder;
         try {
-            builder = documentBuilderFactory.newDocumentBuilder();
+            return documentBuilderFactory.newDocumentBuilder().parse(xmlFile);
         } catch (ParserConfigurationException e) {
-            throw new IllegalArgumentException("Failed to create a " + DocumentBuilder.class.getName() + "instance.", e);
-        }
-
-        try {
-            return builder.parse(xmlFile);
-        } catch (SAXException saxException) {
-            throw new IllegalArgumentException("Failed to parse an XML file (" + xmlFile.getAbsolutePath() + ").",
-                    saxException);
-        } catch (IOException ioException) {
-            throw new IllegalArgumentException("Failed to open an XML file (" + xmlFile.getAbsolutePath() + ").", ioException);
+            throw new IllegalArgumentException("Failed to create a %s instance."
+                    .formatted(DocumentBuilder.class.getSimpleName()), e);
+        } catch (SAXException | IOException exception) {
+            throw new IllegalArgumentException("Failed to parse an XML file (%s)."
+                    .formatted(xmlFile.getAbsolutePath()), exception);
         }
     }
 
     private void removeReferencesToSolverConfig(Document benchmarkSchema, File benchmarkSchemaFile) {
-        boolean solverNamespaceRemoved = false;
-        boolean solverElementRefRemoved = false;
-        boolean importRemoved = false;
+        var solverNamespaceRemoved = false;
+        var solverElementRefRemoved = false;
+        var importRemoved = false;
 
-        NodeList nodeList = benchmarkSchema.getElementsByTagName("*");
-        for (int i = 0; i < nodeList.getLength(); i++) {
-            Node node = nodeList.item(i);
-            Element element = (Element) node;
+        var nodeList = benchmarkSchema.getElementsByTagName("*");
+        for (var i = 0; i < nodeList.getLength(); i++) {
+            var node = Objects.requireNonNull(nodeList.item(i));
+            var element = (Element) node;
 
             if ("xs:schema".equals(node.getNodeName())) { // Remove the solver namespace declaration.
                 element.removeAttribute("xmlns:" + SOLVER_NAMESPACE_PREFIX);
@@ -148,31 +141,30 @@ private void removeReferencesToSolverConfig(Document benchmarkSchema, File bench
          * a successful validation by the resulting XML schema.
          */
         if (!solverElementRefRemoved) {
-            String msg = String.format("An expected reference to the solver element was not found. Check the content of (%s).",
+            var msg = String.format("An expected reference to the solver element was not found. Check the content of (%s).",
                     benchmarkSchemaFile);
             throw new AssertionError(msg);
         }
 
         if (!solverNamespaceRemoved) {
-            String msg = String.format("An expected namespace (%s) declaration was not found. Check the content of (%s).",
+            var msg = String.format("An expected namespace (%s) declaration was not found. Check the content of (%s).",
                     SolverConfig.XML_NAMESPACE, benchmarkSchemaFile);
             throw new AssertionError(msg);
         }
 
         if (!importRemoved) {
-            String msg = String.format("An expected import element was not found. Check the content of (%s).",
-                    benchmarkSchemaFile);
+            var msg =
+                    String.format("An expected import element was not found. Check the content of (%s).", benchmarkSchemaFile);
             throw new AssertionError(msg);
         }
     }
 
     private void copySolverConfigTypes(Document benchmarkSchema, Element solverSchemaRoot) {
-        Element benchmarkSchemaRoot = benchmarkSchema.getDocumentElement();
-        NodeList solverChildNodes = solverSchemaRoot.getChildNodes();
-        for (int i = 0; i < solverChildNodes.getLength(); i++) {
-            Node node = solverChildNodes.item(i);
-            boolean isSolverElementDeclaration =
-                    isXsElement(node) && hasAttribute(node, "name", SolverConfig.XML_ELEMENT_NAME);
+        var benchmarkSchemaRoot = benchmarkSchema.getDocumentElement();
+        var solverChildNodes = solverSchemaRoot.getChildNodes();
+        for (var i = 0; i < solverChildNodes.getLength(); i++) {
+            var node = solverChildNodes.item(i);
+            var isSolverElementDeclaration = isXsElement(node) && hasAttribute(node, "name", SolverConfig.XML_ELEMENT_NAME);
             if (!isSolverElementDeclaration) { // Skip the solver root element.
                 benchmarkSchemaRoot.appendChild(benchmarkSchema.importNode(node, true));
             }
@@ -184,41 +176,31 @@ private boolean isXsElement(Node node) {
     }
 
     private boolean hasAttribute(Node node, String attributeName, String attributeValue) {
-        Objects.requireNonNull(node);
-        Objects.requireNonNull(attributeName);
-        Objects.requireNonNull(attributeValue);
-
-        Attr attribute = ((Element) node).getAttributeNode(attributeName);
+        var attribute = ((Element) node).getAttributeNode(attributeName);
         return (attribute != null && attributeValue.equals(attribute.getValue()));
     }
 
     private void updateNodeAttributes(Node node, Predicate<Attr> attributePredicate, UnaryOperator<String> valueFunction) {
-        Objects.requireNonNull(node);
-        Objects.requireNonNull(attributePredicate);
-        Objects.requireNonNull(valueFunction);
-        for (int i = 0; i < node.getAttributes().getLength(); i++) {
-            Attr attribute = (Attr) node.getAttributes().item(i);
+        for (var i = 0; i < node.getAttributes().getLength(); i++) {
+            var attribute = (Attr) node.getAttributes().item(i);
             if (attributePredicate.test(attribute)) {
                 attribute.setValue(valueFunction.apply(attribute.getValue()));
             }
         }
     }
 
     private Transformer createTransformer() {
-        TransformerFactory transformerFactory = TransformerFactory.newInstance();
-        // Protect the Transformer from XXE attacks.
-        transformerFactory.setAttribute(XMLConstants.ACCESS_EXTERNAL_DTD, "");
-        transformerFactory.setAttribute(XMLConstants.ACCESS_EXTERNAL_STYLESHEET, "");
-        Transformer transformer;
+        var transformerFactory = GenericJaxbIO.createTransformerFactory();
         try {
-            transformer = transformerFactory.newTransformer();
+            var transformer = transformerFactory.newTransformer();
+            transformer.setOutputProperty(OutputKeys.INDENT, "yes");
+            transformer.setOutputProperty(OutputKeys.STANDALONE, "yes");
+            transformer.setOutputProperty(OutputKeys.ENCODING, "utf-8");
+            transformer.setOutputProperty("{http://xml.apache.org/xslt}indent-amount", "2");
+            return transformer;
         } catch (TransformerConfigurationException e) {
-            throw new IllegalArgumentException("Failed to create a " + Transformer.class.getName() + ".", e);
+            throw new IllegalArgumentException("Failed to create a %s."
+                    .formatted(Transformer.class.getSimpleName()), e);
         }
-        transformer.setOutputProperty(OutputKeys.INDENT, "yes");
-        transformer.setOutputProperty(OutputKeys.STANDALONE, "yes");
-        transformer.setOutputProperty(OutputKeys.ENCODING, "utf-8");
-        transformer.setOutputProperty("{http://xml.apache.org/xslt}indent-amount", "2");
-        return transformer;
     }
 }

benchmark/src/test/java/ai/timefold/solver/benchmark/config/PlannerBenchmarkConfigTest.java

@@ -16,7 +16,6 @@
 import ai.timefold.solver.core.testdomain.TestdataSolution;
 import ai.timefold.solver.jackson.impl.domain.solution.JacksonSolutionFileIO;
 import ai.timefold.solver.persistence.common.api.domain.solution.RigidTestdataSolutionFileIO;
-import ai.timefold.solver.persistence.common.api.domain.solution.SolutionFileIO;
 
 import org.apache.commons.io.IOUtils;
 import org.junit.jupiter.api.Test;
@@ -32,7 +31,7 @@ class PlannerBenchmarkConfigTest {
     @ParameterizedTest
     @ValueSource(strings = { TEST_PLANNER_BENCHMARK_CONFIG_WITHOUT_NAMESPACE, TEST_PLANNER_BENCHMARK_CONFIG_WITH_NAMESPACE })
     void xmlConfigFileRemainsSameAfterReadWrite(String xmlBenchmarkConfigResource) throws IOException {
-        PlannerBenchmarkConfigIO xmlIO = new PlannerBenchmarkConfigIO();
+        var xmlIO = new PlannerBenchmarkConfigIO();
         PlannerBenchmarkConfig jaxbBenchmarkConfig;
 
         try (Reader reader = new InputStreamReader(
@@ -44,13 +43,13 @@ void xmlConfigFileRemainsSameAfterReadWrite(String xmlBenchmarkConfigResource) t
 
         Writer stringWriter = new StringWriter();
         xmlIO.write(jaxbBenchmarkConfig, stringWriter);
-        String jaxbString = stringWriter.toString();
+        var jaxbString = stringWriter.toString();
 
-        String originalXml = IOUtils.toString(PlannerBenchmarkConfigTest.class.getResourceAsStream(xmlBenchmarkConfigResource),
+        var originalXml = IOUtils.toString(PlannerBenchmarkConfigTest.class.getResourceAsStream(xmlBenchmarkConfigResource),
                 StandardCharsets.UTF_8);
 
         // During writing the benchmark config, the benchmark element's namespace is removed.
-        String benchmarkElementWithNamespace =
+        var benchmarkElementWithNamespace =
                 PlannerBenchmarkConfig.XML_ELEMENT_NAME + " xmlns=\"" + PlannerBenchmarkConfig.XML_NAMESPACE + "\"";
         if (originalXml.contains(benchmarkElementWithNamespace)) {
             originalXml = originalXml.replace(benchmarkElementWithNamespace, PlannerBenchmarkConfig.XML_ELEMENT_NAME);
@@ -60,8 +59,8 @@ void xmlConfigFileRemainsSameAfterReadWrite(String xmlBenchmarkConfigResource) t
 
     @Test
     void readAndValidateInvalidBenchmarkConfig_failsIndicatingTheIssue() {
-        PlannerBenchmarkConfigIO xmlIO = new PlannerBenchmarkConfigIO();
-        String benchmarkConfigXml = "<plannerBenchmark xmlns=\"https://timefold.ai/xsd/benchmark\">\n"
+        var xmlIO = new PlannerBenchmarkConfigIO();
+        var benchmarkConfigXml = "<plannerBenchmark xmlns=\"https://timefold.ai/xsd/benchmark\">\n"
                 + "  <benchmarkDirectory>data</benchmarkDirectory>\n"
                 + "  <parallelBenchmarkCount>AUTO</parallelBenchmarkCount>\n"
                 + "  <solverBenchmark>\n"
@@ -78,19 +77,20 @@ void readAndValidateInvalidBenchmarkConfig_failsIndicatingTheIssue() {
                 + "  </solverBenchmark>\n"
                 + "</plannerBenchmark>\n";
 
-        StringReader stringReader = new StringReader(benchmarkConfigXml);
+        var stringReader = new StringReader(benchmarkConfigXml);
         assertThatExceptionOfType(TimefoldXmlSerializationException.class)
                 .isThrownBy(() -> xmlIO.read(stringReader))
-                .withRootCauseExactlyInstanceOf(SAXParseException.class)
+                .havingRootCause()
+                .isInstanceOf(SAXParseException.class)
                 .withMessageContaining("solutionKlazz");
     }
 
     @Test
     public void assignCustomSolutionIO() {
-        ProblemBenchmarksConfig pbc = new ProblemBenchmarksConfig();
+        var pbc = new ProblemBenchmarksConfig();
         pbc.setSolutionFileIOClass(RigidTestdataSolutionFileIO.class);
 
-        Class<? extends SolutionFileIO<?>> configured = pbc.getSolutionFileIOClass();
+        var configured = pbc.getSolutionFileIOClass();
         assertThat(configured).isNotNull();
     }
 

core/src/main/java/ai/timefold/solver/core/impl/io/jaxb/ElementNamespaceOverride.java

@@ -1,32 +1,5 @@
 package ai.timefold.solver.core.impl.io.jaxb;
 
-public class ElementNamespaceOverride {
+public record ElementNamespaceOverride(String elementLocalName, String namespaceOverride) {
 
-    public static ElementNamespaceOverride of(String elementLocalName, String namespaceOverride) {
-        return new ElementNamespaceOverride(elementLocalName, namespaceOverride);
-    }
-
-    private final String elementLocalName;
-    private final String namespaceOverride;
-
-    private ElementNamespaceOverride(String elementLocalName, String namespaceOverride) {
-        this.elementLocalName = elementLocalName;
-        this.namespaceOverride = namespaceOverride;
-    }
-
-    public String getElementLocalName() {
-        return elementLocalName;
-    }
-
-    public String getNamespaceOverride() {
-        return namespaceOverride;
-    }
-
-    @Override
-    public String toString() {
-        return "ElementNamespaceOverride{" +
-                "elementLocalName='" + elementLocalName + '\'' +
-                ", namespaceOverride='" + namespaceOverride + '\'' +
-                '}';
-    }
 }