Code/Engine/Foundation/Platform/Win/GuardedAllocation_Win.cpp - dev - redgpu/ezEngine

Code/Engine/Foundation/Platform/Win/GuardedAllocation_Win.cpp
110 строк4 KB
Jan Krassnigg
Refactor of platform specific includes (#1414)
23 окт 2024, 10:09
23 окт 2024, 10:0974f6d97
#include <Foundation/FoundationPCH.h>

#if EZ_ENABLED(EZ_PLATFORM_WINDOWS)

#  include <Foundation/Memory/Policies/AllocPolicyGuarding.h>

#  include <Foundation/Platform/Win/Utils/IncludeWindows.h>

struct AlloctionMetaData
{
  AlloctionMetaData()
  {
    m_uiSize = 0;

    for (ezUInt32 i = 0; i < EZ_ARRAY_SIZE(m_magic); ++i)
    {
      m_magic[i] = 0x12345678;
    }
  }

  ~AlloctionMetaData()
  {
    for (ezUInt32 i = 0; i < EZ_ARRAY_SIZE(m_magic); ++i)
    {
      EZ_ASSERT_DEV(m_magic[i] == 0x12345678, "Magic value has been overwritten. This might be the result of a buffer underrun!");
    }
  }

  size_t m_uiSize;
  ezUInt32 m_magic[32];
};

ezAllocPolicyGuarding::ezAllocPolicyGuarding(ezAllocator* pParent)
{
  EZ_IGNORE_UNUSED(pParent);

  SYSTEM_INFO sysInfo;
  GetSystemInfo(&sysInfo);
  m_uiPageSize = sysInfo.dwPageSize;
}


void* ezAllocPolicyGuarding::Allocate(size_t uiSize, size_t uiAlign)
{
  EZ_ASSERT_DEV(ezMath::IsPowerOf2((ezUInt32)uiAlign), "Alignment must be power of two");
  uiAlign = ezMath::Max<size_t>(uiAlign, EZ_ALIGNMENT_MINIMUM);

  size_t uiAlignedSize = ezMemoryUtils::AlignSize(uiSize, uiAlign);
  size_t uiTotalSize = uiAlignedSize + sizeof(AlloctionMetaData);

  // align to full pages and add one page in front and one in back
  size_t uiPageSize = m_uiPageSize;
  size_t uiFullPageSize = ezMemoryUtils::AlignSize(uiTotalSize, uiPageSize);
  void* pMemory = VirtualAlloc(nullptr, uiFullPageSize + 2 * uiPageSize, MEM_RESERVE, PAGE_NOACCESS);
  EZ_ASSERT_DEV(pMemory != nullptr, "Could not reserve memory pages. Error Code '{0}'", ezArgErrorCode(::GetLastError()));

  // add one page and commit the payload pages
  pMemory = ezMemoryUtils::AddByteOffset(pMemory, uiPageSize);
  void* ptr = VirtualAlloc(pMemory, uiFullPageSize, MEM_COMMIT, PAGE_READWRITE);
  EZ_ASSERT_DEV(ptr != nullptr, "Could not commit memory pages. Error Code '{0}'", ezArgErrorCode(::GetLastError()));

  // store information in meta data
  AlloctionMetaData* metaData = ezMemoryUtils::AddByteOffset(static_cast<AlloctionMetaData*>(ptr), uiFullPageSize - uiTotalSize);
  ezMemoryUtils::Construct<SkipTrivialTypes>(metaData, 1);
  metaData->m_uiSize = uiAlignedSize;

  // finally add offset to the actual payload
  ptr = ezMemoryUtils::AddByteOffset(metaData, sizeof(AlloctionMetaData));
  return ptr;
}

// deactivate analysis warning for VirtualFree flags, it is needed for the specific functionality
EZ_MSVC_ANALYSIS_WARNING_PUSH
EZ_MSVC_ANALYSIS_WARNING_DISABLE(6250)

void ezAllocPolicyGuarding::Deallocate(void* pPtr)
{
  ezLock<ezMutex> lock(m_Mutex);

  if (!m_AllocationsToFreeLater.CanAppend())
  {
    void* pMemory = m_AllocationsToFreeLater.PeekFront();
    EZ_VERIFY(::VirtualFree(pMemory, 0, MEM_RELEASE), "Could not free memory pages. Error Code '{0}'", ezArgErrorCode(::GetLastError()));

    m_AllocationsToFreeLater.PopFront();
  }

  // Retrieve info from meta data first.
  AlloctionMetaData* metaData = ezMemoryUtils::AddByteOffset(static_cast<AlloctionMetaData*>(pPtr), -((std::ptrdiff_t)sizeof(AlloctionMetaData)));
  size_t uiAlignedSize = metaData->m_uiSize;

  ezMemoryUtils::Destruct(metaData, 1);

  // Decommit the pages but do not release the memory yet so use-after-free can be detected.
  size_t uiPageSize = m_uiPageSize;
  size_t uiTotalSize = uiAlignedSize + sizeof(AlloctionMetaData);
  size_t uiFullPageSize = ezMemoryUtils::AlignSize(uiTotalSize, uiPageSize);
  pPtr = ezMemoryUtils::AddByteOffset(pPtr, ((std::ptrdiff_t)uiAlignedSize) - uiFullPageSize);

  EZ_VERIFY(
    ::VirtualFree(pPtr, uiFullPageSize, MEM_DECOMMIT), "Could not decommit memory pages. Error Code '{0}'", ezArgErrorCode(::GetLastError()));

  // Finally store the allocation so we can release it later
  void* pMemory = ezMemoryUtils::AddByteOffset(pPtr, -((std::ptrdiff_t)uiPageSize));
  m_AllocationsToFreeLater.PushBack(pMemory);
}

EZ_MSVC_ANALYSIS_WARNING_POP

#endif