/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.commons.xml.secure;

import static org.apache.commons.xml.secure.AttackTestSupport.assertParseFails;
import static org.apache.commons.xml.secure.AttackTestSupport.assertParseSucceeds;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertSame;
import static org.junit.jupiter.api.Assertions.assertTrue;

import java.io.IOException;
import java.io.StringWriter;
import java.net.URL;

import javax.xml.XMLConstants;
import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.SAXParser;
import javax.xml.parsers.SAXParserFactory;
import javax.xml.stream.XMLInputFactory;
import javax.xml.stream.XMLResolver;
import javax.xml.stream.XMLStreamException;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.URIResolver;
import javax.xml.transform.stream.StreamResult;
import javax.xml.validation.SchemaFactory;

import org.junit.jupiter.api.Assumptions;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import org.w3c.dom.Document;
import org.w3c.dom.bootstrap.DOMImplementationRegistry;
import org.w3c.dom.ls.DOMImplementationLS;
import org.w3c.dom.ls.LSInput;
import org.w3c.dom.ls.LSResourceResolver;
import org.xml.sax.EntityResolver;
import org.xml.sax.InputSource;
import org.xml.sax.SAXException;
import org.xml.sax.XMLReader;

/**
 * Tests that a caller-supplied resolver cannot remove the secure ignore-all floor on any factory.
 *
 * <p>The observable contract on every secure factory is the same: a resource the caller resolves (returns a non-null value) is allowed, but anything the
 * caller does not resolve is resolved to empty content instead of fetched, so a resolver that resolves nothing leaves the block in place. Most
 * factories enforce this with a {@link FallbackIgnoreEntityResolver2}-style floor that consults the caller and returns empty on a {@code null} return; Saxon
 * enforces the equivalent through an ignore-all {@code ResourceResolver} floor on its {@code Configuration}. Every resolver channel is exercised: the SAX/DOM
 * {@link EntityResolver}, the StAX {@link XMLResolver}, the schema {@link LSResourceResolver} and the XSLT {@link URIResolver}.</p>
 */
class EntityResolverFloorTest {

    /** systemId the allow-list resolvers permit (its content carries {@link AttackTestSupport#LEAKED_MARKER}). */
    private static final String ALLOWED = AttackTestSupport.resourceUrl("referenced.txt").toString();

    /** systemId the allow-list resolvers do not resolve (so the floor resolves it to empty; its content carries {@link AttackTestSupport#LEAKED_MARKER}). */
    private static final String UNLISTED = AttackTestSupport.resourceUrl("referenced.xml").toString();

    /** Resolves only {@link #ALLOWED}; returns {@code null} for anything else. */
    private static final EntityResolver ENTITY_ALLOW_LIST = (publicId, systemId) ->
            ALLOWED.equals(systemId) ? new InputSource(new URL(systemId).openStream()) : null;

    /**
     * Allow-all resolver: it denies nothing, resolving whatever {@code systemId} it is handed by opening it as a URL. It nonetheless cannot resolve a bare
     * relative reference such as {@code referenced.xml}, because a plain {@link EntityResolver} (unlike {@link org.xml.sax.ext.EntityResolver2}) is given no
     * base URI and the SAX 2 contract promises it an already-absolutized {@code systemId}. So the resolution fails not from any deny decision but because the
     * resolver was never handed the whole URL: it succeeds only if the floor absolutizes the XInclude href against the base before consulting the caller.
     */
    private static final EntityResolver RESOLVE_ALL = (publicId, systemId) -> {
        final InputSource source = new InputSource(new URL(systemId).openStream());
        source.setSystemId(systemId);
        return source;
    };

    /** Absolute URL of the host document whose {@code xi:include} references {@code referenced.xml} by a relative href. */
    private static final String XINCLUDE_HOST = AttackTestSupport.resourceUrl("with-xinclude.xml").toString();

    /** Resolves only {@link #ALLOWED} to its content stream; returns {@code null} for anything else. */
    private static final XMLResolver STAX_ALLOW_LIST = (publicID, systemID, baseURI, namespace) -> {
        if (!ALLOWED.equals(systemID)) {
            return null;
        }
        try {
            return new URL(systemID).openStream();
        } catch (final IOException e) {
            throw new XMLStreamException(e);
        }
    };

    /** Absolute location of the imported schema the allow-list resolver permits. */
    private static final String ALLOWED_SCHEMA = AttackTestSupport.resourceUrl("included.xsd").toString();

    /** Resolves only the {@code included.xsd} import; returns {@code null} for anything else. */
    private static final LSResourceResolver SCHEMA_ALLOW_LIST = (type, namespaceURI, publicId, systemId, baseURI) ->
            systemId != null && systemId.endsWith("included.xsd") ? lsInput(ALLOWED_SCHEMA) : null;

    /** Resolves only the {@code included.xsl} import; returns {@code null} for anything else. */
    private static final URIResolver XSL_ALLOW_LIST = (href, base) ->
            href != null && href.endsWith("included.xsl") ? AttackTestSupport.resourceSource("included.xsl") : null;

    private static String entityPayload(final String entitySystemId) {
        return "<?xml version=\"1.0\"?>\n"
                + "<!DOCTYPE root [\n  <!ENTITY xxe SYSTEM \"" + entitySystemId + "\">\n]>\n"
                + "<root>&xxe;</root>";
    }

    private static XMLInputFactory externalEntityStaxFactory() {
        final XMLInputFactory factory = SecureXMLInputFactory.newInstance();
        factory.setProperty(XMLInputFactory.IS_SUPPORTING_EXTERNAL_ENTITIES, true);
        factory.setProperty(XMLInputFactory.IS_REPLACING_ENTITY_REFERENCES, true);
        return factory;
    }

    /** An {@link LSInput} naming the resource but carrying no content: a redirect the implementation fetches itself, like an identifier-only {@code InputSource}. */
    private static LSInput identifierOnlyLsInput(final String systemId) {
        return assertDoesNotThrow(() -> {
            final DOMImplementationLS ls = (DOMImplementationLS) DOMImplementationRegistry.newInstance().getDOMImplementation("LS");
            final LSInput input = ls.createLSInput();
            input.setSystemId(systemId);
            return input;
        }, "Failed to build identifier-only LSInput for " + systemId);
    }

    private static LSInput lsInput(final String systemId) {
        return assertDoesNotThrow(() -> {
            final LSInput input = identifierOnlyLsInput(systemId);
            input.setByteStream(new URL(systemId).openStream());
            return input;
        }, "Failed to build LSInput for " + systemId);
    }

    private static DocumentBuilder secureBuilder() throws Exception {
        final DocumentBuilder builder = SecureDocumentBuilderFactory.newInstance().newDocumentBuilder();
        builder.setErrorHandler(AttackTestSupport.STRICT_REPORTER);
        return builder;
    }

    /**
     * A secure {@link TransformerFactory} with a re-throwing error listener. XSLTC and Xalan enforce the block through the
     * {@link FallbackIgnoreURIResolver} floor; Saxon enforces it through the ignore-all resolver floor on its {@code Configuration}. Either way a caller-set
     * resolver that returns {@code null} cannot re-open the fetch. The strict listener turns any reported-and-recovered error into a test failure, so an
     * implementation cannot quietly recover from a floor resolution while the test asserts clean completion.
     */
    private static TransformerFactory secureTransformerFactory() {
        final TransformerFactory factory = SecureTransformerFactory.newInstance();
        factory.setErrorListener(AttackTestSupport.STRICT_REPORTER);
        return factory;
    }

    private static XMLReader secureXMLReader() throws Exception {
        final XMLReader reader = SecureSAXParserFactory.newInstance().newSAXParser().getXMLReader();
        reader.setErrorHandler(AttackTestSupport.STRICT_REPORTER);
        return reader;
    }

    private static DocumentBuilder xIncludeAwareBuilder() throws Exception {
        final DocumentBuilderFactory factory = SecureDocumentBuilderFactory.newInstance();
        factory.setNamespaceAware(true);
        AttackTestSupport.assumeDoesNotThrow(() -> factory.setXIncludeAware(true));
        final DocumentBuilder builder = factory.newDocumentBuilder();
        builder.setErrorHandler(AttackTestSupport.STRICT_REPORTER);
        return builder;
    }

    private static XMLReader xIncludeAwareReader() throws Exception {
        final SAXParserFactory factory = SecureSAXParserFactory.newInstance();
        factory.setNamespaceAware(true);
        AttackTestSupport.assumeDoesNotThrow(() -> factory.setXIncludeAware(true));
        final XMLReader reader = factory.newSAXParser().getXMLReader();
        reader.setErrorHandler(AttackTestSupport.STRICT_REPORTER);
        return reader;
    }

    @Test
    @Tag("dom")
    void domDoesNotLeakUnlisted() throws Exception {
        Assumptions.assumeTrue(AttackTestSupport.DOM_RESOLVES_INTERNAL_ENTITIES, "platform DOM does not resolve user-defined entities");
        final DocumentBuilder builder = secureBuilder();
        builder.setEntityResolver(ENTITY_ALLOW_LIST);
        // The caller returns null for the unlisted entity, so the floor resolves it to empty rather than fetching it: the parse completes (or is rejected)
        // without leaking the entity's content.
        try {
            final Document doc = builder.parse(AttackTestSupport.inputSource(entityPayload(UNLISTED)));
            assertFalse(doc.getDocumentElement().getTextContent().contains(AttackTestSupport.LEAKED_MARKER), "unlisted external entity leaked into the DOM");
        } catch (final SAXException blocked) {
            // Acceptable: the reference was rejected at parse rather than resolved to empty.
        }
    }

    @Test
    @Tag("dom")
    void domResolvesAllowListed() throws Exception {
        Assumptions.assumeTrue(AttackTestSupport.DOM_RESOLVES_INTERNAL_ENTITIES, "platform DOM does not resolve user-defined entities");
        final DocumentBuilder builder = secureBuilder();
        builder.setEntityResolver(ENTITY_ALLOW_LIST);
        final Document doc = builder.parse(AttackTestSupport.inputSource(entityPayload(ALLOWED)));
        assertTrue(doc.getDocumentElement().getTextContent().contains(AttackTestSupport.LEAKED_MARKER),
                "allow-listed external entity should resolve through the caller's resolver");
    }

    @Test
    @Tag("dom")
    void domResolvesRelativeXIncludeSibling() throws Exception {
        final DocumentBuilder builder = xIncludeAwareBuilder();
        builder.setEntityResolver(RESOLVE_ALL);
        final Document doc = builder.parse(XINCLUDE_HOST);
        assertTrue(doc.getDocumentElement().getTextContent().contains(AttackTestSupport.LEAKED_MARKER),
                "relative XInclude sibling should resolve through the caller's resolver after the floor absolutizes the href");
    }

    @Test
    @Tag("sax")
    void saxParseWithHandlerDoesNotBypass() throws Exception {
        // SAXParser.parse(source, handler) installs the handler as the reader's entity resolver; the handler does not resolve it (returns null), so the
        // ignore-all floor must still resolve the external entity to empty rather than letting the parser fetch it.
        final SAXParser parser = SecureSAXParserFactory.newInstance().newSAXParser();
        final StringBuilder text = new StringBuilder();
        try {
            parser.parse(AttackTestSupport.inputSource(entityPayload(ALLOWED)), AttackTestSupport.capturingHandler(text));
        } catch (final SAXException e) {
            return; // blocked at parse: acceptable
        }
        assertFalse(text.toString().contains(AttackTestSupport.LEAKED_MARKER), "parse(source, handler) leaked the external entity:\n" + text);
    }

    @Test
    @Tag("sax")
    void saxReaderDoesNotLeakUnlisted() throws Exception {
        final XMLReader reader = secureXMLReader();
        reader.setEntityResolver(ENTITY_ALLOW_LIST);
        // The caller returns null for the unlisted entity, so the floor resolves it to empty rather than fetching it.
        final String text;
        try {
            text = AttackTestSupport.captureCharacters(reader, entityPayload(UNLISTED));
        } catch (final SAXException blocked) {
            return; // Acceptable: rejected at parse rather than resolved to empty.
        }
        assertFalse(text.contains(AttackTestSupport.LEAKED_MARKER), "unlisted external entity leaked:\n" + text);
    }

    @Test
    @Tag("sax")
    void saxReaderResolvesAllowListed() throws Exception {
        final XMLReader reader = secureXMLReader();
        reader.setEntityResolver(ENTITY_ALLOW_LIST);
        final String text = AttackTestSupport.captureCharacters(reader, entityPayload(ALLOWED));
        assertTrue(text.contains(AttackTestSupport.LEAKED_MARKER),
                "allow-listed external entity should resolve through the caller's resolver");
    }

    @Test
    @Tag("sax")
    void saxResolvesRelativeXIncludeSibling() throws Exception {
        final XMLReader reader = xIncludeAwareReader();
        reader.setEntityResolver(RESOLVE_ALL);
        final String text = AttackTestSupport.captureCharacters(reader, new InputSource(XINCLUDE_HOST));
        assertTrue(text.contains(AttackTestSupport.LEAKED_MARKER),
                "relative XInclude sibling should resolve through the caller's resolver after the floor absolutizes the href");
    }

    @Test
    @Tag("schema")
    void schemaDeniesUnlisted() {
        assertParseFails(() -> {
            final SchemaFactory factory = SecureSchemaFactory.newInstance(XMLConstants.W3C_XML_SCHEMA_NS_URI);
            factory.setResourceResolver((type, namespaceURI, publicId, systemId, baseURI) -> null);
            factory.newSchema(AttackTestSupport.resourceSource("with-import.xsd"));
        }, "Schema import", SAXException.class, SecurityException.class);
    }

    @Test
    @Tag("schema")
    void schemaFetchesIdentifierOnlyOptIn() {
        // A non-null return is an opt-in even without content: the implementation fetches the named resource itself, mirroring the entity floor's contract.
        assertParseSucceeds(() -> {
            final SchemaFactory factory = SecureSchemaFactory.newInstance(XMLConstants.W3C_XML_SCHEMA_NS_URI);
            factory.setResourceResolver((type, namespaceURI, publicId, systemId, baseURI) ->
                    systemId != null && systemId.endsWith("included.xsd") ? identifierOnlyLsInput(ALLOWED_SCHEMA) : null);
            factory.newSchema(AttackTestSupport.resourceSource("with-import.xsd"));
        }, "Schema import via identifier-only LSInput");
    }

    @Test
    @Tag("schema")
    void schemaResolvesAllowListed() {
        // with-import.xsd references an element defined only in the imported included.xsd, so it compiles only if the import is resolved.
        assertParseSucceeds(() -> {
            final SchemaFactory factory = SecureSchemaFactory.newInstance(XMLConstants.W3C_XML_SCHEMA_NS_URI);
            factory.setResourceResolver(SCHEMA_ALLOW_LIST);
            factory.newSchema(AttackTestSupport.resourceSource("with-import.xsd"));
        }, "Schema import via caller resolver");
    }

    @Test
    @Tag("stax")
    void staxCallerCannotRemoveFloor() throws Exception {
        // A caller resolver that resolves nothing must not re-open external fetches: the floor still resolves the reference to empty rather than fetching it.
        final XMLInputFactory factory = externalEntityStaxFactory();
        factory.setXMLResolver((publicID, systemID, baseURI, namespace) -> null);
        try {
            assertFalse(AttackTestSupport.captureStaxEventText(factory, entityPayload(ALLOWED)).contains(AttackTestSupport.LEAKED_MARKER),
                    "floor was bypassed and the entity leaked");
        } catch (final XMLStreamException blocked) {
            // Acceptable: rejected at parse rather than resolved to empty.
        }
    }

    @Test
    @Tag("stax")
    void staxDoesNotLeakUnlisted() throws Exception {
        final XMLInputFactory factory = externalEntityStaxFactory();
        factory.setXMLResolver(STAX_ALLOW_LIST);
        // The caller returns null for the unlisted entity, so the floor resolves it to empty rather than fetching it.
        try {
            assertFalse(AttackTestSupport.captureStaxEventText(factory, entityPayload(UNLISTED)).contains(AttackTestSupport.LEAKED_MARKER),
                    "unlisted external entity leaked");
        } catch (final XMLStreamException blocked) {
            // Acceptable: rejected at parse rather than resolved to empty.
        }
    }

    @Test
    @Tag("stax")
    void staxGetXMLResolverReportsCallerUnwrapped() {
        final XMLInputFactory factory = SecureXMLInputFactory.newInstance();
        final XMLResolver caller = (publicID, systemID, baseURI, namespace) -> null;
        factory.setXMLResolver(caller);
        assertSame(caller, factory.getXMLResolver(), "getXMLResolver should report the caller's resolver, not the floor wrapper");
    }

    @Test
    @Tag("stax")
    void staxResolvesAllowListed() throws Exception {
        final XMLInputFactory factory = externalEntityStaxFactory();
        factory.setXMLResolver(STAX_ALLOW_LIST);
        assertTrue(AttackTestSupport.captureStaxEventText(factory, entityPayload(ALLOWED)).contains(AttackTestSupport.LEAKED_MARKER),
                "allow-listed external entity should resolve through the caller's resolver");
    }

    @Test
    @Tag("trax")
    void transformerDoesNotLeakUnlisted() throws Exception {
        final TransformerFactory factory = secureTransformerFactory();
        factory.setURIResolver((href, base) -> null);
        // Deterministic on every implementation: XSLTC and Xalan compile the empty document the URIResolver floor
        // returns, Saxon the EmptySource its Configuration floor returns, so the import contributes nothing.
        final StringWriter sink = new StringWriter();
        factory.newTemplates(AttackTestSupport.resourceSource("with-import.xsl")).newTransformer().transform(AttackTestSupport.streamSource("<root/>"),
                new StreamResult(sink));
        assertFalse(sink.toString().contains(AttackTestSupport.LEAKED_MARKER), "unlisted stylesheet import leaked");
    }

    @Test
    @Tag("trax")
    void transformerParsesOptedInDocumentSecured() throws Exception {
        // Same contract on the runtime document() channel, which reaches a different internal reader than the compile-time import.
        final TransformerFactory factory = secureTransformerFactory();
        factory.setURIResolver(
                (href, base) -> href != null && href.endsWith("referenced.xml") ? AttackTestSupport.resourceSource("referenced-with-entity.xml") : null);
        // Same undeclared-entity outcome as the import above: skipped, never expanded.
        final StringWriter sink = new StringWriter();
        factory.newTemplates(AttackTestSupport.resourceSource("with-document.xsl")).newTransformer().transform(AttackTestSupport.streamSource("<root/>"),
                new StreamResult(sink));
        assertFalse(sink.toString().contains(AttackTestSupport.LEAKED_MARKER), "opted-in document() resource leaked its external entity");
    }

    @Test
    @Tag("trax")
    void transformerParsesOptedInImportSecured() throws Exception {
        // The opted-in module carries an external DTD reference; parsed on the floor the DTD is empty, so its entity cannot expand into the output.
        final TransformerFactory factory = secureTransformerFactory();
        factory.setURIResolver(
                (href, base) -> href != null && href.endsWith("included.xsl") ? AttackTestSupport.resourceSource("included-with-entity.xsl") : null);
        // The emptied DTD leaves the entity undeclared — only a validity violation when an external subset is
        // declared — so every non-validating parser skips it and the transform deterministically completes.
        final StringWriter sink = new StringWriter();
        factory.newTemplates(AttackTestSupport.resourceSource("with-import.xsl")).newTransformer().transform(AttackTestSupport.streamSource("<root/>"),
                new StreamResult(sink));
        assertFalse(sink.toString().contains(AttackTestSupport.LEAKED_MARKER), "opted-in stylesheet import leaked its external entity");
    }

    @Test
    @Tag("trax")
    void transformerResolvesAllowListed() {
        // with-import.xsl imports included.xsl, so it compiles only if the import is resolved.
        final TransformerFactory factory = secureTransformerFactory();
        factory.setURIResolver(XSL_ALLOW_LIST);
        assertParseSucceeds(() -> factory.newTemplates(AttackTestSupport.resourceSource("with-import.xsl")), "Stylesheet import via caller resolver");
    }
}
