heap underflow detector

heap-defender.cpp

#include <string>
#include <vector>
#include <Windows.h>

#include "SDK/plugin.h"
#include "jump-x86.h"

typedef void (*logprintf_t)(const char *format, ...);
static logprintf_t logprintf;

extern void *pAMXFunctions;

static JumpX86 amx_release_hook;

static void Init(void **ppData) {
	pAMXFunctions = ppData[PLUGIN_DATA_AMX_EXPORTS];
	logprintf = (logprintf_t)ppData[PLUGIN_DATA_LOGPRINTF];
}

std::string GetFileName(const std::string &path) {
	std::string::size_type lastSep = path.find_last_of("/\\");
	if (lastSep != std::string::npos) {
		return path.substr(lastSep + 1);
	}
	return path;
}

std::string GetModulePath(void *address, std::size_t maxLength) {
	std::vector<char> name(maxLength + 1);
	if (address != 0) {
		MEMORY_BASIC_INFORMATION mbi;
		VirtualQuery(address, &mbi, sizeof(mbi));
		GetModuleFileNameA((HMODULE)mbi.AllocationBase, &name[0], maxLength);
	}
	return std::string(&name[0]);
}

static int __cdecl amx_Release_impl(AMX *amx, cell amx_addr, void *caller) {
	JumpX86::ScopedRemove scoped_remove(&amx_release_hook);

 	if (amx_addr < amx->hlw) {
		std::string plugin_name = GetFileName(GetModulePath(caller, FILENAME_MAX));
		logprintf("WARNING: '%s' is doing a heap underflow (%x < %x)", plugin_name.c_str(), amx_addr, amx->hlw);
	}

	return amx_Release(amx, amx_addr);
}

static __declspec(naked) int __cdecl amx_Release_stub(AMX *amx, cell amx_addr) {
	__asm push ebp
	__asm mov ebp, esp

	__asm push dword ptr [ebp + 4]
	__asm push dword ptr [ebp + 12]
	__asm push dword ptr [ebp + 8]
	__asm call amx_Release_impl
	__asm add esp, 12

	__asm pop ebp
	__asm ret
}

PLUGIN_EXPORT unsigned int PLUGIN_CALL Supports() {
	return SUPPORTS_VERSION;
}

PLUGIN_EXPORT bool PLUGIN_CALL Load(void **ppData) {
	Init(ppData);

	amx_release_hook.Install(
		reinterpret_cast<void**>(pAMXFunctions)[PLUGIN_AMX_EXPORT_Release],
		amx_Release_stub
	);

	logprintf("  Heap Defender is OK.");
	return true;
}