Activated defense against XML eXternal Entity (XXE) attacks

Author: vruusmann

pmml-model/src/main/java/org/jpmml/model/SAXUtil.java

@@ -31,13 +31,23 @@ private SAXUtil(){
 	static
 	public SAXSource createFilteredSource(InputStream is, XMLFilter... filters) throws SAXException {
 		XMLReader reader = XMLReaderFactory.createXMLReader();
+		reader.setFeature("http://apache.org/xml/features/disallow-doctype-decl", true);
+
+		reader = createFilteredReader(reader, filters);
+
+		return new SAXSource(reader, new InputSource(is));
+	}
+
+	static
+	public XMLReader createFilteredReader(XMLReader reader, XMLFilter... filters){
+		XMLReader result = reader;
 
 		for(XMLFilter filter : filters){
-			filter.setParent(reader);
+			filter.setParent(result);
 
-			reader = filter;
+			result = filter;
 		}
 
-		return new SAXSource(reader, new InputSource(is));
+		return result;
 	}
 }

pmml-model/src/test/java/org/jpmml/model/XXEAttackTest.java

@@ -8,17 +8,15 @@
 import java.util.List;
 
 import javax.xml.bind.UnmarshalException;
-import javax.xml.transform.sax.SAXSource;
+import javax.xml.transform.Source;
 import javax.xml.transform.stream.StreamSource;
 
-import org.dmg.pmml.Extension;
 import org.dmg.pmml.PMML;
 import org.junit.Test;
-import org.xml.sax.InputSource;
-import org.xml.sax.XMLReader;
-import org.xml.sax.helpers.XMLReaderFactory;
+import org.xml.sax.SAXParseException;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 public class XXEAttackTest {
@@ -30,30 +28,31 @@ public void unmarshal() throws Exception {
 		System.setProperty("javax.xml.accessExternalDTD", "file");
 
 		try(InputStream is = ResourceUtil.getStream(XXEAttackTest.class);){
-			pmml = JAXBUtil.unmarshalPMML(new StreamSource(is));
+			Source source = new StreamSource(is);
+
+			pmml = JAXBUtil.unmarshalPMML(source);
 		} finally {
 			System.clearProperty("javax.xml.accessExternalDTD");
 		}
 
-		List<Extension> extensions = pmml.getExtensions();
-		assertEquals(1, extensions.size());
+		List<?> content = ExtensionUtil.getContent(pmml);
 
-		Extension extension = extensions.get(0);
-		assertEquals(Arrays.asList("lol"), extension.getContent());
+		assertEquals(Arrays.asList("lol"), content);
 	}
 
 	@Test
 	public void unmarshalSecurely() throws Exception {
 
 		try(InputStream is = ResourceUtil.getStream(XXEAttackTest.class)){
-			XMLReader reader = XMLReaderFactory.createXMLReader();
-			reader.setFeature("http://apache.org/xml/features/disallow-doctype-decl", true);
+			Source source = SAXUtil.createFilteredSource(is);
 
-			JAXBUtil.unmarshalPMML(new SAXSource(reader, new InputSource(is)));
+			JAXBUtil.unmarshalPMML(source);
 
 			fail();
 		} catch(UnmarshalException ue){
-			// Ignored
+			Throwable cause = ue.getCause();
+
+			assertTrue(cause instanceof SAXParseException);
 		}
 	}
 }