/*
 * Demo of handling malicious XML with Expat (tested with Expat 2.0.1)
 * v2.0 2008-09-13
 *
 * 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>

// Config
int const MAX_BYTES_PER_ENTITY_VALUE = 100000;
int const MAX_LOOKUPS_PER_ENTITY_VALUE = 30;
int const MAX_OUTPUT_INPUT_RATIO = 20;


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

struct EntityInfo {
	int valueLen;
	int maxLookups;

	EntityInfo(int valueLen, int maxLookups)
			: valueLen(valueLen), maxLookups(maxLookups) {

	}
};

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


// Global vars
XML_Index totalCharDataBytesServed = 0;
MapType entityNameToValueLen;


XML_Char *
dupString(XML_Char const * source) {
	size_t const len = ::strlen(source);
	XML_Char * const dup = new XML_Char[len + 1];
	::strcpy(dup, source);
	return dup;
}

void
initMap() {
	// Register default entities
	EntityInfo info(1, 0);
	entityNameToValueLen.insert(PairType(dupString("amp"), info));
	entityNameToValueLen.insert(PairType(dupString("lt"), info));
	entityNameToValueLen.insert(PairType(dupString("gt"), info));
	entityNameToValueLen.insert(PairType(dupString("apos"), info));
	entityNameToValueLen.insert(PairType(dupString("quot"), info));
}

XML_Parser
getParser(void * userData) {
	return reinterpret_cast<XML_Parser>(userData);
}

void
panic(void * userData, XML_Char const * diagonis) {
	::puts("\n  PANIC:");
	::printf("    %s\n", diagonis);
	::puts("    -> Content considered malicious XML");
	::puts("    -> Aborting");
	::XML_StopParser(getParser(userData), XML_FALSE);
}

void
handleCharacterData(void *userData, const XML_Char *s, int len) {
	::puts("BEGIN handleCharacterData");

	totalCharDataBytesServed += (len * sizeof(XML_Char));
	XML_Index const byteIndex = XML_GetCurrentByteIndex(getParser(userData));

	XML_Char * toPrint = new XML_Char[len + 1];
	::strncpy(toPrint, s, len);
	toPrint[len] = '\0';
	::printf("  \"%s\"\n", toPrint);
	::printf("  Byte index:              %u\n", byteIndex);
	::printf("  Char data bytes served:  %u\n", totalCharDataBytesServed);
	delete [] toPrint;

#if 1
	// Prevent
	if (totalCharDataBytesServed > MAX_OUTPUT_INPUT_RATIO * byteIndex) {
		panic(userData, "Output/input ratio too large");
	}
#endif

	::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++;
	}
}

EntityInfo
getEntityInfo(XML_Char const * entityName) {
	static EntityInfo errorInfo(0, 0);
	MapType::iterator found = entityNameToValueLen.find(entityName);
	if (found != entityNameToValueLen.end()) {
		return found->second;
	} else {
		return errorInfo;
	}
}

void
setEntityInfo(XML_Char const * name, EntityInfo const & info) {
	entityNameToValueLen.insert(PairType(name, info));
}

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;
	int maxLookups = 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) {
			EntityInfo const info = getEntityInfo(entityRefname);
			valueLen += info.valueLen;
			maxLookups += info.maxLookups + 1;
		} else {
			valueLen += (afterSemiColon - walker);
			break;
		}
		walker = afterSemiColon;
	}

	int const bytesNeeded = valueLen * sizeof(XML_Char);
	::printf("  Length in bytes:  %d\n", bytesNeeded);
	::printf("  Maximum lookups:  %d\n", maxLookups);
	EntityInfo const info(valueLen, maxLookups);
	setEntityInfo(entityName, info);

#if 1
	// Prevent
	if (bytesNeeded > MAX_BYTES_PER_ENTITY_VALUE) {
		panic(userData, "Entity takes too much space");
	} else if (maxLookups > MAX_LOOKUPS_PER_ENTITY_VALUE) {
		panic(userData, "Entity requires too many lookups");
	}
#endif

	::puts("END\n");
}

int
main() {
	initMap();

	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"
# if 0
	"<!ENTITY laugh0 \"ha\">\n"
# else
#  if 0
	"<!ENTITY laugh0 \"\">\n"
#  else
	"<!ENTITY laugh0 \"" "hahahahahahahaha" "hahahahahahahaha" "hahahahahahahaha"
	"hahahahahahahaha" "hahahahahahahaha" "hahahahahahahaha" "hahahahahahahaha"
	"hahahahahahahaha" "hahahahahahahaha" "hahahahahahahaha" "hahahahahahahaha"
	"hahahahahahahaha" "hahahahahahahaha" "hahahahahahahaha" "hahahahahahahaha"
	"\">\n"
#  endif
# endif
	"<!ENTITY laugh1 \"&laugh0;&laugh0;\">\n"
	"<!ENTITY laugh2 \"&laugh1;&laugh1;\">\n"
# if 0
	"<!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"
# else
	"]>\n"
	"<foo>\n"
	"&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;\n"
	"&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;\n"
	"&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;\n"
	"&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;\n"
	"&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;\n"
	"&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;\n"
	"&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;\n"
	"&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;\n"
	"&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;\n"
	"&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;\n"
	"&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;\n"
	"&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;\n"
	"&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;&laugh2;\n"
	"</foo>"
# endif
#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;
}