Merge pull request #1402 from microsoft/main

FI: Main -> Develop

1 files changed, 149 insertions, 0 deletions

diff --git a/network/wlan/wificx/drivercode/memorymanagement.cpp b/network/wlan/wificx/drivercode/memorymanagement.cpp
new file mode 100644
index 00000000..088bad88
--- /dev/null
+++ b/network/wlan/wificx/drivercode/memorymanagement.cpp
@@ -0,0 +1,149 @@
+// Copyright (C) Microsoft Corporation. All rights reserved.
+#include "precomp.h"
+
+#ifndef _KERNEL_MODE
+#include <intrin.h> // for _ReturnAddress for user mode
+#endif // UM
+
+typedef struct _PLACEMENT_NEW_ALLOCATION_CONTEXT
+{
+    size_t cbMaxSize;
+    _Field_size_bytes_(cbMaxSize) void* pbBuffer;
+} PLACEMENT_NEW_ALLOCATION_CONTEXT, * PPLACEMENT_NEW_ALLOCATION_CONTEXT;
+typedef const PLACEMENT_NEW_ALLOCATION_CONTEXT* PCPLACEMENT_NEW_ALLOCATION_CONTEXT;
+
+// for FreeWdfMemoryBuffer to correct get
+// the handle to free the memory
+struct WIFI_IHV_MEMORY_HEADER
+{
+    size_t HeaderSize;
+    WDFMEMORY WdfMemoryHandle;
+};
+
+// for tracking memory leaks
+struct WIFI_IHV_MEMORY_CONTEXT
+{
+    size_t ContextSize;
+    void* pvCaller;
+};
+
+WDF_DECLARE_CONTEXT_TYPE_WITH_NAME(WIFI_IHV_MEMORY_CONTEXT, GetWificxIhvMemoryContextFromHandle);
+
+
+void* AllocateWdfMemoryBuffer(size_t Size, _In_ void* CallerForMemoryLeakTracking)
+{
+    size_t totalSize = 0;
+    if (!NT_SUCCESS(Wifi::SizeTAddSafe(Size, sizeof(WIFI_IHV_MEMORY_HEADER), &totalSize)))
+    {
+        NT_ASSERT(FALSE);
+        WFCError("Failed to calculate total size");
+        return nullptr;
+    }
+
+    // allocate the context memory to store the WIFI_IHV_MEMORY_CONTEXT
+    WDF_OBJECT_ATTRIBUTES attributes;
+    WDF_OBJECT_ATTRIBUTES_INIT_CONTEXT_TYPE(&attributes, WIFI_IHV_MEMORY_CONTEXT);
+    WDFMEMORY wdfMemoryBufferHandle = WDF_NO_HANDLE;
+    void* pWdfMemoryBuffer = nullptr;
+
+    if (!NT_SUCCESS(WdfMemoryCreate(&attributes, NonPagedPoolNx, WIFI_DRIVER_DEFAULT_POOL_TAG, totalSize, &wdfMemoryBufferHandle, &pWdfMemoryBuffer)))
+    {
+        WFCError("Failed to allocate memory buffer");
+        return nullptr;
+    }
+
+    RtlZeroMemory(pWdfMemoryBuffer, totalSize);
+
+    // store the wdf memory handle in the header for FreeWdfMemoryBuffer to find the handle
+    WIFI_IHV_MEMORY_HEADER* pMemoryHeader = static_cast<WIFI_IHV_MEMORY_HEADER*>(pWdfMemoryBuffer);
+    pMemoryHeader->HeaderSize = sizeof(WIFI_IHV_MEMORY_HEADER);
+    pMemoryHeader->WdfMemoryHandle = wdfMemoryBufferHandle;
+
+    // store the caller and caller's caller for tracking memory leaks
+    auto memoryContext = GetWificxIhvMemoryContextFromHandle(wdfMemoryBufferHandle);
+    RtlZeroMemory(memoryContext, sizeof(WIFI_IHV_MEMORY_CONTEXT));
+    memoryContext->ContextSize = sizeof(WIFI_IHV_MEMORY_CONTEXT);
+    memoryContext->pvCaller = CallerForMemoryLeakTracking;
+
+    // hide the wdf memory header to the caller.
+    const ULONG_PTR wdfMemoryWithHeaderBuffer = reinterpret_cast<ULONG_PTR>(pWdfMemoryBuffer);
+    ULONG_PTR wdfMemoryPayloadOnlyBuffer{ 0 };
+    if (!NT_SUCCESS(Wifi::ULongPtrAddSafe(wdfMemoryWithHeaderBuffer, sizeof(WIFI_IHV_MEMORY_HEADER), &wdfMemoryPayloadOnlyBuffer)))
+    {
+        NT_ASSERT(FALSE);
+        WFCError("Failed to calculate payload buffer address");
+        return nullptr;
+    }
+
+    return reinterpret_cast<void*>(wdfMemoryPayloadOnlyBuffer);
+}
+
+void FreeWdfMemoryBuffer(_In_opt_ void* pBuffer)
+{
+    if (pBuffer == nullptr) // existing wdi code not checking for nullptr before delete, so we have to leave it as is
+    {
+        return;
+    }
+    const ULONG_PTR wdfMemoryPayloadOnlyBuffer = reinterpret_cast<ULONG_PTR>(pBuffer);
+    ULONG_PTR wdfMemoryWithHeaderBuffer{ 0 };
+    if (!NT_SUCCESS(Wifi::ULongPtrSubSafe(wdfMemoryPayloadOnlyBuffer, sizeof(WIFI_IHV_MEMORY_HEADER), &wdfMemoryWithHeaderBuffer)))
+    {
+        NT_ASSERT(FALSE);
+        WFCError("Failed to calculate header buffer address");
+        return;
+    }
+
+    const auto wificxMemoryHeader = static_cast<WIFI_IHV_MEMORY_HEADER*>(reinterpret_cast<void*>(wdfMemoryWithHeaderBuffer));
+
+    // This is memory corruption detection logic,
+    // keep it, don't remove it.
+    NT_ASSERT(wificxMemoryHeader->HeaderSize == sizeof(WIFI_IHV_MEMORY_HEADER));
+
+    WdfObjectDelete(wificxMemoryHeader->WdfMemoryHandle);
+}
+
+_Ret_writes_bytes_maybenull_(_Size) void* PlacementNewHelper(size_t _Size, PCPLACEMENT_NEW_ALLOCATION_CONTEXT AllocationContext)
+{
+    if (_Size <= AllocationContext->cbMaxSize)
+    {
+        RtlZeroMemory(AllocationContext->pbBuffer, _Size);
+        return AllocationContext->pbBuffer;
+    }
+    WFCError(
+        "Placement operator new called with insufficient buffer space (desired: %Iu, availible: %Iu)", _Size, AllocationContext->cbMaxSize);
+    return nullptr;
+}
+
+void* __cdecl operator new(size_t Size) noexcept
+{
+    return AllocateWdfMemoryBuffer(Size, _ReturnAddress());
+}
+
+__forceinline void* __cdecl operator new(size_t _Size, ULONG_PTR AllocationContext) noexcept // for WIFICX TLV
+{
+    if (AllocationContext != 0)
+    {
+        return PlacementNewHelper(_Size, (PPLACEMENT_NEW_ALLOCATION_CONTEXT)AllocationContext);
+    }
+    return AllocateWdfMemoryBuffer(_Size, _ReturnAddress());
+}
+
+void __cdecl operator delete(void* pData) noexcept
+{
+    FreeWdfMemoryBuffer(pData);
+}
+
+void __cdecl operator delete[](void* pData) noexcept
+{
+    FreeWdfMemoryBuffer(pData);
+}
+
+void __cdecl operator delete(void* pData, ULONG_PTR) noexcept // For WIFICX TLV
+{
+    FreeWdfMemoryBuffer(pData);
+}
+
+void __cdecl operator delete[](void* pData, ULONG_PTR) noexcept // For WIFICX TLV
+{
+    FreeWdfMemoryBuffer(pData);
+}
\ No newline at end of file