/*
 *  MinHook - Minimalistic API Hook Library
 *  Copyright (C) 2009 Tsuda Kageyu. All rights reserved.
 *
 *  Redistribution and use in source and binary forms, with or without
 *  modification, are permitted provided that the following conditions
 *  are met:
 *
 *  1. Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *  2. Redistributions in binary form must reproduce the above copyright
 *     notice, this list of conditions and the following disclaimer in the
 *     documentation and/or other materials provided with the distribution.
 *  3. The name of the author may not be used to endorse or promote products
 *     derived from this software without specific prior written permission.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
 *  IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 *  OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 *  IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
 *  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
 *  NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 *  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 *  THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 *  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 *  THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <cassert>
#include <vector>
#include <algorithm>
#include <Windows.h>

#include "buffer.h"

namespace MinHook { namespace
{
	struct MEMORY_BLOCK
	{
		void*	pAddress;
		DWORD	protect;
		size_t	uncommittedSize;
		size_t	uncommittedCount;
		size_t	fixedSize;
		size_t	fixedCount;
	};

	template <typename T>
	bool operator <(const MEMORY_BLOCK& lhs, const T& rhs);
	template <typename T>
	bool operator <(const T& lhs, const MEMORY_BLOCK& rhs);
	bool operator <(const MEMORY_BLOCK& lhs, const MEMORY_BLOCK& rhs);

	void*			AllocateBuffer(void* const pOrigin, DWORD protect, size_t size);
	MEMORY_BLOCK*	GetMemoryBlock(void* const pOrigin, DWORD protect, size_t capacity);

	const size_t BlockSize = 0x10000;

#if defined _M_X64
	intptr_t gMinAddress;
	intptr_t gMaxAddress;
#endif
	std::vector<MEMORY_BLOCK> gMemoryBlocks;
}}

namespace MinHook
{
	void InitializeBuffer()
	{
#if defined _M_X64
		SYSTEM_INFO si;
		GetSystemInfo(&si);

		gMinAddress = reinterpret_cast<intptr_t>(si.lpMinimumApplicationAddress);
		gMaxAddress = reinterpret_cast<intptr_t>(si.lpMaximumApplicationAddress);
#endif
	}

	void UninitializeBuffer()
	{
		for (size_t i = 0, count = gMemoryBlocks.size(); i < count; ++i)
		{
			MEMORY_BLOCK& block = gMemoryBlocks[i];
			VirtualFree(block.pAddress, 0, MEM_RELEASE);
		}

		std::vector<MEMORY_BLOCK> v;
		gMemoryBlocks.swap(v);
	}

	void* AllocateCodeBuffer(void* const pOrigin, size_t size)
	{
		assert(("AllocateBuffer", (size > 0)));

		return AllocateBuffer(pOrigin, PAGE_EXECUTE_READ, size);
	}

	void* AllocateDataBuffer(void* const pOrigin, size_t size)
	{
		assert(("AllocateBuffer", (size > 0)));

		return AllocateBuffer(pOrigin, PAGE_READONLY, size);
	}

	void FreeBuffer(void* const pBuffer)
	{
		for (size_t i = 0, count = gMemoryBlocks.size(); i < count; ++i)
		{
			MEMORY_BLOCK& block = gMemoryBlocks[i];
			void* pBlockEnd = reinterpret_cast<char*>(block.pAddress) + block.fixedSize;

			if (pBuffer >= block.pAddress && pBuffer < pBlockEnd)
			{
				assert(("FreeBuffer", (block.uncommittedSize == 0 && block.uncommittedCount == 0)));

				block.fixedCount--;

				if (block.fixedCount == 0)
				{
					VirtualFree(block.pAddress, 0, MEM_RELEASE);
					gMemoryBlocks.erase(gMemoryBlocks.begin()+i);
				}

				return;
			}
		}

		assert(("FreeBuffer", 0));
	}

	void RollbackBuffer()
	{
		for (size_t i = 0, count = gMemoryBlocks.size(); i < count; ++i)
		{
			MEMORY_BLOCK& block = gMemoryBlocks[i];
			if (block.uncommittedSize == 0)
			{
				continue;
			}

			block.uncommittedSize = 0;
			block.uncommittedCount = 0;

			if (block.fixedCount == 0)
			{
				VirtualFree(block.pAddress, 0, MEM_RELEASE);
				gMemoryBlocks.erase(gMemoryBlocks.begin()+i);
				i--;
				count--;
			}
		}
	}

	void CommitBuffer()
	{
		for (size_t i = 0, count = gMemoryBlocks.size(); i < count; ++i)
		{
			MEMORY_BLOCK& block = gMemoryBlocks[i];
			if (block.uncommittedSize == 0)
			{
				continue;
			}

			void* pBuffer = reinterpret_cast<char*>(block.pAddress) + block.fixedSize;
			size_t size = block.uncommittedSize;
			DWORD op;
			VirtualProtect(pBuffer, size, block.protect, &op);

			block.fixedSize += size;
			block.uncommittedSize = 0;

			block.fixedCount += block.uncommittedCount;
			block.uncommittedCount = 0;
		}
	}
}

namespace MinHook { namespace
{
	void* AllocateBuffer(void* const pOrigin, DWORD protect, size_t size)
	{
		assert(("AllocateBuffer", (protect == PAGE_EXECUTE_READ || protect == PAGE_READONLY)));
		assert(("AllocateBuffer", (size > 0)));

		// Round up to the alignment.
		size = (size + TYPE_ALIGNMENT(void*) - 1) & ~(TYPE_ALIGNMENT(void*) - 1);

		MEMORY_BLOCK* pBlock = GetMemoryBlock(pOrigin, protect, size);
		if (pBlock == NULL)
		{
			return NULL;
		}

		void* pBuffer = reinterpret_cast<char*>(pBlock->pAddress) + pBlock->fixedSize + pBlock->uncommittedSize;
		if (VirtualAlloc(pBuffer, size, MEM_COMMIT, pBlock->protect) == NULL)
		{
			return NULL;
		}

		DWORD oldProtect;
		// PAGE_EXECUTE_READ -> PAGE_EXECUTE_READWRITE, PAGE_READONLY -> PAGE_READWRITE
		if (!VirtualProtect(pBuffer, size, (pBlock->protect << 1), &oldProtect))
		{
			return NULL;
		}

		pBlock->uncommittedSize += size;
		pBlock->uncommittedCount++;
		return pBuffer;
	}

	MEMORY_BLOCK* GetMemoryBlock(void* const pOrigin, DWORD protect, size_t capacity)
	{
		assert(("GetMemoryBlock", (protect == PAGE_EXECUTE_READ || protect == PAGE_READONLY)));
		assert(("GetMemoryBlock", (capacity > 0)));

		typedef std::vector<MEMORY_BLOCK>::iterator mb_iter;

#if defined _M_X64
		intptr_t minAddr = gMinAddress;
		intptr_t maxAddr = gMaxAddress;
		if (pOrigin != NULL)
		{
			// pOrigin ± 512MB
			minAddr = std::max<intptr_t>(minAddr, reinterpret_cast<intptr_t>(pOrigin) - 0x20000000);
			maxAddr = std::min<intptr_t>(maxAddr, reinterpret_cast<intptr_t>(pOrigin) + 0x20000000);
		}
#endif

		// Look the registered blocks for a reachable one.
		MEMORY_BLOCK* pBlock = NULL;
		{
			mb_iter ib = gMemoryBlocks.begin();
			mb_iter ie = gMemoryBlocks.end();
#if defined _M_X64
			if (pOrigin != NULL)
			{
				// Ignore the blocks too far.
				ib = std::lower_bound(ib, ie, minAddr);
				ie = std::lower_bound(ib, ie, maxAddr);
			}
#endif
			for (mb_iter i = ib; i != ie; ++i)
			{
				if (i->protect == protect && i->fixedSize + i->uncommittedSize + capacity <= BlockSize)
				{
					return &(*i);
				}
			}
		}

		// Alloc a new block if not found.
		void* pAlloc = NULL;
#if defined _M_X64
		if (pOrigin != NULL)
		{
			// Seek a unallocated area from the center to the outside.
			intptr_t min = minAddr / BlockSize;
			intptr_t max = maxAddr / BlockSize;
			int rel = 0;
			MEMORY_BASIC_INFORMATION mi = { 0 };
			for (int i = 0; i < (max - min + 1); ++i)
			{
				rel = -rel + (i & 1);
				void* pQuery = reinterpret_cast<void*>(((min + max) / 2 + rel) * BlockSize);
				VirtualQuery(pQuery, &mi, sizeof(mi));
				if (mi.State == MEM_FREE)
				{
					pAlloc = VirtualAlloc(pQuery, BlockSize, MEM_RESERVE, protect);
					if (pAlloc != NULL)
					{
						break;
					}
				}
			}
		}
		else
#endif		// The address doesn't matter in x86 mode.
		{
			pAlloc = VirtualAlloc(NULL, BlockSize, MEM_RESERVE, protect);
		}

		if (pAlloc != NULL)
		{
			MEMORY_BLOCK block = { 0 };
			block.pAddress = pAlloc;
			block.protect = protect;

#if defined _M_X64
			mb_iter i = std::lower_bound(gMemoryBlocks.begin(), gMemoryBlocks.end(), pAlloc);
#elif defined _M_IX86
			mb_iter i = gMemoryBlocks.end();
#endif
			i = gMemoryBlocks.insert(i, block);

			return &(*i);
		}

		return NULL;
	}

	template <typename T>
	bool operator <(const MEMORY_BLOCK& lhs, const T& rhs)
	{
		return lhs.pAddress < reinterpret_cast<void*>(rhs);
	}

	template <typename T>
	bool operator <(const T& lhs, const MEMORY_BLOCK& rhs)
	{
		return reinterpret_cast<void*>(lhs) < rhs.pAddress;
	}

	bool operator <(const MEMORY_BLOCK& lhs, const MEMORY_BLOCK& rhs)
	{
		return lhs.pAddress < rhs.pAddress;
	}
}}