api: use a logging alloc function

Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>

8 files changed, 96 insertions, 131 deletions

diff --git a/api/adapter.c b/api/adapter.c
index 6a0af9d..12b02fe 100644
--- a/api/adapter.c
+++ b/api/adapter.c
@@ -61,18 +61,14 @@ static _Return_type_success_(return != NULL) SP_DRVINFO_DETAIL_DATA_W *GetAdapte
     DWORD Size = sizeof(SP_DRVINFO_DETAIL_DATA_W) + 0x100;
     for (;;)
     {
-        SP_DRVINFO_DETAIL_DATA_W *DrvInfoDetailData = HeapAlloc(ModuleHeap, 0, Size);
+        SP_DRVINFO_DETAIL_DATA_W *DrvInfoDetailData = Alloc(Size);
         if (!DrvInfoDetailData)
-        {
-            LOG(WINTUN_LOG_ERR, L"Out of memory");
-            SetLastError(ERROR_OUTOFMEMORY);
             return NULL;
-        }
         DrvInfoDetailData->cbSize = sizeof(SP_DRVINFO_DETAIL_DATA_W);
         if (SetupDiGetDriverInfoDetailW(DevInfo, DevInfoData, DrvInfoData, DrvInfoDetailData, Size, &Size))
             return DrvInfoDetailData;
         DWORD LastError = GetLastError();
-        HeapFree(ModuleHeap, 0, DrvInfoDetailData);
+        Free(DrvInfoDetailData);
         if (LastError != ERROR_INSUFFICIENT_BUFFER)
         {
             SetLastError(LOG_ERROR(L"Failed", LastError));
@@ -90,17 +86,13 @@ static _Return_type_success_(return != NULL) void *GetDeviceRegistryProperty(
 {
     for (;;)
     {
-        BYTE *Data = HeapAlloc(ModuleHeap, 0, *BufLen);
+        BYTE *Data = Alloc(*BufLen);
         if (!Data)
-        {
-            LOG(WINTUN_LOG_ERR, L"Out of memory");
-            SetLastError(ERROR_OUTOFMEMORY);
             return NULL;
-        }
         if (SetupDiGetDeviceRegistryPropertyW(DevInfo, DevInfoData, Property, ValueType, Data, *BufLen, BufLen))
             return Data;
         DWORD LastError = GetLastError();
-        HeapFree(ModuleHeap, 0, Data);
+        Free(Data);
         if (LastError != ERROR_INSUFFICIENT_BUFFER)
         {
             SetLastError(LOG_ERROR(L"Querying property failed", LastError));
@@ -129,7 +121,7 @@ static _Return_type_success_(return != NULL)
         LOG(WINTUN_LOG_ERR, L"Property is not a string");
         LastError = ERROR_INVALID_DATATYPE;
     }
-    HeapFree(ModuleHeap, 0, Buf);
+    Free(Buf);
     SetLastError(LastError);
     return NULL;
 }
@@ -154,7 +146,7 @@ static _Return_type_success_(return != NULL)
         LOG(WINTUN_LOG_ERR, L"Property is not a string");
         LastError = ERROR_INVALID_DATATYPE;
     }
-    HeapFree(ModuleHeap, 0, Buf);
+    Free(Buf);
     SetLastError(LastError);
     return NULL;
 }
@@ -178,7 +170,7 @@ IsOurAdapter(_In_ HDEVINFO DevInfo, _In_ SP_DEVINFO_DATA *DevInfoData)
         return FALSE;
     }
     BOOL IsOurs = IsOurHardwareID(Hwids);
-    HeapFree(ModuleHeap, 0, Hwids);
+    Free(Hwids);
     return IsOurs;
 }
 
@@ -197,13 +189,9 @@ static _Return_type_success_(return != INVALID_HANDLE_VALUE) HANDLE OpenDeviceOb
         SetLastError(LOG_ERROR(L"Failed to query associated instances size", LastError));
         return INVALID_HANDLE_VALUE;
     }
-    WCHAR *Interfaces = HeapAlloc(ModuleHeap, 0, InterfacesLen * sizeof(WCHAR));
+    WCHAR *Interfaces = Alloc(InterfacesLen * sizeof(WCHAR));
     if (!Interfaces)
-    {
-        LOG(WINTUN_LOG_ERR, L"Out of memory");
-        SetLastError(ERROR_OUTOFMEMORY);
         return INVALID_HANDLE_VALUE;
-    }
     HANDLE Handle = INVALID_HANDLE_VALUE;
     LastError = CM_MapCrToWin32Err(
         CM_Get_Device_Interface_ListW(
@@ -228,7 +216,7 @@ static _Return_type_success_(return != INVALID_HANDLE_VALUE) HANDLE OpenDeviceOb
         NULL);
     LastError = Handle != INVALID_HANDLE_VALUE ? ERROR_SUCCESS : LOG_LAST_ERROR(L"Failed to connect to adapter");
 cleanupBuf:
-    HeapFree(ModuleHeap, 0, Interfaces);
+    Free(Interfaces);
     if (LastError != ERROR_SUCCESS)
         SetLastError(LastError);
     return Handle;
@@ -247,13 +235,9 @@ static _Return_type_success_(return != FALSE) BOOL
         LOG_ERROR(L"Failed to query instance ID size", LastError);
         return FALSE;
     }
-    WCHAR *InstanceId = HeapAlloc(ModuleHeap, HEAP_ZERO_MEMORY, sizeof(*InstanceId) * RequiredBytes);
+    WCHAR *InstanceId = Zalloc(sizeof(*InstanceId) * RequiredBytes);
     if (!InstanceId)
-    {
-        LOG(WINTUN_LOG_ERR, L"Out of memory");
-        SetLastError(ERROR_OUTOFMEMORY);
         return FALSE;
-    }
     if (!SetupDiGetDeviceInstanceIdW(DevInfo, DevInfoData, InstanceId, RequiredBytes, &RequiredBytes))
     {
         LastError = LOG_LAST_ERROR(L"Failed to get instance ID");
@@ -276,7 +260,7 @@ static _Return_type_success_(return != FALSE) BOOL
         LastError = LOG_LAST_ERROR(L"Failed to perform ioctl");
     CloseHandle(NdisHandle);
 cleanupInstanceId:
-    HeapFree(ModuleHeap, 0, InstanceId);
+    Free(InstanceId);
     return RET_ERROR(TRUE, LastError);
 }
 
@@ -290,19 +274,15 @@ static _Return_type_success_(return != FALSE) BOOL
     DWORD LastError = ERROR_SUCCESS;
     for (DWORD EnumIndex = 0;; ++EnumIndex)
     {
-        SP_DEVINFO_DATA_LIST *DeviceNode = HeapAlloc(ModuleHeap, 0, sizeof(SP_DEVINFO_DATA_LIST));
+        SP_DEVINFO_DATA_LIST *DeviceNode = Alloc(sizeof(SP_DEVINFO_DATA_LIST));
         if (!DeviceNode)
-        {
-            LOG(WINTUN_LOG_ERR, L"Out of memory");
-            SetLastError(ERROR_OUTOFMEMORY);
             return FALSE;
-        }
         DeviceNode->Data.cbSize = sizeof(SP_DEVINFO_DATA);
         if (!SetupDiEnumDeviceInfo(DevInfo, EnumIndex, &DeviceNode->Data))
         {
             if (GetLastError() == ERROR_NO_MORE_ITEMS)
             {
-                HeapFree(ModuleHeap, 0, DeviceNode);
+                Free(DeviceNode);
                 break;
             }
             goto cleanupDeviceNode;
@@ -333,7 +313,7 @@ static _Return_type_success_(return != FALSE) BOOL
         continue;
 
     cleanupDeviceNode:
-        HeapFree(ModuleHeap, 0, DeviceNode);
+        Free(DeviceNode);
     }
     return RET_ERROR(TRUE, LastError);
 }
@@ -431,7 +411,7 @@ static _Return_type_success_(return != FALSE) BOOL
         LOG(WINTUN_LOG_ERR, L"NetCfgInstanceId is not a GUID");
         LastError = ERROR_INVALID_DATA;
     }
-    HeapFree(ModuleHeap, 0, ValueStr);
+    Free(ValueStr);
 cleanupKey:
     RegCloseKey(Key);
     return RET_ERROR(TRUE, LastError);
@@ -524,9 +504,9 @@ IsPoolMember(_In_z_ const WCHAR *Pool, _In_ HDEVINFO DevInfo, _In_ SP_DEVINFO_DA
         goto cleanupFriendlyName;
     }
 cleanupFriendlyName:
-    HeapFree(ModuleHeap, 0, FriendlyName);
+    Free(FriendlyName);
 cleanupDeviceDesc:
-    HeapFree(ModuleHeap, 0, DeviceDesc);
+    Free(DeviceDesc);
     SetLastError(LastError);
     return Ret;
 }
@@ -543,11 +523,10 @@ static _Return_type_success_(return != NULL) WINTUN_ADAPTER
     }
 
     DWORD LastError;
-    WINTUN_ADAPTER *Adapter = HeapAlloc(ModuleHeap, 0, sizeof(WINTUN_ADAPTER));
+    WINTUN_ADAPTER *Adapter = Alloc(sizeof(WINTUN_ADAPTER));
     if (!Adapter)
     {
-        LOG(WINTUN_LOG_ERR, L"Out of memory");
-        LastError = ERROR_OUTOFMEMORY;
+        LastError = GetLastError();
         goto cleanupKey;
     }
 
@@ -560,11 +539,11 @@ static _Return_type_success_(return != NULL) WINTUN_ADAPTER
     if (FAILED(CLSIDFromString(ValueStr, &Adapter->CfgInstanceID)))
     {
         LOG(WINTUN_LOG_ERR, L"NetCfgInstanceId is not Adapter GUID");
-        HeapFree(ModuleHeap, 0, ValueStr);
+        Free(ValueStr);
         LastError = ERROR_INVALID_DATA;
         goto cleanupAdapter;
     }
-    HeapFree(ModuleHeap, 0, ValueStr);
+    Free(ValueStr);
 
     if (!RegistryQueryDWORD(Key, L"NetLuidIndex", &Adapter->LuidIndex, TRUE))
     {
@@ -596,7 +575,7 @@ static _Return_type_success_(return != NULL) WINTUN_ADAPTER
     return Adapter;
 
 cleanupAdapter:
-    HeapFree(ModuleHeap, 0, Adapter);
+    Free(Adapter);
 cleanupKey:
     RegCloseKey(Key);
     SetLastError(LastError);
@@ -624,7 +603,7 @@ static _Return_type_success_(return != FALSE) BOOL
 void WINAPI
 WintunFreeAdapter(_In_ WINTUN_ADAPTER *Adapter)
 {
-    HeapFree(ModuleHeap, 0, Adapter);
+    Free(Adapter);
 }
 
 _Return_type_success_(return != NULL) WINTUN_ADAPTER *WINAPI
@@ -1022,7 +1001,7 @@ static _Return_type_success_(return != FALSE) BOOL
         goto cleanupPaths;
     }
 cleanupPaths:
-    HeapFree(ModuleHeap, 0, Paths);
+    Free(Paths);
 cleanupTcpipAdapterRegKey:
     RegCloseKey(TcpipAdapterRegKey);
     return RET_ERROR(TRUE, LastError);
@@ -1037,13 +1016,9 @@ static _Return_type_success_(return != 0) DWORD VersionOfFile(_In_z_ const WCHAR
         LOG_LAST_ERROR(L"Failed to query version info size");
         return 0;
     }
-    VOID *VersionInfo = HeapAlloc(ModuleHeap, 0, Len);
+    VOID *VersionInfo = Alloc(Len);
     if (!VersionInfo)
-    {
-        LOG(WINTUN_LOG_ERR, L"Out of memory");
-        SetLastError(ERROR_OUTOFMEMORY);
         return 0;
-    }
     DWORD LastError = ERROR_SUCCESS, Version = 0;
     VS_FIXEDFILEINFO *FixedInfo;
     UINT FixedInfoLen = sizeof(*FixedInfo);
@@ -1064,7 +1039,7 @@ static _Return_type_success_(return != 0) DWORD VersionOfFile(_In_z_ const WCHAR
         LastError = ERROR_VERSION_PARSE_ERROR;
     }
 out:
-    HeapFree(ModuleHeap, 0, VersionInfo);
+    Free(VersionInfo);
     return RET_ERROR(Version, LastError);
 }
 
@@ -1114,17 +1089,13 @@ WintunGetVersion(void)
     ULONG BufferSize = 128 * 1024;
     for (;;)
     {
-        Modules = HeapAlloc(ModuleHeap, 0, BufferSize);
+        Modules = Alloc(BufferSize);
         if (!Modules)
-        {
-            LOG(WINTUN_LOG_ERR, L"Out of memory");
-            SetLastError(ERROR_OUTOFMEMORY);
             return 0;
-        }
         NTSTATUS Status = NtQuerySystemInformation(SystemModuleInformation, Modules, BufferSize, &BufferSize);
         if (NT_SUCCESS(Status))
             break;
-        HeapFree(ModuleHeap, 0, Modules);
+        Free(Modules);
         if (Status == STATUS_INFO_LENGTH_MISMATCH)
             continue;
         LOG(WINTUN_LOG_ERR, L"Failed to enumerate drivers");
@@ -1148,7 +1119,7 @@ WintunGetVersion(void)
     }
     LastError = ERROR_FILE_NOT_FOUND;
 cleanupModules:
-    HeapFree(ModuleHeap, 0, Modules);
+    Free(Modules);
     return RET_ERROR(Version, LastError);
 }
 
@@ -1212,7 +1183,7 @@ static _Return_type_success_(return != FALSE) BOOL SelectDriver(
                 if (DevInfoExistingAdapters == INVALID_HANDLE_VALUE)
                 {
                     LastError = LOG_LAST_ERROR(L"Failed to get present adapters");
-                    HeapFree(ModuleHeap, 0, DrvInfoDetailData);
+                    Free(DrvInfoDetailData);
                     goto cleanupExistingAdapters;
                 }
                 _Analysis_assume_(DevInfoExistingAdapters != NULL);
@@ -1237,7 +1208,7 @@ static _Return_type_success_(return != FALSE) BOOL SelectDriver(
         DriverDate = DrvInfoData.DriverDate;
         DriverVersion = DrvInfoData.DriverVersion;
     next:
-        HeapFree(ModuleHeap, 0, DrvInfoDetailData);
+        Free(DrvInfoDetailData);
     }
 
     if (DriverVersion)
@@ -1340,7 +1311,7 @@ cleanupExistingAdapters:
         while (ExistingAdapters)
         {
             SP_DEVINFO_DATA_LIST *Next = ExistingAdapters->Next;
-            HeapFree(ModuleHeap, 0, ExistingAdapters);
+            Free(ExistingAdapters);
             ExistingAdapters = Next;
         }
     }
@@ -1494,7 +1465,7 @@ static _Return_type_success_(return != NULL) WINTUN_ADAPTER *CreateAdapter(
         LastError = LOG(WINTUN_LOG_ERR, L"Failed to get NetCfgInstanceId");
         goto cleanupNetDevRegKey;
     }
-    HeapFree(ModuleHeap, 0, DummyStr);
+    Free(DummyStr);
     DWORD DummyDWORD;
     if (!RegistryQueryDWORDWait(NetDevRegKey, L"NetLuidIndex", WAIT_FOR_REGISTRY_TIMEOUT, &DummyDWORD))
     {
@@ -1534,7 +1505,7 @@ static _Return_type_success_(return != NULL) WINTUN_ADAPTER *CreateAdapter(
         LastError = LOG(WINTUN_LOG_ERR, L"Failed to get IpConfig");
         goto cleanupTcpipAdapterRegKey;
     }
-    HeapFree(ModuleHeap, 0, DummyStr);
+    Free(DummyStr);
 
     HKEY TcpipInterfaceRegKey;
     WCHAR TcpipInterfaceRegPath[MAX_REG_PATH];
@@ -1602,7 +1573,7 @@ cleanupTcpipAdapterRegKey:
 cleanupAdapter:
     if (LastError != ERROR_SUCCESS)
     {
-        HeapFree(ModuleHeap, 0, Adapter);
+        Free(Adapter);
         Adapter = NULL;
     }
 cleanupNetDevRegKey:
@@ -1856,7 +1827,7 @@ _Return_type_success_(return != FALSE) BOOL WINAPI
                 LastError = LastError != ERROR_SUCCESS ? LastError : GetLastError();
             }
         }
-        HeapFree(ModuleHeap, 0, DriverDetail);
+        Free(DriverDetail);
     }
     SetupDiDestroyDriverInfoList(DeviceInfoSet, NULL, SPDIT_CLASSDRIVER);
 cleanupDeviceInfoSet:
@@ -1905,7 +1876,7 @@ _Return_type_success_(return != FALSE) BOOL WINAPI
             break;
         }
         Continue = Func(Adapter, Param);
-        HeapFree(ModuleHeap, 0, Adapter);
+        Free(Adapter);
     }
     SetupDiDestroyDeviceInfoList(DevInfo);
 cleanupMutex:
diff --git a/api/entry.h b/api/entry.h
index 4f925b1..ac37680 100644
--- a/api/entry.h
+++ b/api/entry.h
@@ -26,8 +26,6 @@
 #endif
 #pragma warning(disable : 4127) /* conditional expression is constant */
 
-#define RET_ERROR(Ret, Error) ((Error) == ERROR_SUCCESS ? (Ret) : (SetLastError(Error), 0))
-
 extern HINSTANCE ResourceModule;
 extern HANDLE ModuleHeap;
 extern SECURITY_ATTRIBUTES SecurityAttributes;
diff --git a/api/logger.c b/api/logger.c
index fc43706..d86b7a8 100644
--- a/api/logger.c
+++ b/api/logger.c
@@ -5,6 +5,7 @@
 
 #include "logger.h"
 #include <Windows.h>
+#include <wchar.h>
 
 static BOOL CALLBACK
 NopLogger(_In_ WINTUN_LOGGER_LEVEL Level, _In_z_ const WCHAR *LogLine)
@@ -32,8 +33,9 @@ LoggerLog(_In_ WINTUN_LOGGER_LEVEL Level, _In_z_ const WCHAR *Function, _In_z_ c
     {
         WCHAR Combined[0x400];
         if (_snwprintf_s(Combined, _countof(Combined), _TRUNCATE, L"%s: %s", Function, LogLine) == -1)
-            return LastError;
-        Logger(Level, Combined);
+            Logger(Level, LogLine);
+        else
+            Logger(Level, Combined);
     }
     else
         Logger(Level, LogLine);
diff --git a/api/logger.h b/api/logger.h
index b0458c5..189ff72 100644
--- a/api/logger.h
+++ b/api/logger.h
@@ -6,6 +6,7 @@
 #pragma once
 
 #include "wintun.h"
+#include "entry.h"
 #include <Windows.h>
 
 extern WINTUN_LOGGER_CALLBACK_FUNC Logger;
@@ -36,3 +37,28 @@ LoggerLastError(_In_z_ const WCHAR *Function, _In_z_ const WCHAR *Prefix)
 #define LOG(lvl, msg) (LoggerLog((lvl), _L(__FUNCTION__), msg))
 #define LOG_ERROR(msg, err) (LoggerError(_L(__FUNCTION__), msg, (err)))
 #define LOG_LAST_ERROR(msg) (LoggerLastError(_L(__FUNCTION__), msg))
+
+#define RET_ERROR(Ret, Error) ((Error) == ERROR_SUCCESS ? (Ret) : (SetLastError(Error), 0))
+
+static inline _Return_type_success_(return != NULL) _Ret_maybenull_
+    _Post_writable_byte_size_(Size) void *LoggerAlloc(_In_z_ const WCHAR *Function, _In_ DWORD Flags, _In_ SIZE_T Size)
+{
+    void *Data = HeapAlloc(ModuleHeap, Flags, Size);
+    if (!Data)
+    {
+        LoggerLog(WINTUN_LOG_ERR, Function, L"Out of memory");
+        SetLastError(ERROR_OUTOFMEMORY);
+    }
+    return Data;
+}
+#define Alloc(Size) LoggerAlloc(_L(__FUNCTION__), 0, Size)
+#define Zalloc(Size) LoggerAlloc(_L(__FUNCTION__), HEAP_ZERO_MEMORY, Size)
+static inline void
+Free(void *Ptr)
+{
+    if (!Ptr)
+        return;
+    DWORD LastError = GetLastError();
+    HeapFree(ModuleHeap, 0, Ptr);
+    SetLastError(LastError);
+}
\ No newline at end of file
diff --git a/api/namespace.c b/api/namespace.c
index 44105ca..509173c 100644
--- a/api/namespace.c
+++ b/api/namespace.c
@@ -22,21 +22,16 @@ static _Return_type_success_(
     int Len = NormalizeString(NormForm, Source, -1, NULL, 0);
     for (;;)
     {
-        WCHAR *Str = HeapAlloc(ModuleHeap, 0, sizeof(WCHAR) * Len);
+        WCHAR *Str = Alloc(sizeof(WCHAR) * Len);
         if (!Str)
-        {
-            LOG(WINTUN_LOG_ERR, L"Out of memory");
-            SetLastError(ERROR_OUTOFMEMORY);
             return NULL;
-        }
         Len = NormalizeString(NormForm, Source, -1, Str, Len);
         if (Len > 0)
             return Str;
-        DWORD LastError = GetLastError();
-        HeapFree(ModuleHeap, 0, Str);
-        if (LastError != ERROR_INSUFFICIENT_BUFFER)
+        Free(Str);
+        if (GetLastError() != ERROR_INSUFFICIENT_BUFFER)
         {
-            SetLastError(LOG_ERROR(L"Failed", LastError));
+            LOG_LAST_ERROR(L"Failed");
             return NULL;
         }
         Len = -Len;
@@ -169,7 +164,7 @@ _Return_type_success_(return != NULL) HANDLE NamespaceTakePoolMutex(_In_z_ const
     {
     case WAIT_OBJECT_0:
     case WAIT_ABANDONED:
-        HeapFree(ModuleHeap, 0, PoolNorm);
+        Free(PoolNorm);
         BCryptDestroyHash(Sha256);
         return Mutex;
     }
@@ -177,7 +172,7 @@ _Return_type_success_(return != NULL) HANDLE NamespaceTakePoolMutex(_In_z_ const
     LastError = ERROR_GEN_FAILURE;
     CloseHandle(Mutex);
 cleanupPoolNorm:
-    HeapFree(ModuleHeap, 0, PoolNorm);
+    Free(PoolNorm);
 cleanupSha256:
     BCryptDestroyHash(Sha256);
     SetLastError(LastError);
diff --git a/api/registry.c b/api/registry.c
index 523016e..5ec1848 100644
--- a/api/registry.c
+++ b/api/registry.c
@@ -92,16 +92,12 @@ _Return_type_success_(return != FALSE) BOOL RegistryGetString(_Inout_ WCHAR **Bu
     if (wcsnlen(*Buf, Len) >= Len)
     {
         /* String is missing zero-terminator. */
-        WCHAR *BufZ = HeapAlloc(ModuleHeap, 0, ((size_t)Len + 1) * sizeof(WCHAR));
+        WCHAR *BufZ = Alloc(((size_t)Len + 1) * sizeof(WCHAR));
         if (!BufZ)
-        {
-            LOG(WINTUN_LOG_ERR, L"Out of memory");
-            SetLastError(ERROR_OUTOFMEMORY);
             return FALSE;
-        }
         wmemcpy(BufZ, *Buf, Len);
         BufZ[Len] = 0;
-        HeapFree(ModuleHeap, 0, *Buf);
+        Free(*Buf);
         *Buf = BufZ;
     }
 
@@ -116,28 +112,23 @@ _Return_type_success_(return != FALSE) BOOL RegistryGetString(_Inout_ WCHAR **Bu
     Len = Len * 2 + 64;
     for (;;)
     {
-        WCHAR *Expanded = HeapAlloc(ModuleHeap, 0, Len * sizeof(WCHAR));
+        WCHAR *Expanded = Alloc(Len * sizeof(WCHAR));
         if (!Expanded)
-        {
-            LOG(WINTUN_LOG_ERR, L"Out of memory");
-            SetLastError(ERROR_OUTOFMEMORY);
             return FALSE;
-        }
         DWORD Result = ExpandEnvironmentStringsW(*Buf, Expanded, Len);
         if (!Result)
         {
-            DWORD LastError = LOG_LAST_ERROR(L"Failed to expand environment variables");
-            HeapFree(ModuleHeap, 0, Expanded);
-            SetLastError(LastError);
+            LOG_LAST_ERROR(L"Failed to expand environment variables");
+            Free(Expanded);
             return FALSE;
         }
         if (Result > Len)
         {
-            HeapFree(ModuleHeap, 0, Expanded);
+            Free(Expanded);
             Len = Result;
             continue;
         }
-        HeapFree(ModuleHeap, 0, *Buf);
+        Free(*Buf);
         *Buf = Expanded;
         return TRUE;
     }
@@ -153,33 +144,25 @@ _Return_type_success_(return != FALSE) BOOL
             if (i > Len)
             {
                 /* Missing string and list terminators. */
-                WCHAR *BufZ = HeapAlloc(ModuleHeap, 0, ((size_t)Len + 2) * sizeof(WCHAR));
+                WCHAR *BufZ = Alloc(((size_t)Len + 2) * sizeof(WCHAR));
                 if (!BufZ)
-                {
-                    LOG(WINTUN_LOG_ERR, L"Out of memory");
-                    SetLastError(ERROR_OUTOFMEMORY);
                     return FALSE;
-                }
                 wmemcpy(BufZ, *Buf, Len);
                 BufZ[Len] = 0;
                 BufZ[Len + 1] = 0;
-                HeapFree(ModuleHeap, 0, *Buf);
+                Free(*Buf);
                 *Buf = BufZ;
                 return TRUE;
             }
             if (i == Len)
             {
                 /* Missing list terminator. */
-                WCHAR *BufZ = HeapAlloc(ModuleHeap, 0, ((size_t)Len + 1) * sizeof(WCHAR));
+                WCHAR *BufZ = Alloc(((size_t)Len + 1) * sizeof(WCHAR));
                 if (!BufZ)
-                {
-                    LOG(WINTUN_LOG_ERR, L"Out of memory");
-                    SetLastError(ERROR_OUTOFMEMORY);
                     return FALSE;
-                }
                 wmemcpy(BufZ, *Buf, Len);
                 BufZ[Len] = 0;
-                HeapFree(ModuleHeap, 0, *Buf);
+                Free(*Buf);
                 *Buf = BufZ;
                 return TRUE;
             }
@@ -192,16 +175,12 @@ _Return_type_success_(return != FALSE) BOOL
     if (!RegistryGetString(Buf, Len, ValueType))
         return FALSE;
     Len = (DWORD)wcslen(*Buf) + 1;
-    WCHAR *BufZ = HeapAlloc(ModuleHeap, 0, ((size_t)Len + 1) * sizeof(WCHAR));
+    WCHAR *BufZ = Alloc(((size_t)Len + 1) * sizeof(WCHAR));
     if (!BufZ)
-    {
-        LOG(WINTUN_LOG_ERR, L"Out of memory");
-        SetLastError(ERROR_OUTOFMEMORY);
         return FALSE;
-    }
     wmemcpy(BufZ, *Buf, Len);
     BufZ[Len] = 0;
-    HeapFree(ModuleHeap, 0, *Buf);
+    Free(*Buf);
     *Buf = BufZ;
     return TRUE;
 }
@@ -215,17 +194,13 @@ static _Return_type_success_(return != NULL) void *RegistryQuery(
 {
     for (;;)
     {
-        BYTE *p = HeapAlloc(ModuleHeap, 0, *BufLen);
+        BYTE *p = Alloc(*BufLen);
         if (!p)
-        {
-            LOG(WINTUN_LOG_ERR, L"Out of memory");
-            SetLastError(ERROR_OUTOFMEMORY);
             return NULL;
-        }
         LSTATUS LastError = RegQueryValueExW(Key, Name, NULL, ValueType, p, BufLen);
         if (LastError == ERROR_SUCCESS)
             return p;
-        HeapFree(ModuleHeap, 0, p);
+        Free(p);
         if (LastError != ERROR_MORE_DATA)
         {
             if (Log)
@@ -256,7 +231,7 @@ _Return_type_success_(
         LOG(WINTUN_LOG_ERR, L"Value is not a string");
         LastError = ERROR_INVALID_DATATYPE;
     }
-    HeapFree(ModuleHeap, 0, Value);
+    Free(Value);
     SetLastError(LastError);
     return NULL;
 }
diff --git a/api/rundll32_i.c b/api/rundll32_i.c
index 5f9de48..1b0a8d4 100644
--- a/api/rundll32_i.c
+++ b/api/rundll32_i.c
@@ -131,11 +131,10 @@ static _Return_type_success_(return != FALSE) BOOL ExecuteRunDll32(
         goto cleanupDelete;
     }
     size_t CommandLineLen = 10 + MAX_PATH + 2 + wcslen(Arguments) + 1;
-    WCHAR *CommandLine = HeapAlloc(ModuleHeap, 0, CommandLineLen * sizeof(WCHAR));
+    WCHAR *CommandLine = Alloc(CommandLineLen * sizeof(WCHAR));
     if (!CommandLine)
     {
-        LOG(WINTUN_LOG_ERR, L"Out of memory");
-        LastError = ERROR_OUTOFMEMORY;
+        LastError = GetLastError();
         goto cleanupDelete;
     }
     if (_snwprintf_s(CommandLine, CommandLineLen, _TRUNCATE, L"rundll32 \"%.*s\",%s", MAX_PATH, DllPath, Arguments) ==
@@ -218,7 +217,7 @@ cleanupPipes:
     CloseHandle(StreamWStderr);
     CloseHandle(StreamRStdout);
     CloseHandle(StreamWStdout);
-    HeapFree(ModuleHeap, 0, CommandLine);
+    Free(CommandLine);
 cleanupDelete:
     DeleteFileW(DllPath);
 cleanupDirectory:
diff --git a/api/session.c b/api/session.c
index c4e1c04..867ac04 100644
--- a/api/session.c
+++ b/api/session.c
@@ -73,11 +73,10 @@ _Return_type_success_(return != NULL) TUN_SESSION *WINAPI
     WintunStartSession(_In_ const WINTUN_ADAPTER *Adapter, _In_ DWORD Capacity)
 {
     DWORD LastError;
-    TUN_SESSION *Session = HeapAlloc(ModuleHeap, HEAP_ZERO_MEMORY, sizeof(TUN_SESSION));
+    TUN_SESSION *Session = Zalloc(sizeof(TUN_SESSION));
     if (!Session)
     {
-        LOG(WINTUN_LOG_ERR, L"Out of memory");
-        LastError = ERROR_OUTOFMEMORY;
+        LastError = GetLastError();
         goto out;
     }
     const ULONG RingSize = TUN_RING_SIZE(Capacity);
@@ -146,7 +145,7 @@ cleanupToken:
 cleanupAllocatedRegion:
     VirtualFree(AllocatedRegion, 0, MEM_RELEASE);
 cleanupRings:
-    HeapFree(ModuleHeap, 0, Session);
+    Free(Session);
 out:
     SetLastError(LastError);
     return NULL;
@@ -162,7 +161,7 @@ WintunEndSession(_In_ TUN_SESSION *Session)
     CloseHandle(Session->Descriptor.Send.TailMoved);
     CloseHandle(Session->Descriptor.Receive.TailMoved);
     VirtualFree(Session->Descriptor.Send.Ring, 0, MEM_RELEASE);
-    HeapFree(ModuleHeap, 0, Session);
+    Free(Session);
 }
 
 HANDLE WINAPI