Re: A way to handle malicious XML with Expat / was Re: Handling malicious XML with Expat - what options do I have?

/*
 * Demo of handling malicious XML with Expat (tested with Expat 2.0.1)
 * v3.0 2008-09-16
 *
 * 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 <[email protected]>
 */

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

// Config, needs adjustment to match your use cases
int const MAX_BYTES_PER_ENTITY          = 100000;
int const MAX_TOTAL_LOOKUPS_PER_ENTITY  = 10000;
int const MAX_LOOKUP_DEPTH_PER_ENTITY   = 10;


struct EntityInfo {
	int valueLen;
	int totalLookup;
	int lookupDepth;

	EntityInfo(int valueLen, int totalLookup, int lookupDepth)
			: valueLen(valueLen), totalLookup(totalLookup),
			lookupDepth(lookupDepth) { }
};

typedef std::basic_string<XML_Char> StringType;
typedef std::map<StringType, EntityInfo> MapType;
typedef std::pair<StringType, EntityInfo> PairType;


// Global vars
MapType entityNameToValueLen;


void
initMap() {
	// Register default entities
	EntityInfo info(1, 0, 0);
	entityNameToValueLen.insert(PairType(StringType("amp"), info));
	entityNameToValueLen.insert(PairType(StringType("lt"), info));
	entityNameToValueLen.insert(PairType(StringType("gt"), info));
	entityNameToValueLen.insert(PairType(StringType("apos"), info));
	entityNameToValueLen.insert(PairType(StringType("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");
	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++;
	}
}

EntityInfo
getEntityInfo(XML_Char const * entityName) {
	MapType::iterator found = entityNameToValueLen.find(
			StringType(entityName));
	assert(found != entityNameToValueLen.end());
	return found->second;
}

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 totalLookup = 0;
	int lookupDepth = 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;
			totalLookup += info.totalLookup + 1;
			int const minLookupDepth = info.lookupDepth + 1;
			if (lookupDepth < minLookupDepth) {
				lookupDepth = minLookupDepth;
			}
			delete[] entityRefname;
		} else {
			valueLen += (afterSemiColon - walker);
			break;
		}
		walker = afterSemiColon;
	}

	int const bytesNeeded = valueLen * sizeof(XML_Char);
	::printf("  Length in bytes:  %d\n", bytesNeeded);
	::printf("  Total lookups:  %d\n", totalLookup);
	::printf("  Lookup depth:  %d\n", lookupDepth);
	EntityInfo const info(valueLen, totalLookup, lookupDepth);
	setEntityInfo(entityName, info);

#if 1
	// Prevent
	if (bytesNeeded > MAX_BYTES_PER_ENTITY) {
		panic(userData, "Entity takes too much space");
	} else if (totalLookup > MAX_TOTAL_LOOKUPS_PER_ENTITY) {
		panic(userData, "Entity requires too many lookups");
	} else if (lookupDepth > MAX_LOOKUP_DEPTH_PER_ENTITY) {
		panic(userData, "Entity requires too deep lookup");
	}
#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 1
	"<!ENTITY laugh0 \"ha\">\n"
# else
	"<!ENTITY laugh0 \"\">\n"
# endif
	"<!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",
				static_cast<int>(::XML_GetCurrentLineNumber(parser)),
				static_cast<int>(::XML_GetCurrentColumnNumber(parser)),
				::XML_ErrorString(::XML_GetErrorCode(parser)));
	}

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