diff --git a/sourcehook/sh_list.h b/sourcehook/sh_list.h index d5f37e8..391df29 100644 --- a/sourcehook/sh_list.h +++ b/sourcehook/sh_list.h @@ -103,6 +103,20 @@ namespace SourceHook insert(begin(), obj); } + void push_sorted(const T &obj) + { + iterator iter; + for (iter = begin(); iter != end(); ++iter) + { + if (obj < *iter) + { + insert(iter, obj); + return; + } + } + push_back(obj); + } + size_t size() const { return m_Size; @@ -301,6 +315,7 @@ namespace SourceHook return *this; } }; + }; //NAMESPACE #endif //_INCLUDE_CSDM_LIST_H diff --git a/sourcehook/sh_memory.h b/sourcehook/sh_memory.h index d626bb9..a8e7077 100644 --- a/sourcehook/sh_memory.h +++ b/sourcehook/sh_memory.h @@ -43,6 +43,7 @@ # error Unsupported OS/Compiler # endif +#include "sh_list.h" namespace SourceHook { @@ -181,6 +182,258 @@ namespace SourceHook #endif } } + + /* + Class which lets us allocate memory regions in special pages only meant for on the fly code generation. + + If we alloc with malloc and then set the page access type to read/exec only, other regions returned by + malloc that are in the same page would lose their write access as well and the process could crash. + + Allocating one page per code generation session is usually a waste of memory and on some plattforms also + a waste of virtual address space (Windows’ VirtualAlloc has a granularity of 64K). + + + IMPORTANT: the memory that Alloc() returns is not a in a defined state! + It could be in read+exec OR read+write mode. + -> call SetRE() or SetRW() before using allocated memory! + */ + class CPageAlloc + { + struct AllocationUnit + { + size_t begin_offset; + size_t size; + + AllocationUnit(size_t p_offs, size_t p_size) : begin_offset(p_offs), size(p_size) + { + } + + bool operator < (const AllocationUnit &other) const + { + return begin_offset < other.begin_offset; + } + }; + + typedef List AUList; + + struct AllocatedRegion + { + void *startPtr; + size_t size; + bool isolated; // may contain only one AU + AUList allocUnits; + + bool TryAlloc(size_t reqsize, void * &outAddr) + { + // Check for isolated + if (isolated && !allocUnits.empty()) + return false; + + // Find the smallest gap where req fits + size_t lastend = 0; + size_t smallestgap_pos = size + 1; + size_t smallestgap_size = size + 1; + + for (AUList::iterator iter = allocUnits.begin(); iter != allocUnits.end(); ++iter) + { + if (iter->begin_offset - lastend >= reqsize) + { + if (iter->begin_offset - lastend < smallestgap_size) + { + smallestgap_size = iter->begin_offset - lastend; + smallestgap_pos = lastend; + } + } + lastend = iter->begin_offset + iter->size; + } + + if (size - lastend >= reqsize) + { + if (size - lastend < smallestgap_size) + { + smallestgap_size = size - lastend; + smallestgap_pos = lastend; + } + } + + if (smallestgap_pos < size) + { + outAddr = reinterpret_cast(reinterpret_cast(startPtr) + smallestgap_pos); + allocUnits.push_sorted( AllocationUnit(smallestgap_pos, reqsize) ); + return true; + } + else + { + return false; + } + } + + bool TryFree(void *addr) + { + if (addr < startPtr || addr >= reinterpret_cast(reinterpret_cast(startPtr) + size)) + return false; + + size_t offs = reinterpret_cast(addr) - reinterpret_cast(startPtr); + + for (AUList::iterator iter = allocUnits.begin(); iter != allocUnits.end(); ++iter) + { + if (iter->begin_offset == offs) + { + allocUnits.erase(iter); + return true; + } + } + + return false; + } + + bool Contains(void *addr) + { + return addr >= startPtr && addr < reinterpret_cast(reinterpret_cast(startPtr) + size); + } + + void FreeRegion() + { +#ifdef __linux__ + munmap(startPtr, size); +#else + VirtualFree(startPtr, 0, MEM_RELEASE); +#endif + } + }; + + typedef List ARList; + + size_t m_PageSize; + ARList m_Regions; + + bool AddRegion(size_t minSize, bool isolated) + { + AllocatedRegion newRegion; + newRegion.startPtr = 0; + newRegion.isolated = isolated; + + // Compute real size -> align up to m_PageSize boundary + + newRegion.size = minSize - (minSize % m_PageSize); + if (newRegion.size < minSize) + newRegion.size += m_PageSize; + +#ifdef __linux__ + newRegion.startPtr = mmap(0, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAN_ANONYMOUS, -1); +#else + newRegion.startPtr = VirtualAlloc(NULL, newRegion.size, MEM_COMMIT, PAGE_READWRITE); +#endif + + if (newRegion.startPtr) + { + m_Regions.push_back(newRegion); + return true; + } + else + { + return false; + } + + } + + void *AllocPriv(size_t size, bool isolated) + { + void *addr; + + if (!isolated) + { + for (ARList::iterator iter = m_Regions.begin(); iter != m_Regions.end(); ++iter) + { + if (iter->TryAlloc(size, addr)) + return addr; + } + } + + if (!AddRegion(size, isolated)) + return NULL; + + bool tmp = m_Regions.back().TryAlloc(size, addr); + SH_ASSERT(tmp, ("TryAlloc fails after AddRegion")); + return addr; + } + + public: + CPageAlloc() + { +#ifdef __linux__ + m_PageSize = sysconf(_SC_PAGESIZE); +#else + SYSTEM_INFO sysInfo; + GetSystemInfo(&sysInfo); + m_PageSize = sysInfo.dwPageSize; +#endif + } + + ~CPageAlloc() + { + // Free all regions + for (ARList::iterator iter = m_Regions.begin(); iter != m_Regions.end(); ++iter) + { + iter->FreeRegion(); + } + } + + void *Alloc(size_t size) + { + return AllocPriv(size, false); + } + + void *AllocIsolated(size_t size) + { + return AllocPriv(size, true); + } + + void Free(void *ptr) + { + for (ARList::iterator iter = m_Regions.begin(); iter != m_Regions.end(); ++iter) + { + if (iter->TryFree(ptr)) + { + if (iter->allocUnits.empty()) + { + iter->FreeRegion(); + m_Regions.erase(iter); + } + break; + } + } + } + + void SetRE(void *ptr) + { + for (ARList::iterator iter = m_Regions.begin(); iter != m_Regions.end(); ++iter) + { + if (iter->Contains(ptr)) + { + SetMemAccess(iter->startPtr, iter->size, SH_MEM_READ | SH_MEM_EXEC); + break; + } + } + } + + void SetRW(void *ptr) + { + for (ARList::iterator iter = m_Regions.begin(); iter != m_Regions.end(); ++iter) + { + if (iter->Contains(ptr)) + { + SetMemAccess(iter->startPtr, iter->size, SH_MEM_READ | SH_MEM_WRITE); + break; + } + } + } + + size_t GetPageSize() + { + return m_PageSize; + } + }; } #endif diff --git a/sourcehook/sourcehook_hookmangen.cpp b/sourcehook/sourcehook_hookmangen.cpp index e00a50d..ef5c6a3 100644 --- a/sourcehook/sourcehook_hookmangen.cpp +++ b/sourcehook/sourcehook_hookmangen.cpp @@ -25,6 +25,8 @@ namespace SourceHook { namespace Impl { + CPageAlloc GenBuffer::ms_Allocator; + template jit_int32_t DownCastPtr(T ptr) { @@ -1123,6 +1125,9 @@ namespace SourceHook // m_HookfuncVfnPtr is a pointer to a void* because SH expects a pointer // into the hookman's vtable *m_HookfuncVfnptr = reinterpret_cast(m_HookFunc.GetData()); + + m_HookFunc.SetRE(); + return m_HookFunc.GetData(); } @@ -1211,8 +1216,7 @@ namespace SourceHook IA32_Pop_Reg(&m_PubFunc, REG_EBP); IA32_Return(&m_PubFunc); - SetMemAccess(reinterpret_cast(m_PubFunc.GetData()), m_PubFunc.GetSize(), - SH_MEM_READ | SH_MEM_EXEC); + m_PubFunc.SetRE(); return m_PubFunc; } diff --git a/sourcehook/sourcehook_hookmangen.h b/sourcehook/sourcehook_hookmangen.h index 4bfeea7..1d95daf 100644 --- a/sourcehook/sourcehook_hookmangen.h +++ b/sourcehook/sourcehook_hookmangen.h @@ -42,84 +42,84 @@ namespace SourceHook class GenBuffer { - unsigned char *m_pData; - jitoffs_t m_Size; - jitoffs_t m_AllocatedSize; - - public: - GenBuffer() : m_pData(NULL), m_Size(0), m_AllocatedSize(0) - { - } - ~GenBuffer() - { - clear(); - } - jitoffs_t GetSize() - { - return m_Size; - } - unsigned char *GetData() - { - return m_pData; - } - - template void push(PT what) - { - push((const unsigned char *)&what, sizeof(PT)); - } - - void push(const unsigned char *data, jitoffs_t size) - { - jitoffs_t newSize = m_Size + size; - if (newSize > m_AllocatedSize) - { - m_AllocatedSize = newSize > m_AllocatedSize*2 ? newSize : m_AllocatedSize*2; - if (m_AllocatedSize < 64) - m_AllocatedSize = 64; - - unsigned char *newBuf; - try - { - //!!!! Better use of pages! or something! - newBuf = reinterpret_cast(VirtualAlloc(NULL, m_AllocatedSize, MEM_COMMIT, PAGE_READWRITE)); - } - catch (std::bad_alloc) - { - newBuf = NULL; - } - if (!newBuf) - { - SH_ASSERT(0, ("bad_alloc: couldn't allocate 0x%08X bytes of memory\n", m_AllocatedSize)); - return; - } - memcpy((void*)newBuf, (const void*)m_pData, m_Size); - if (m_pData) - VirtualFree(m_pData, 0, MEM_RELEASE); - m_pData = newBuf; - } - memcpy((void*)(m_pData + m_Size), (const void*)data, size); - m_Size = newSize; - } - - template void rewrite(jitoffs_t offset, PT what) - { - rewrite(offset, (const unsigned char *)&what, sizeof(PT)); - } - - void rewrite(jitoffs_t offset, const unsigned char *data, jitoffs_t size) + static CPageAlloc ms_Allocator; + + unsigned char *m_pData; + jitoffs_t m_Size; + jitoffs_t m_AllocatedSize; + + public: + GenBuffer() : m_pData(NULL), m_Size(0), m_AllocatedSize(0) + { + } + ~GenBuffer() + { + clear(); + } + jitoffs_t GetSize() + { + return m_Size; + } + unsigned char *GetData() + { + return m_pData; + } + + template void push(PT what) + { + push((const unsigned char *)&what, sizeof(PT)); + } + + void push(const unsigned char *data, jitoffs_t size) + { + jitoffs_t newSize = m_Size + size; + if (newSize > m_AllocatedSize) + { + m_AllocatedSize = newSize > m_AllocatedSize*2 ? newSize : m_AllocatedSize*2; + if (m_AllocatedSize < 64) + m_AllocatedSize = 64; + + unsigned char *newBuf; + newBuf = reinterpret_cast(ms_Allocator.Alloc(m_AllocatedSize)); + ms_Allocator.SetRW(newBuf); + if (!newBuf) + { + SH_ASSERT(0, ("bad_alloc: couldn't allocate 0x%08X bytes of memory\n", m_AllocatedSize)); + return; + } + memcpy((void*)newBuf, (const void*)m_pData, m_Size); + if (m_pData) + ms_Allocator.Free(reinterpret_cast(m_pData)); + m_pData = newBuf; + } + memcpy((void*)(m_pData + m_Size), (const void*)data, size); + m_Size = newSize; + } + + template void rewrite(jitoffs_t offset, PT what) + { + rewrite(offset, (const unsigned char *)&what, sizeof(PT)); + } + + void rewrite(jitoffs_t offset, const unsigned char *data, jitoffs_t size) { SH_ASSERT(offset + size <= m_AllocatedSize, ("rewrite too far")); memcpy((void*)(m_pData + offset), (const void*)data, size); - } - - void clear() - { - if (m_pData) - VirtualFree(m_pData, 0, MEM_RELEASE); - m_pData = NULL; - m_Size = 0; - m_AllocatedSize = 0; + } + + void clear() + { + if (m_pData) + ms_Allocator.Free(reinterpret_cast(m_pData)); + m_pData = NULL; + m_Size = 0; + m_AllocatedSize = 0; + } + + void SetRE() + { + ms_Allocator.SetRE(reinterpret_cast(m_pData)); } operator unsigned char *() @@ -149,7 +149,6 @@ namespace SourceHook { offs = get_outputpos() - offs; } - // :TODO: real buffer which uses virtualalloc correctly }; class GenContext diff --git a/sourcehook/test/main.cpp b/sourcehook/test/main.cpp index c7946bf..bdf3216 100644 --- a/sourcehook/test/main.cpp +++ b/sourcehook/test/main.cpp @@ -47,6 +47,7 @@ DECL_TEST(Ref); DECL_TEST(RefRet); DECL_TEST(VPHooks); DECL_TEST(HookManGen); +DECL_TEST(CPageAlloc); // in testhookmangen.cpp int main(int argc, char *argv[]) { @@ -70,6 +71,7 @@ int main(int argc, char *argv[]) DO_TEST(RefRet); DO_TEST(VPHooks); DO_TEST(HookManGen); + DO_TEST(CPageAlloc); cout << endl << "----" << endl << "Passed: " << passed << endl << "Failed: " << failed << endl; cout << "Total: " << passed + failed << endl; diff --git a/sourcehook/test/msvc8/test.vcproj b/sourcehook/test/msvc8/test.vcproj index ae768e5..a11c8e1 100644 --- a/sourcehook/test/msvc8/test.vcproj +++ b/sourcehook/test/msvc8/test.vcproj @@ -371,6 +371,10 @@ RelativePath="..\..\sourcehook.cpp" > + + @@ -395,6 +399,10 @@ RelativePath="..\testbail2.cpp" > + + @@ -431,6 +439,10 @@ RelativePath="..\testrefret.cpp" > + + + + @@ -460,6 +476,14 @@ RelativePath="..\..\sourcehook.h" > + + + + @@ -500,6 +524,10 @@ /> + + The difference is at least one page + CHECK_COND(static_cast(abs(test4[1] - test4[0])) >= ps, "Part 3.1"); + CHECK_COND(static_cast(abs(test4[2] - test4[1])) >= ps, "Part 3.2"); + CHECK_COND(static_cast(abs(test4[3] - test4[2])) >= ps, "Part 3.3"); + + CHECK_COND(static_cast(abs(test4[5] - test4[4])) >= ps, "Part 3.4"); + CHECK_COND(static_cast(abs(test4[6] - test4[5])) >= ps, "Part 3.5"); + CHECK_COND(static_cast(abs(test4[7] - test4[6])) >= ps, "Part 3.6"); + + // Thus i can set everything except for test4[2] to RE and still write to test4[2] + + alloc.SetRW(test4[2]); + + alloc.SetRE(test4[0]); + alloc.SetRE(test4[1]); + alloc.SetRE(test4[3]); + alloc.SetRE(test4[4]); + alloc.SetRE(test4[5]); + alloc.SetRE(test4[6]); + alloc.SetRE(test4[7]); + + memset((void*)test4[2], 0, ps / 4); + + return true; +} \ No newline at end of file