/*
 * Demo of handling malicious XML with Expat (tested with Expat 2.0.1)
 * v1.0 2008-09-12
 *
 * Copyright (c) 2008 Sebastian Pipping
 *
 * == The MIT License ==
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 *
 * Sebastian Pipping <webmaster@hartwork.org>
 */

#include <expat.h>
#include <cstdio>
#include <cstring>
#include <map>

const int MAX_BYTES_PER_ENTITY_VALUE=100000;

struct XmlCharStringCompare {
	bool operator()(XML_Char const * s1,
			XML_Char const * s2) const {
		return strcmp(s1, s2) < 0;
	}
};

typedef std::map<XML_Char const *, int, XmlCharStringCompare> MapType;
typedef std::pair<XML_Char const *, int> PairType;


MapType entityNameToValueLen;


void
handleCharacterData(void *userData, const XML_Char *s, int len) {
	::puts("BEGIN handleCharacterData");
	XML_Char * toPrint = new XML_Char[len + 1];
	::strncpy(toPrint, s, len);
	toPrint[len] = '\0';
	::printf("  \"%s\"\n", toPrint);
	delete [] toPrint;
	::puts("END\n");
}

XML_Char *
makeString(XML_Char const * first, XML_Char const * afterLast) {
	size_t const len = afterLast - first;
	XML_Char * dup = new XML_Char[len + 1];
	::strncpy(dup, first, len);
	dup[len] = '\0';
	return dup;
}

XML_Char *
nextEntityRefMalloc(XML_Char const * start,
		XML_Char const * & atAmpersand,
		XML_Char const * & afterSemiColon) {
	XML_Char const * walker = start;
	while (true) {
		switch (walker[0]) {
		case '\0':
			// No complete entity found
			atAmpersand = start;
			afterSemiColon = walker;
			return NULL;

		case '&':
			// Entity start found
			atAmpersand = walker;
			break;

		case ';':
			// Entity stop found
			if (atAmpersand != NULL) {
				afterSemiColon = walker + 1;
				return makeString(atAmpersand + 1, walker);
			}
			break;
		}
		walker++;
	}
}

int
getValueLenForEntity(XML_Char const * entityName) {
	MapType::iterator found = entityNameToValueLen.find(entityName);
	if (found != entityNameToValueLen.end()) {
		return found->second;
	} else {
		return 0;
	}
}

void
setValueLenForEntity(XML_Char const * entityName, int valueLen) {
	entityNameToValueLen.insert(PairType(entityName, valueLen));
}

void
handleEntityDeclaration(void *userData, const XML_Char *entityName,
		int is_parameter_entity, const XML_Char *value,
		int value_length, const XML_Char *base, const XML_Char *systemId,
		const XML_Char *publicId, const XML_Char *notationName) {
	::puts("BEGIN handleEntityDeclaration");
	::printf("  %s := \"%s\"\n", entityName, value);

	XML_Char const * walker = value;
	int valueLen = 0;
	while (walker[0] != '\0') {
		XML_Char const * atAmpersand = NULL;
		XML_Char const * afterSemiColon = NULL;
		XML_Char * entityRefname = nextEntityRefMalloc(walker,
				atAmpersand, afterSemiColon);
		valueLen += (atAmpersand - walker);
		if (entityRefname != NULL) {
			valueLen += getValueLenForEntity(entityRefname);
		} else {
			valueLen += (afterSemiColon - walker);
			break;
		}
		walker = afterSemiColon;
	}

	::printf("  Length is %d\n", valueLen);
	setValueLenForEntity(entityName, valueLen);

	::puts("END\n");

	if (valueLen > MAX_BYTES_PER_ENTITY_VALUE) {
		::printf("Content consided malicious XML, aborting\n", valueLen);
		XML_Parser const parser = reinterpret_cast<XML_Parser>(userData);
		::XML_StopParser(parser, XML_FALSE);
	}
}

int
main() {
	char const * const document =
#if 0
		"<!DOCTYPE d [\n"
		"\t<!ENTITY a1 \"1a1\">\n"
		"\t<!ENTITY a2 \"2&a1;2&a1;2\">\n"
		"\t<!ENTITY a4 \"4&a2;4&a2;4\">\n"
		"]>\n"
		"<t>&a4;</t>\n"
#else
	// From http://www.cogsci.ed.ac.uk/~richard/billion-laughs.xml
	"<?xml version=\"1.0\"?>\n"
	"<!DOCTYPE billion [\n"
	"<!ELEMENT billion (#PCDATA)>\n"
	"<!ENTITY laugh0 \"ha\">\n"
	"<!ENTITY laugh1 \"&laugh0;&laugh0;\">\n"
	"<!ENTITY laugh2 \"&laugh1;&laugh1;\">\n"
	"<!ENTITY laugh3 \"&laugh2;&laugh2;\">\n"
	"<!ENTITY laugh4 \"&laugh3;&laugh3;\">\n"
	"<!ENTITY laugh5 \"&laugh4;&laugh4;\">\n"
	"<!ENTITY laugh6 \"&laugh5;&laugh5;\">\n"
	"<!ENTITY laugh7 \"&laugh6;&laugh6;\">\n"
	"<!ENTITY laugh8 \"&laugh7;&laugh7;\">\n"
	"<!ENTITY laugh9 \"&laugh8;&laugh8;\">\n"
	"<!ENTITY laugh10 \"&laugh9;&laugh9;\">\n"
	"<!ENTITY laugh11 \"&laugh10;&laugh10;\">\n"
	"<!ENTITY laugh12 \"&laugh11;&laugh11;\">\n"
	"<!ENTITY laugh13 \"&laugh12;&laugh12;\">\n"
	"<!ENTITY laugh14 \"&laugh13;&laugh13;\">\n"
	"<!ENTITY laugh15 \"&laugh14;&laugh14;\">\n"
	"<!ENTITY laugh16 \"&laugh15;&laugh15;\">\n"
	"<!ENTITY laugh17 \"&laugh16;&laugh16;\">\n"
	"<!ENTITY laugh18 \"&laugh17;&laugh17;\">\n"
	"<!ENTITY laugh19 \"&laugh18;&laugh18;\">\n"
	"<!ENTITY laugh20 \"&laugh19;&laugh19;\">\n"
	"<!ENTITY laugh21 \"&laugh20;&laugh20;\">\n"
	"<!ENTITY laugh22 \"&laugh21;&laugh21;\">\n"
	"<!ENTITY laugh23 \"&laugh22;&laugh22;\">\n"
	"<!ENTITY laugh24 \"&laugh23;&laugh23;\">\n"
	"<!ENTITY laugh25 \"&laugh24;&laugh24;\">\n"
	"<!ENTITY laugh26 \"&laugh25;&laugh25;\">\n"
	"<!ENTITY laugh27 \"&laugh26;&laugh26;\">\n"
	"<!ENTITY laugh28 \"&laugh27;&laugh27;\">\n"
	"<!ENTITY laugh29 \"&laugh28;&laugh28;\">\n"
	"<!ENTITY laugh30 \"&laugh29;&laugh29;\">\n"
	"]>\n"
	"<billion>&laugh30;</billion>\n"
#endif
	;

	XML_Parser const parser = ::XML_ParserCreate(NULL);
	::XML_SetCharacterDataHandler(parser, handleCharacterData);
	::XML_SetEntityDeclHandler(parser, handleEntityDeclaration);
	::XML_UseParserAsHandlerArg(parser);

	XML_Status const res = ::XML_Parse(parser, document, strlen(document), 1);
	bool const good = (res == XML_STATUS_OK);
	if (good) {
		::puts("All good.");
	} else {
		::printf("Error (Line %d, column %d): %s\n",
				::XML_GetCurrentLineNumber(parser),
				::XML_GetCurrentColumnNumber(parser),
				::XML_ErrorString(::XML_GetErrorCode(parser)));
	}

	::XML_ParserFree(parser);
	return good ? 0 : 1;
}