#include <windows.h>
#include "buffer.h"
#define MEMORY_BLOCK_SIZE 0x1000
#define MAX_MEMORY_RANGE 0x40000000
#define PAGE_EXECUTE_FLAGS \
(PAGE_EXECUTE | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY)
typedef struct _MEMORY_SLOT
{
union
{
struct _MEMORY_SLOT *pNext;
UINT8 buffer[MEMORY_SLOT_SIZE];
};
} MEMORY_SLOT, *PMEMORY_SLOT;
typedef struct _MEMORY_BLOCK
{
struct _MEMORY_BLOCK *pNext;
PMEMORY_SLOT pFree; UINT usedCount;
} MEMORY_BLOCK, *PMEMORY_BLOCK;
PMEMORY_BLOCK g_pMemoryBlocks;
VOID InitializeBuffer(VOID)
{
}
VOID UninitializeBuffer(VOID)
{
PMEMORY_BLOCK pBlock = g_pMemoryBlocks;
g_pMemoryBlocks = NULL;
while (pBlock)
{
PMEMORY_BLOCK pNext = pBlock->pNext;
VirtualFree(pBlock, 0, MEM_RELEASE);
pBlock = pNext;
}
}
#if defined(_M_X64) || defined(__x86_64__)
static LPVOID FindPrevFreeRegion(LPVOID pAddress, LPVOID pMinAddr, DWORD dwAllocationGranularity)
{
ULONG_PTR tryAddr = (ULONG_PTR)pAddress;
tryAddr -= tryAddr % dwAllocationGranularity;
tryAddr -= dwAllocationGranularity;
while (tryAddr >= (ULONG_PTR)pMinAddr)
{
MEMORY_BASIC_INFORMATION mbi;
if (VirtualQuery((LPVOID)tryAddr, &mbi, sizeof(mbi)) == 0)
break;
if (mbi.State == MEM_FREE)
return (LPVOID)tryAddr;
if ((ULONG_PTR)mbi.AllocationBase < dwAllocationGranularity)
break;
tryAddr = (ULONG_PTR)mbi.AllocationBase - dwAllocationGranularity;
}
return NULL;
}
#endif
#if defined(_M_X64) || defined(__x86_64__)
static LPVOID FindNextFreeRegion(LPVOID pAddress, LPVOID pMaxAddr, DWORD dwAllocationGranularity)
{
ULONG_PTR tryAddr = (ULONG_PTR)pAddress;
tryAddr -= tryAddr % dwAllocationGranularity;
tryAddr += dwAllocationGranularity;
while (tryAddr <= (ULONG_PTR)pMaxAddr)
{
MEMORY_BASIC_INFORMATION mbi;
if (VirtualQuery((LPVOID)tryAddr, &mbi, sizeof(mbi)) == 0)
break;
if (mbi.State == MEM_FREE)
return (LPVOID)tryAddr;
tryAddr = (ULONG_PTR)mbi.BaseAddress + mbi.RegionSize;
tryAddr += dwAllocationGranularity - 1;
tryAddr -= tryAddr % dwAllocationGranularity;
}
return NULL;
}
#endif
static PMEMORY_BLOCK GetMemoryBlock(LPVOID pOrigin)
{
PMEMORY_BLOCK pBlock;
#if defined(_M_X64) || defined(__x86_64__)
ULONG_PTR minAddr;
ULONG_PTR maxAddr;
SYSTEM_INFO si;
GetSystemInfo(&si);
minAddr = (ULONG_PTR)si.lpMinimumApplicationAddress;
maxAddr = (ULONG_PTR)si.lpMaximumApplicationAddress;
if ((ULONG_PTR)pOrigin > MAX_MEMORY_RANGE && minAddr < (ULONG_PTR)pOrigin - MAX_MEMORY_RANGE)
minAddr = (ULONG_PTR)pOrigin - MAX_MEMORY_RANGE;
if (maxAddr > (ULONG_PTR)pOrigin + MAX_MEMORY_RANGE)
maxAddr = (ULONG_PTR)pOrigin + MAX_MEMORY_RANGE;
maxAddr -= MEMORY_BLOCK_SIZE - 1;
#else
(void) (pOrigin);
#endif
for (pBlock = g_pMemoryBlocks; pBlock != NULL; pBlock = pBlock->pNext)
{
#if defined(_M_X64) || defined(__x86_64__)
if ((ULONG_PTR)pBlock < minAddr || (ULONG_PTR)pBlock >= maxAddr)
continue;
#endif
if (pBlock->pFree != NULL)
return pBlock;
}
#if defined(_M_X64) || defined(__x86_64__)
{
LPVOID pAlloc = pOrigin;
while ((ULONG_PTR)pAlloc >= minAddr)
{
pAlloc = FindPrevFreeRegion(pAlloc, (LPVOID)minAddr, si.dwAllocationGranularity);
if (pAlloc == NULL)
break;
pBlock = (PMEMORY_BLOCK)VirtualAlloc(
pAlloc, MEMORY_BLOCK_SIZE, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
if (pBlock != NULL)
break;
}
}
if (pBlock == NULL)
{
LPVOID pAlloc = pOrigin;
while ((ULONG_PTR)pAlloc <= maxAddr)
{
pAlloc = FindNextFreeRegion(pAlloc, (LPVOID)maxAddr, si.dwAllocationGranularity);
if (pAlloc == NULL)
break;
pBlock = (PMEMORY_BLOCK)VirtualAlloc(
pAlloc, MEMORY_BLOCK_SIZE, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
if (pBlock != NULL)
break;
}
}
#else
pBlock = (PMEMORY_BLOCK)VirtualAlloc(
NULL, MEMORY_BLOCK_SIZE, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
#endif
if (pBlock != NULL)
{
PMEMORY_SLOT pSlot = (PMEMORY_SLOT)pBlock + 1;
pBlock->pFree = NULL;
pBlock->usedCount = 0;
do
{
pSlot->pNext = pBlock->pFree;
pBlock->pFree = pSlot;
pSlot++;
} while ((ULONG_PTR)pSlot - (ULONG_PTR)pBlock <= MEMORY_BLOCK_SIZE - MEMORY_SLOT_SIZE);
pBlock->pNext = g_pMemoryBlocks;
g_pMemoryBlocks = pBlock;
}
return pBlock;
}
LPVOID AllocateBuffer(LPVOID pOrigin)
{
PMEMORY_SLOT pSlot;
PMEMORY_BLOCK pBlock = GetMemoryBlock(pOrigin);
if (pBlock == NULL)
return NULL;
pSlot = pBlock->pFree;
pBlock->pFree = pSlot->pNext;
pBlock->usedCount++;
#ifdef _DEBUG
memset(pSlot, 0xCC, sizeof(MEMORY_SLOT));
#endif
return pSlot;
}
VOID FreeBuffer(LPVOID pBuffer)
{
PMEMORY_BLOCK pBlock = g_pMemoryBlocks;
PMEMORY_BLOCK pPrev = NULL;
ULONG_PTR pTargetBlock = ((ULONG_PTR)pBuffer / MEMORY_BLOCK_SIZE) * MEMORY_BLOCK_SIZE;
while (pBlock != NULL)
{
if ((ULONG_PTR)pBlock == pTargetBlock)
{
PMEMORY_SLOT pSlot = (PMEMORY_SLOT)pBuffer;
#ifdef _DEBUG
memset(pSlot, 0x00, sizeof(*pSlot));
#endif
pSlot->pNext = pBlock->pFree;
pBlock->pFree = pSlot;
pBlock->usedCount--;
if (pBlock->usedCount == 0)
{
if (pPrev)
pPrev->pNext = pBlock->pNext;
else
g_pMemoryBlocks = pBlock->pNext;
VirtualFree(pBlock, 0, MEM_RELEASE);
}
break;
}
pPrev = pBlock;
pBlock = pBlock->pNext;
}
}
BOOL IsExecutableAddress(LPVOID pAddress)
{
MEMORY_BASIC_INFORMATION mi;
VirtualQuery(pAddress, &mi, sizeof(mi));
return (mi.State == MEM_COMMIT && (mi.Protect & PAGE_EXECUTE_FLAGS));
}