// Copyright Epic Games, Inc. All Rights Reserved.

#include <zencore/memory/mallocstomp.h>

#if ZEN_WITH_MALLOC_STOMP

#	include <zencore/memory/align.h>
#	include <zencore/xxhash.h>

#	if ZEN_PLATFORM_LINUX
#		include <sys/mman.h>
#	endif

#	if ZEN_PLATFORM_WINDOWS
#		include <zencore/windows.h>
#	endif

#	if ZEN_PLATFORM_WINDOWS
// MallocStomp can keep virtual address range reserved after memory block is freed, while releasing the physical memory.
// This dramatically increases accuracy of use-after-free detection, but consumes significant amount of memory for the OS page table.
// Virtual memory limit for a process on Win10 is 128 TB, which means we can afford to keep virtual memory reserved for a very long time.
// Running Infiltrator demo consumes ~700MB of virtual address space per second.
#		define MALLOC_STOMP_KEEP_VIRTUAL_MEMORY 1
#	else
#		define MALLOC_STOMP_KEEP_VIRTUAL_MEMORY 0
#	endif

// 64-bit ABIs on x86_64 expect a 16-byte alignment
#	define STOMPALIGNMENT 16U

namespace zen {

struct FMallocStomp::FAllocationData
{
	/** Pointer to the full allocation. Needed so the OS knows what to free. */
	void* FullAllocationPointer;
	/** Full size of the allocation including the extra page. */
	size_t FullSize;
	/** Size of the allocation requested. */
	size_t Size;
	/** Sentinel used to check for underrun. */
	size_t Sentinel;

	/** Calculate the expected sentinel value for this allocation data. */
	size_t CalculateSentinel() const
	{
		XXH3_128 Xxh = XXH3_128::HashMemory(this, offsetof(FAllocationData, Sentinel));

		size_t Hash;
		memcpy(&Hash, Xxh.Hash, sizeof(Hash));

		return Hash;
	}
};

FMallocStomp::FMallocStomp(const bool InUseUnderrunMode) : PageSize(4096 /* TODO: make dynamic */), bUseUnderrunMode(InUseUnderrunMode)
{
}

void*
FMallocStomp::Malloc(size_t Size, uint32_t Alignment)
{
	void* Result = TryMalloc(Size, Alignment);

	if (Result == nullptr)
	{
		OutOfMemory(Size, Alignment);
	}

	return Result;
}

void*
FMallocStomp::TryMalloc(size_t Size, uint32_t Alignment)
{
	if (Size == 0U)
	{
		Size = 1U;
	}

	Alignment = Max<uint32_t>(Alignment, STOMPALIGNMENT);

	constexpr static size_t AllocationDataSize = sizeof(FAllocationData);

	const size_t AlignedSize		 = Alignment ? ((Size + Alignment - 1) & -(int32_t)Alignment) : Size;
	const size_t AlignmentSize		 = Alignment > PageSize ? Alignment - PageSize : 0;
	const size_t AllocFullPageSize	 = (AlignedSize + AlignmentSize + AllocationDataSize + PageSize - 1) & ~(PageSize - 1);
	const size_t TotalAllocationSize = AllocFullPageSize + PageSize;

#	if ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
	void* FullAllocationPointer = mmap(nullptr, TotalAllocationSize, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANON, -1, 0);
#	elif ZEN_PLATFORM_WINDOWS && MALLOC_STOMP_KEEP_VIRTUAL_MEMORY
	// Allocate virtual address space from current block using linear allocation strategy.
	// If there is not enough space, try to allocate new block from OS. Report OOM if block allocation fails.
	void* FullAllocationPointer = nullptr;

	{
		RwLock::ExclusiveLockScope _(Lock);

		if (VirtualAddressCursor + TotalAllocationSize <= VirtualAddressMax)
		{
			FullAllocationPointer = (void*)(VirtualAddressCursor);
		}
		else
		{
			const size_t ReserveSize = Max(VirtualAddressBlockSize, TotalAllocationSize);

			// Reserve a new block of virtual address space that will be linearly sub-allocated
			// We intentionally don't keep track of reserved blocks, as we never need to explicitly release them.
			FullAllocationPointer = VirtualAlloc(nullptr, ReserveSize, MEM_RESERVE, PAGE_NOACCESS);

			VirtualAddressCursor = uintptr_t(FullAllocationPointer);
			VirtualAddressMax	 = VirtualAddressCursor + ReserveSize;
		}

		VirtualAddressCursor += TotalAllocationSize;
	}
#	else
	void* FullAllocationPointer = FPlatformMemory::BinnedAllocFromOS(TotalAllocationSize);
#	endif	// PLATFORM_UNIX || PLATFORM_MAC

	if (!FullAllocationPointer)
	{
		return nullptr;
	}

	void* ReturnedPointer = nullptr;

	ZEN_ASSERT_SLOW(IsAligned(FullAllocationPointer, PageSize));

	if (bUseUnderrunMode)
	{
		ReturnedPointer				= Align((uint8_t*)FullAllocationPointer + PageSize + AllocationDataSize, Alignment);
		void* AllocDataPointerStart = static_cast<FAllocationData*>(ReturnedPointer) - 1;
		ZEN_ASSERT_SLOW(AllocDataPointerStart >= FullAllocationPointer);

#	if ZEN_PLATFORM_WINDOWS && MALLOC_STOMP_KEEP_VIRTUAL_MEMORY
		// Commit physical pages to the used range, leaving the first page unmapped.
		void* CommittedMemory = VirtualAlloc(AllocDataPointerStart, AllocationDataSize + AlignedSize, MEM_COMMIT, PAGE_READWRITE);
		if (!CommittedMemory)
		{
			// Failed to allocate and commit physical memory pages.
			return nullptr;
		}
		ZEN_ASSERT(CommittedMemory == AlignDown(AllocDataPointerStart, PageSize));
#	else
		// Page protect the first page, this will cause the exception in case there is an underrun.
		FPlatformMemory::PageProtect((uint8*)AlignDown(AllocDataPointerStart, PageSize) - PageSize, PageSize, false, false);
#	endif
	}  //-V773
	else
	{
		ReturnedPointer			 = AlignDown((uint8_t*)FullAllocationPointer + AllocFullPageSize - AlignedSize, Alignment);
		void* ReturnedPointerEnd = (uint8_t*)ReturnedPointer + AlignedSize;
		ZEN_ASSERT_SLOW(IsAligned(ReturnedPointerEnd, PageSize));

		void* AllocDataPointerStart = static_cast<FAllocationData*>(ReturnedPointer) - 1;
		ZEN_ASSERT_SLOW(AllocDataPointerStart >= FullAllocationPointer);

#	if ZEN_PLATFORM_WINDOWS && MALLOC_STOMP_KEEP_VIRTUAL_MEMORY
		// Commit physical pages to the used range, leaving the last page unmapped.
		void* CommitPointerStart = AlignDown(AllocDataPointerStart, PageSize);
		void* CommittedMemory	 = VirtualAlloc(CommitPointerStart,
												size_t((uint8_t*)ReturnedPointerEnd - (uint8_t*)CommitPointerStart),
												MEM_COMMIT,
												PAGE_READWRITE);
		if (!CommittedMemory)
		{
			// Failed to allocate and commit physical memory pages.
			return nullptr;
		}
		ZEN_ASSERT(CommittedMemory == CommitPointerStart);
#	else
		// Page protect the last page, this will cause the exception in case there is an overrun.
		FPlatformMemory::PageProtect(ReturnedPointerEnd, PageSize, false, false);
#	endif
	}  //-V773

	ZEN_ASSERT_SLOW(IsAligned(FullAllocationPointer, PageSize));
	ZEN_ASSERT_SLOW(IsAligned(TotalAllocationSize, PageSize));
	ZEN_ASSERT_SLOW(IsAligned(ReturnedPointer, Alignment));
	ZEN_ASSERT_SLOW((uint8_t*)ReturnedPointer + AlignedSize <= (uint8_t*)FullAllocationPointer + TotalAllocationSize);

	FAllocationData& AllocationData = static_cast<FAllocationData*>(ReturnedPointer)[-1];
	AllocationData					= {FullAllocationPointer, TotalAllocationSize, AlignedSize, 0};
	AllocationData.Sentinel			= AllocationData.CalculateSentinel();

	return ReturnedPointer;
}

void*
FMallocStomp::Realloc(void* InPtr, size_t NewSize, uint32_t Alignment)
{
	void* Result = TryRealloc(InPtr, NewSize, Alignment);

	if (Result == nullptr && NewSize)
	{
		OutOfMemory(NewSize, Alignment);
	}

	return Result;
}

void*
FMallocStomp::TryRealloc(void* InPtr, size_t NewSize, uint32_t Alignment)
{
	if (NewSize == 0U)
	{
		Free(InPtr);
		return nullptr;
	}

	void* ReturnPtr = nullptr;

	if (InPtr != nullptr)
	{
		ReturnPtr = TryMalloc(NewSize, Alignment);

		if (ReturnPtr != nullptr)
		{
			FAllocationData* AllocDataPtr = reinterpret_cast<FAllocationData*>(reinterpret_cast<uint8_t*>(InPtr) - sizeof(FAllocationData));
			memcpy(ReturnPtr, InPtr, Min(AllocDataPtr->Size, NewSize));
			Free(InPtr);
		}
	}
	else
	{
		ReturnPtr = TryMalloc(NewSize, Alignment);
	}

	return ReturnPtr;
}

void
FMallocStomp::Free(void* InPtr)
{
	if (InPtr == nullptr)
	{
		return;
	}

	FAllocationData* AllocDataPtr = reinterpret_cast<FAllocationData*>(InPtr);
	AllocDataPtr--;

	// Check the sentinel to verify that the allocation data is intact.
	if (AllocDataPtr->Sentinel != AllocDataPtr->CalculateSentinel())
	{
		// There was a memory underrun related to this allocation.
		ZEN_DEBUG_BREAK();
	}

#	if ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
	munmap(AllocDataPtr->FullAllocationPointer, AllocDataPtr->FullSize);
#	elif ZEN_PLATFORM_WINDOWS && MALLOC_STOMP_KEEP_VIRTUAL_MEMORY
	// Unmap physical memory, but keep virtual address range reserved to catch use-after-free errors.

	VirtualFree(AllocDataPtr->FullAllocationPointer, AllocDataPtr->FullSize, MEM_DECOMMIT);

#	else
	FPlatformMemory::BinnedFreeToOS(AllocDataPtr->FullAllocationPointer, AllocDataPtr->FullSize);
#	endif	// PLATFORM_UNIX || PLATFORM_MAC
}

bool
FMallocStomp::GetAllocationSize(void* Original, size_t& SizeOut)
{
	if (Original == nullptr)
	{
		SizeOut = 0U;
	}
	else
	{
		FAllocationData* AllocDataPtr = reinterpret_cast<FAllocationData*>(Original);
		AllocDataPtr--;
		SizeOut = AllocDataPtr->Size;
	}

	return true;
}

}  // namespace zen

#endif	// WITH_MALLOC_STOMP