/src/test/java/SaxParserBehaviourTest.java

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.StringReader;
import java.io.UnsupportedEncodingException;

import javax.xml.XMLConstants;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.parsers.SAXParser;
import javax.xml.parsers.SAXParserFactory;

import com.atlassian.security.xml.UntrustedXmlParserFactory;
import org.junit.Test;
import org.w3c.dom.Document;
import org.xml.sax.EntityResolver;
import org.xml.sax.InputSource;
import org.xml.sax.SAXException;
import org.xml.sax.SAXNotRecognizedException;
import org.xml.sax.SAXParseException;
import org.xml.sax.XMLReader;
import org.xml.sax.helpers.DefaultHandler;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;

/**
 * There are two ways to use SAX parsers, through {@link SAXParser} and through {@link XMLReader}.
 * In both cases, you must ensure the entity resolver is set to an implementation that throws IOException
 * for all resolution attempts.
 */
public class SaxParserBehaviourTest
{
    static SAXParser createSaxParser() throws ParserConfigurationException, SAXException
    {
        SAXParserFactory spf = SAXParserFactory.newInstance();
        spf.setFeature("http://apache.org/xml/features/nonvalidating/load-external-dtd", false);

        // This will turn any attempt to use a DTD into failure
//        spf.setFeature("http://apache.org/xml/features/disallow-doctype-decl", true);

        // Only necessary for bundled non-JDK Xerces
        spf.setFeature(XMLConstants.FEATURE_SECURE_PROCESSING, true);

        // Note that it's the EntityResolver in DefaultHandler that looks up entities.
        // You must throw an IOException in the resolver to prevent it hitting the specified path

        return spf.newSAXParser();
    }

    private static InputSource EMPTY_INPUT_SOURCE = new InputSource(new ByteArrayInputStream(new byte[0]));

    private static EntityResolver emptyEntityResolver = new EntityResolver()
    {
        public InputSource resolveEntity(String publicId, String systemId) throws SAXException, IOException
        {
            return EMPTY_INPUT_SOURCE;
        }
    };

    static XMLReader createXmlReader() throws ParserConfigurationException, SAXException
    {
        return UntrustedXmlParserFactory.newXmlReader();

//        XMLReader xr = createSaxParser().getXMLReader();
//        xr.setEntityResolver(emptyEntityResolver);
//        return xr;
    }

    static XMLReader createNamespaceAwareXmlReader() throws ParserConfigurationException, SAXException
    {
        return UntrustedXmlParserFactory.newNamespaceAwareXmlReader();
    }

    @Test(expected = SAXNotRecognizedException.class)
    public void askingForUnknownAttributesFails() throws Exception
    {
        SAXParserFactory spf = SAXParserFactory.newInstance();
        spf.setFeature("no-attribute-with-this-identifier", false);
    }

    static InputSource in(String content) throws UnsupportedEncodingException
    {
        return new InputSource(new ByteArrayInputStream(content.getBytes("us-ascii")));
    }

    @Test
    public void parseDocumentExpandsAmpersand() throws Exception
    {
        XMLReader xr = createXmlReader();

        StringGathererHandler sgh = new StringGathererHandler();

        xr.setContentHandler(sgh);
        xr.parse(in(SampleXmlDocuments.AMPERSAND_DOCUMENT));

        assertEquals("&", sgh.toString());
    }

    static class StringGathererHandler extends DefaultHandler
    {
        private StringBuilder sb = new StringBuilder();

        @Override
        public void characters(char[] ch, int start, int length) throws SAXException
        {
            sb.append(ch, start, length);
        }

        @Override
        public String toString()
        {
            return sb.toString();
        }

        @Override
        public InputSource resolveEntity(String publicId, String systemId) throws IOException, SAXException
        {
            throw new IOException();
        }
    }

    static class StringGathererContentHandler extends DefaultHandler
    {
        private StringBuilder sb = new StringBuilder();

        @Override
        public void characters(char[] ch, int start, int length) throws SAXException
        {
            sb.append(ch, start, length);
        }

        @Override
        public String toString()
        {
            return sb.toString();
        }
    }

    static class FailingEntityResolver implements EntityResolver
    {
        public InputSource resolveEntity(String publicId, String systemId) throws SAXException, IOException
        {
            throw new IOException();
        }
    }

    @Test(expected = SAXParseException.class, timeout = 1000)
    public void parseBillionLaughsDoesNotExhaustMemory() throws Exception
    {
        StringGathererHandler sgh = new StringGathererHandler();

        XMLReader xr = createXmlReader();
        xr.setContentHandler(sgh);
        xr.parse(in(SampleXmlDocuments.BILLION_LAUGHS));

        assertEquals("", sgh.toString());
    }

    @Test
    public void externalEntityIsNotRead() throws Exception
    {
        StringGathererHandler sgh = new StringGathererHandler();

        XMLReader xr = createXmlReader();
        xr.setContentHandler(sgh);
        xr.parse(in(SampleXmlDocuments.externalResourceEntity()));

        assertEquals("", sgh.toString());
    }

    @Test
    public void externalEntityIsNotIncludedInResultUsingXmlReader() throws Exception
    {
        StringGathererContentHandler sgh = new StringGathererContentHandler();

        XMLReader xr = createXmlReader();
        xr.setContentHandler(sgh);
        xr.parse(in(SampleXmlDocuments.externalResourceEntity()));

        assertEquals("", sgh.toString());
    }

    @Test
    public void externalEntityIsNotReadUsingXmlReader() throws Exception
    {
        HttpAttemptDetector detector = new HttpAttemptDetector();

        new Thread(detector).start();

        StringGathererContentHandler sgh = new StringGathererContentHandler();

        XMLReader xr = createXmlReader();
        xr.setContentHandler(sgh);
        xr.parse(in(SampleXmlDocuments.externalResourceEntity(detector.getUrl())));

        assertFalse(detector.wasAttempted());
    }

    @Test
    public void dtdUriPointsToFile() throws Exception
    {
        StringGathererHandler sgh = new StringGathererHandler();

        XMLReader xr = createXmlReader();
        xr.setContentHandler(sgh);
        xr.parse(in(SampleXmlDocuments.EXTERNAL_DTD));

        assertEquals("", sgh.toString());
    }

    @Test
    public void dtdUriPointsToUrl() throws Exception
    {
        HttpAttemptDetector detector = new HttpAttemptDetector();

        new Thread(detector).start();

        String s = SampleXmlDocuments.externalUrlDtd(detector.getUrl());

        StringGathererHandler sgh = new StringGathererHandler();

        XMLReader xr = createXmlReader();
        xr.setContentHandler(sgh);
        xr.parse(in(s));

        assertFalse("I don't want to see HTTP connection attempts", detector.wasAttempted());
        assertEquals("", sgh.toString());
    }
}