feat: _set_allocator

kext/kmalloc.h

@@ -1,6 +1,6 @@
 /*
  * PROJECT:   Universal C++ RunTime (UCXXRT)
- * FILE:      malloc_km.h
+ * FILE:      kmalloc.h
  * DATA:      2022/06/17
  *
  * PURPOSE:   Universal C++ RunTime
@@ -17,7 +17,7 @@
 extern "C" _CRT_HYBRIDPATCHABLE __declspec(noinline) _CRTRESTRICT
 void* __cdecl kmalloc(
     _In_ _CRT_GUARDOVERFLOW size_t size,
-    _In_ int pool_type,
+    _In_ int pool,
     _In_ unsigned long tag
 );
 
@@ -31,6 +31,23 @@ extern "C" _CRT_HYBRIDPATCHABLE __declspec(noinline) _CRTRESTRICT
 void* __cdecl kcalloc(
     _In_ _CRT_GUARDOVERFLOW size_t const count,
     _In_ _CRT_GUARDOVERFLOW size_t const size,
-    _In_ int          pool_type,
+    _In_ int pool,
     _In_ unsigned long tag
 );
+
+extern "C" _CRT_HYBRIDPATCHABLE __declspec(noinline) _CRTRESTRICT
+void* __cdecl krealloc(
+    void* const block,
+    size_t const size,
+    int pool,
+    unsigned long tag
+);
+
+extern "C" __declspec(noinline) _CRTRESTRICT
+void* __cdecl krecalloc(
+    void* const block,
+    size_t const count,
+    size_t const size,
+    int pool,
+    unsigned long tag
+);

kext/knew.h

@@ -1,6 +1,6 @@
 /*
  * PROJECT:   Universal C++ RunTime (UCXXRT)
- * FILE:      new_km.h
+ * FILE:      knew.h
  * DATA:      2022/06/17
  *
  * PURPOSE:   Universal C++ RunTime

msvc/ucxxrt.vcxproj

@@ -341,17 +341,18 @@
     </ClCompile>
     <ClCompile Include="..\src\ucrt\heap\align.cpp" />
     <ClCompile Include="..\src\ucrt\heap\calloc.cpp" />
-    <ClCompile Include="..\src\ucrt\heap\kcalloc.cpp" />
     <ClCompile Include="..\src\ucrt\heap\expand.cpp" />
     <ClCompile Include="..\src\ucrt\heap\free.cpp" />
     <ClCompile Include="..\src\ucrt\heap\kfree.cpp" />
     <ClCompile Include="..\src\ucrt\heap\malloc.cpp" />
     <ClCompile Include="..\src\ucrt\heap\kmalloc.cpp" />
+    <ClCompile Include="..\src\ucrt\heap\moverride.cpp" />
     <ClCompile Include="..\src\ucrt\heap\msize.cpp" />
     <ClCompile Include="..\src\ucrt\heap\new_handler.cpp" />
     <ClCompile Include="..\src\ucrt\heap\new_mode.cpp" />
     <ClCompile Include="..\src\ucrt\heap\realloc.cpp" />
     <ClCompile Include="..\src\ucrt\heap\recalloc.cpp" />
+    <ClCompile Include="..\src\ucrt\heap\stdalloc.cpp" />
     <ClCompile Include="..\src\ucrt\internal\initialization.cpp" />
     <ClCompile Include="..\src\ucrt\internal\locks.cpp" />
     <ClCompile Include="..\src\ucrt\internal\per_thread_data.cpp" />

msvc/ucxxrt.vcxproj.filters

@@ -671,15 +671,18 @@
     <ClCompile Include="..\src\ucrt\heap\kfree.cpp">
       <Filter>ucxxrt\ucrt\heap</Filter>
     </ClCompile>
-    <ClCompile Include="..\src\ucrt\heap\kcalloc.cpp">
-      <Filter>ucxxrt\ucrt\heap</Filter>
-    </ClCompile>
     <ClCompile Include="..\src\crt\vcruntime\delete_km.cpp">
       <Filter>ucxxrt\crt\vcruntime</Filter>
     </ClCompile>
     <ClCompile Include="..\src\crt\vcruntime\new_km.cpp">
       <Filter>ucxxrt\crt\vcruntime</Filter>
     </ClCompile>
+    <ClCompile Include="..\src\ucrt\heap\moverride.cpp">
+      <Filter>ucxxrt\ucrt\heap</Filter>
+    </ClCompile>
+    <ClCompile Include="..\src\ucrt\heap\stdalloc.cpp">
+      <Filter>ucxxrt\ucrt\heap</Filter>
+    </ClCompile>
   </ItemGroup>
   <ItemGroup>
     <MARMASM Include="..\src\crt\arm\chkstk.asm">

src/crt/vcruntime/internal_shared.h

@@ -19,9 +19,12 @@ _CRT_BEGIN_C_HEADER
 
 
 
-extern ULONG __ucxxrt_tag;
-extern IMAGE_DOS_HEADER __ImageBase;
+extern ULONG                __ucxxrt_tag;
+extern IMAGE_DOS_HEADER     __ImageBase;
 
+extern volatile _malloc_t   __override_malloc;
+extern volatile _free_t     __override_free;
+extern volatile _msize_t    __override_msize;
 
 
 //-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+

src/crt/vcruntime/utility.cpp

@@ -147,6 +147,8 @@ extern "C" void __cdecl __scrt_initialize_memory()
 {
     // Nx
     ExInitializeDriverRuntime(DrvRtPoolNxOptIn);
+
+    _set_allocator(nullptr, nullptr, nullptr);
 }
 
 extern "C" bool __cdecl __scrt_initialize_crt()

src/ucrt/heap/free.cpp

@@ -17,8 +17,5 @@
 // with either other function or vice versa.
 extern "C" _CRT_HYBRIDPATCHABLE __declspec(noinline) void __cdecl free(void* const block)
 {
-    if (block)
-    {
-        ExFreePoolWithTag(block, __ucxxrt_tag);
-    }
+    return __override_free(block);
 }

src/ucrt/heap/kfree.cpp

@@ -1,6 +1,6 @@
 /*
  * PROJECT:   Universal C++ RunTime (UCXXRT)
- * FILE:      free_km.cpp
+ * FILE:      kfree.cpp
  * DATA:      2022/06/17
  *
  * PURPOSE:   Universal C++ RunTime

src/ucrt/heap/kmalloc.cpp

@@ -1,6 +1,6 @@
 /*
  * PROJECT:   Universal C++ RunTime (UCXXRT)
- * FILE:      malloc_km.cpp
+ * FILE:      kmalloc.cpp
  * DATA:      2022/06/17
  *
  * PURPOSE:   Universal C++ RunTime
@@ -15,18 +15,39 @@
 #include <new.h>
 
 
+extern"C" __declspec(noinline) void* __cdecl ExReallocatePoolWithTag(
+    _In_ SIZE_T OldSize,
+    _In_ SIZE_T NewSize,
+    _In_ PVOID  OldBlock,
+    _In_ __drv_strictTypeMatch(__drv_typeExpr) POOL_TYPE PoolType,
+    _In_ ULONG Tag
+)
+{
+    if (OldSize == 0)
+    {
+        return nullptr;
+    }
+
+    #pragma warning(suppress: 4996)
+    void* const NewBlock = ExAllocatePoolWithTag(PoolType, NewSize, Tag);
+    if (NewBlock)
+    {
+        memset (NewBlock, 0, NewSize);
+        memmove(NewBlock, OldBlock, OldSize);
+
+        ExFreePoolWithTag(OldBlock, Tag);
+        return NewBlock;
+    }
+
+    return nullptr;
+}
+
+extern "C" _CRT_HYBRIDPATCHABLE __declspec(noinline) void __cdecl kfree(void* const block, unsigned long tag);
 
-// Allocates a block of memory of size 'size' bytes in the heap.  If allocation
-// fails, nullptr is returned.
-//
-// This function supports patching and therefore must be marked noinline.
-// Both _kmalloc_dbg and _kmalloc_base must also be marked noinline
-// to prevent identical COMDAT folding from substituting calls to malloc
-// with either other function or vice versa.
 extern "C" _CRT_HYBRIDPATCHABLE __declspec(noinline) _CRTRESTRICT void* __cdecl kmalloc(
-    size_t      const   size,
-    int                 pool_type,
-    unsigned    long    tag
+    size_t const size,
+    int pool,
+    unsigned long tag
 )
 {
     // Ensure that the requested size is not too large:
@@ -38,7 +59,7 @@ extern "C" _CRT_HYBRIDPATCHABLE __declspec(noinline) _CRTRESTRICT void* __cdecl
     for (;;)
     {
         #pragma warning(suppress: 4996)
-        void* const block = ExAllocatePoolWithTag((POOL_TYPE)pool_type, actual_size, tag);
+        void* const block = ExAllocatePoolWithTag((POOL_TYPE)pool, actual_size, tag);
         if (block)
             return block;
 
@@ -53,3 +74,109 @@ extern "C" _CRT_HYBRIDPATCHABLE __declspec(noinline) _CRTRESTRICT void* __cdecl
         // The new handler was successful; try to allocate again...
     }
 }
+
+extern "C" _CRT_HYBRIDPATCHABLE __declspec(noinline) _CRTRESTRICT void* __cdecl kcalloc(
+    size_t const count,
+    size_t const size,
+    int pool,
+    unsigned long tag
+)
+{
+    // Ensure that (count * size) does not overflow
+    _VALIDATE_RETURN_NOEXC(count == 0 || (_HEAP_MAXREQ / count) >= size, ENOMEM, nullptr);
+
+    // Ensure that we allocate a nonzero block size:
+    size_t const requested_block_size = count * size;
+    size_t const actual_block_size = requested_block_size == 0
+        ? 1
+        : requested_block_size;
+
+    for (;;)
+    {
+        void* const block = kmalloc(actual_block_size, pool, tag);
+
+        // If allocation succeeded, return the pointer to the new block:
+        if (block)
+        {
+            memset(block, 0, actual_block_size);
+            return block;
+        }
+
+        // Otherwise, see if we need to call the new handler, and if so call it.
+        // If the new handler fails, just return nullptr:
+        if (_query_new_mode() == 0 || !_callnewh(actual_block_size))
+        {
+            errno = ENOMEM;
+            return nullptr;
+        }
+
+        // The new handler was successful; try to allocate aagain...
+    }
+}
+
+extern "C" _CRT_HYBRIDPATCHABLE __declspec(noinline) _CRTRESTRICT void* __cdecl krealloc(
+    void* const block,
+    size_t const size,
+    int pool,
+    unsigned long tag
+)
+{
+    // If the block is a nullptr, just call malloc:
+    if (block == nullptr)
+        return kmalloc(size, pool, tag);
+
+    // If the new size is 0, just call free and return nullptr:
+    if (size == 0)
+    {
+        kfree(block, tag);
+        return nullptr;
+    }
+
+    // Ensure that the requested size is not too large:
+    _VALIDATE_RETURN_NOEXC(_HEAP_MAXREQ >= size, ENOMEM, nullptr);
+
+    for (;;)
+    {
+        void* const new_block = ExReallocatePoolWithTag(_msize(block), size, block, NonPagedPool, __ucxxrt_tag);
+        if (new_block)
+        {
+            return new_block;
+        }
+
+        // Otherwise, see if we need to call the new handler, and if so call it.
+        // If the new handler fails, just return nullptr:
+        if (_query_new_mode() == 0 || !_callnewh(size))
+        {
+            errno = ENOMEM;
+            return nullptr;
+        }
+
+        // The new handler was successful; try to allocate again...
+    }
+}
+
+extern "C" __declspec(noinline) _CRTRESTRICT void* __cdecl krecalloc(
+    void* const block,
+    size_t const count,
+    size_t const size,
+    int pool,
+    unsigned long tag
+)
+{
+    // Ensure that (count * size) does not overflow
+    _VALIDATE_RETURN_NOEXC(count == 0 || (_HEAP_MAXREQ / count) >= size, ENOMEM, nullptr);
+
+    size_t const old_block_size = block != nullptr ? _msize(block) : 0;
+    size_t const new_block_size = count * size;
+
+    void* const new_block = krealloc(block, new_block_size, pool, tag);
+
+    // If the reallocation succeeded and the new block is larger, zero-fill the
+    // new bytes:
+    if (new_block != nullptr && old_block_size < new_block_size)
+    {
+        memset(static_cast<char*>(new_block) + old_block_size, 0, new_block_size - old_block_size);
+    }
+
+    return new_block;
+}

src/ucrt/heap/malloc.cpp

@@ -28,8 +28,7 @@ extern "C" _CRT_HYBRIDPATCHABLE __declspec(noinline) _CRTRESTRICT void* __cdecl
 
     for (;;)
     {
-        #pragma warning(suppress: 4996)
-        void* const block = ExAllocatePoolWithTag(NonPagedPool, actual_size, __ucxxrt_tag);
+        void* const block = __override_malloc(size);
         if (block)
             return block;