diff --git a/src/patches/memorymanager.cpp b/src/patches/memorymanager.cpp index 035c7e4..32dc30e 100644 --- a/src/patches/memorymanager.cpp +++ b/src/patches/memorymanager.cpp @@ -80,34 +80,58 @@ namespace namespace MemoryManager { - void* Allocate(RE::MemoryManager*, std::size_t a_size, std::uint32_t a_alignment, bool a_alignmentRequired) - { - if (a_size > 0) - return a_alignmentRequired ? - scalable_aligned_malloc(a_size, a_alignment) : - scalable_malloc(a_size); - else - return g_trash; + void *Allocate(RE::MemoryManager *, std::size_t a_size, std::uint32_t a_alignment, bool a_alignmentRequired) { + if (a_size > 0) { + void *return_ptr = a_alignmentRequired ? scalable_aligned_malloc(a_size + 0x100, a_alignment) : scalable_malloc(a_size + 0x100); + uint64_t *return_ptr_offset = (uint64_t *) (((uint64_t) return_ptr) + 0x100); + *(uint64_t *) return_ptr = 0x1337DEADDEAD1337; + //logger::trace("MemoryManagerStats::Allocate ptr {} (internal ptr {}) diff {}.", reinterpret_cast(return_ptr_offset), reinterpret_cast(return_ptr), reinterpret_cast(return_ptr_offset) - reinterpret_cast(return_ptr)); + return (void *) return_ptr_offset; + } + else return g_trash; } - void Deallocate(RE::MemoryManager*, void* a_mem, bool a_alignmentRequired) - { - if (a_mem != g_trash) - a_alignmentRequired ? - scalable_aligned_free(a_mem) : - scalable_free(a_mem); + void Deallocate(RE::MemoryManager *, void *a_mem, bool a_alignmentRequired) { + if (a_mem != g_trash) { + if (a_mem) { + std::atomic *original_mem = (std::atomic *) (((uint64_t) a_mem) - 0x100); + uint64_t expected = 0x1337DEADDEAD1337; + if (original_mem->compare_exchange_strong(expected, (uint64_t) 0x0)) { + //logger::trace("MemoryManagerStats::Deallocate ptr {} (internal ptr {}) diff {}.", reinterpret_cast(a_mem), reinterpret_cast(original_mem), reinterpret_cast(a_mem) - reinterpret_cast(original_mem)); + a_alignmentRequired ? scalable_aligned_free(original_mem) : scalable_free(original_mem); + } + else { + logger::warn("MemoryManager::Deallocate error 202: already deallocated for addr {}.", reinterpret_cast(a_mem)); + } + } + } } - void* Reallocate(RE::MemoryManager* a_self, void* a_oldMem, std::size_t a_newSize, std::uint32_t a_alignment, bool a_alignmentRequired) - { + void *Reallocate(RE::MemoryManager *a_self, void *a_oldMem, std::size_t a_newSize, std::uint32_t a_alignment, + bool a_alignmentRequired) { if (a_oldMem == g_trash) return Allocate(a_self, a_newSize, a_alignment, a_alignmentRequired); - else - return a_alignmentRequired ? - scalable_aligned_realloc(a_oldMem, a_newSize, a_alignment) : - scalable_realloc(a_oldMem, a_newSize); + else if (a_oldMem) { + void *original_mem = (std::atomic *) (((uint64_t) a_oldMem) - 0x100); + + //logger::trace("MemoryManagerStats::Reallocate ptr {} (internal ptr {}) diff {}.", reinterpret_cast(a_oldMem), reinterpret_cast(original_mem), reinterpret_cast(a_oldMem) - reinterpret_cast(original_mem)); + + std::atomic *original_cookie = (std::atomic *) (((uint64_t) a_oldMem) - 0x100); + uint64_t expected = 0x1337DEADDEAD1337; + if (original_cookie->compare_exchange_strong(expected, (uint64_t) 0x1337DEADDEAD1337)) { + void *new_mem = a_alignmentRequired ? scalable_aligned_realloc(original_mem, a_newSize + 0x100, a_alignment) + : scalable_realloc(original_mem, a_newSize + 0x100); + *(uint64_t *) new_mem = (uint64_t) 0x1337DEADDEAD1337; + return (void *) (((uint64_t) new_mem) + 0x100); + } + else + { + logger::warn("MemoryManager::Reallocate error 302: already deallocated for addr {}.", reinterpret_cast(a_oldMem)); + return g_trash; + } + } + return g_trash; } - void ReplaceAllocRoutines() { using tuple_t = std::tuple; @@ -158,9 +182,9 @@ namespace { void* Allocate(RE::ScrapHeap*, std::size_t a_size, std::size_t a_alignment) { - return a_size > 0 ? - scalable_aligned_malloc(a_size, a_alignment) : - g_trash; + // TODO: std::size_t -> std::uint32_t + //logger::trace("ScrapHeap::Allocate a_size: {}; a_alignment: {}; static_cast a_alignment: {}.", a_size, a_alignment, static_cast(a_alignment)); + return MemoryManager::Allocate(nullptr, a_size, static_cast(a_alignment), true); } RE::ScrapHeap* Ctor(RE::ScrapHeap* a_this) @@ -172,8 +196,8 @@ namespace void Deallocate(RE::ScrapHeap*, void* a_mem) { - if (a_mem != g_trash) - scalable_aligned_free(a_mem); + //logger::trace("ScrapHeap::Deallocate {}.", reinterpret_cast(a_mem)); + return MemoryManager::Deallocate(nullptr,a_mem,true); } void WriteHooks()