Move around HEAP API wrappers to a more correct spot, and add ANSI string support for heap strings (#477)

* Move around heap wrappers to a correct spot, and add ANSI string support for process heap strings

* Address code review feedback and fix formatting

* Fix COM unit tests

---------

Co-authored-by: Duncan Horn <[email protected]>

Author: GabeDeBacker

include/wil/resource.h

@@ -3717,6 +3717,14 @@ namespace details
     {
         ::HeapFree(::GetProcessHeap(), 0, p);
     }
+
+    struct heap_allocator
+    {
+        static _Ret_opt_bytecap_(size) void* allocate(size_t size) WI_NOEXCEPT
+        {
+            return ::HeapAlloc(::GetProcessHeap(), HEAP_ZERO_MEMORY, size);
+        }
+    };
 } // namespace details
 /// @endcond
 
@@ -3747,25 +3755,6 @@ struct mapview_deleter
     }
 };
 
-template <typename T = void>
-using unique_process_heap_ptr = wistd::unique_ptr<details::ensure_trivially_destructible_t<T>, process_heap_deleter>;
-
-typedef unique_any<PWSTR, decltype(&details::FreeProcessHeap), details::FreeProcessHeap> unique_process_heap_string;
-
-/// @cond
-namespace details
-{
-    template <>
-    struct string_allocator<unique_process_heap_string>
-    {
-        static _Ret_opt_bytecap_(size) void* allocate(size_t size) WI_NOEXCEPT
-        {
-            return ::HeapAlloc(::GetProcessHeap(), HEAP_ZERO_MEMORY, size);
-        }
-    };
-} // namespace details
-/// @endcond
-
 /** Manages a typed pointer allocated with VirtualAlloc
 A specialization of wistd::unique_ptr<> that frees via VirtualFree(p, 0, MEM_RELEASE).
 */
@@ -4061,6 +4050,93 @@ typedef weak_any<shared_hfind_change> weak_hfind_change;
 
 #endif // __WIL_WINBASE_STL
 
+#if (defined(_HEAPAPI_H_) && !defined(__WIL__WIL_HEAP_API) && WINAPI_FAMILY_PARTITION(WINAPI_PARTITION_APP | WINAPI_PARTITION_SYSTEM | WINAPI_PARTITION_GAMES) && !defined(WIL_KERNEL_MODE)) || \
+    defined(WIL_DOXYGEN)
+/// @cond
+#define __WIL__WIL_HEAP_API
+/// @endcond
+
+template <typename T = void>
+using unique_process_heap_ptr = wistd::unique_ptr<details::ensure_trivially_destructible_t<T>, process_heap_deleter>;
+typedef unique_any<void*, decltype(&details::FreeProcessHeap), details::FreeProcessHeap> unique_process_heap;
+typedef unique_any<PWSTR, decltype(&details::FreeProcessHeap), details::FreeProcessHeap> unique_process_heap_string;
+
+#ifndef WIL_NO_ANSI_STRINGS
+typedef unique_any<PSTR, decltype(&wil::details::FreeProcessHeap), wil::details::FreeProcessHeap> unique_process_heap_ansistring;
+#endif // WIL_NO_ANSI_STRINGS
+
+/// @cond
+namespace details
+{
+    template <>
+    struct string_allocator<wil::unique_process_heap_string> : heap_allocator
+    {
+    };
+
+#ifndef WIL_NO_ANSI_STRINGS
+    template <>
+    struct string_allocator<unique_process_heap_ansistring> : heap_allocator
+    {
+    };
+#endif
+} // namespace details
+/// @endcond
+
+inline auto make_process_heap_string_nothrow(
+    _When_((source != nullptr) && length != static_cast<size_t>(-1), _In_reads_(length))
+        _When_((source != nullptr) && length == static_cast<size_t>(-1), _In_z_) PCWSTR source,
+    size_t length = static_cast<size_t>(-1)) WI_NOEXCEPT
+{
+    return make_unique_string_nothrow<unique_process_heap_string>(source, length);
+}
+
+inline auto make_process_heap_string_failfast(
+    _When_((source != nullptr) && length != static_cast<size_t>(-1), _In_reads_(length))
+        _When_((source != nullptr) && length == static_cast<size_t>(-1), _In_z_) PCWSTR source,
+    size_t length = static_cast<size_t>(-1)) WI_NOEXCEPT
+{
+    return make_unique_string_failfast<unique_process_heap_string>(source, length);
+}
+
+#ifndef WIL_NO_ANSI_STRINGS
+inline auto make_process_heap_ansistring_nothrow(
+    _When_((source != nullptr) && length != static_cast<size_t>(-1), _In_reads_(length))
+        _When_((source != nullptr) && length == static_cast<size_t>(-1), _In_z_) PCSTR source,
+    size_t length = static_cast<size_t>(-1)) WI_NOEXCEPT
+{
+    return make_unique_ansistring_nothrow<unique_process_heap_ansistring>(source, length);
+}
+
+inline auto make_process_heap_ansistring_failfast(
+    _When_((source != nullptr) && length != static_cast<size_t>(-1), _In_reads_(length))
+        _When_((source != nullptr) && length == static_cast<size_t>(-1), _In_z_) PCSTR source,
+    size_t length = static_cast<size_t>(-1)) WI_NOEXCEPT
+{
+    return make_unique_ansistring_failfast<unique_process_heap_ansistring>(source, length);
+}
+#endif // WIL_NO_ANSI_STRINGS
+
+#ifdef WIL_ENABLE_EXCEPTIONS
+inline auto make_process_heap_string(
+    _When_((source != nullptr) && length != static_cast<size_t>(-1), _In_reads_(length))
+        _When_((source != nullptr) && length == static_cast<size_t>(-1), _In_z_) PCWSTR source,
+    size_t length = static_cast<size_t>(-1))
+{
+    return make_unique_string<unique_process_heap_string>(source, length);
+}
+
+#ifndef WIL_NO_ANSI_STRINGS
+inline auto make_process_heap_ansistring(
+    _When_((source != nullptr) && length != static_cast<size_t>(-1), _In_reads_(length))
+        _When_((source != nullptr) && length == static_cast<size_t>(-1), _In_z_) PCSTR source,
+    size_t length = static_cast<size_t>(-1))
+{
+    return make_unique_ansistring<unique_process_heap_ansistring>(source, length);
+}
+#endif // WIL_NO_ANSI_STRINGS
+#endif // WIL_ENABLE_EXCEPTIONS
+#endif // _HEAPAPI_H_
+
 #if (defined(__WIL_WINBASE_) && defined(__NOTHROW_T_DEFINED) && !defined(__WIL_WINBASE_NOTHROW_T_DEFINED_STL) && defined(WIL_RESOURCE_STL) && WINAPI_FAMILY_PARTITION(WINAPI_PARTITION_DESKTOP)) || \
     defined(WIL_DOXYGEN)
 /// @cond
@@ -4508,32 +4584,6 @@ namespace details
 } // namespace details
 /// @endcond
 
-inline auto make_process_heap_string_nothrow(
-    _When_((source != nullptr) && length != static_cast<size_t>(-1), _In_reads_(length))
-        _When_((source != nullptr) && length == static_cast<size_t>(-1), _In_z_) PCWSTR source,
-    size_t length = static_cast<size_t>(-1)) WI_NOEXCEPT
-{
-    return make_unique_string_nothrow<unique_process_heap_string>(source, length);
-}
-
-inline auto make_process_heap_string_failfast(
-    _When_((source != nullptr) && length != static_cast<size_t>(-1), _In_reads_(length))
-        _When_((source != nullptr) && length == static_cast<size_t>(-1), _In_z_) PCWSTR source,
-    size_t length = static_cast<size_t>(-1)) WI_NOEXCEPT
-{
-    return make_unique_string_failfast<unique_process_heap_string>(source, length);
-}
-
-#ifdef WIL_ENABLE_EXCEPTIONS
-inline auto make_process_heap_string(
-    _When_((source != nullptr) && length != static_cast<size_t>(-1), _In_reads_(length))
-        _When_((source != nullptr) && length == static_cast<size_t>(-1), _In_z_) PCWSTR source,
-    size_t length = static_cast<size_t>(-1))
-{
-    return make_unique_string<unique_process_heap_string>(source, length);
-}
-#endif // WIL_ENABLE_EXCEPTIONS
-
 typedef unique_any_handle_null<decltype(&::HeapDestroy), ::HeapDestroy> unique_hheap;
 typedef unique_any<DWORD, decltype(&::TlsFree), ::TlsFree, details::pointer_access_all, DWORD, DWORD, TLS_OUT_OF_INDEXES, DWORD> unique_tls;
 typedef unique_any<PSECURITY_DESCRIPTOR, decltype(&::LocalFree), ::LocalFree> unique_hlocal_security_descriptor;

tests/ComTests.cpp

@@ -3067,9 +3067,6 @@ TEST_CASE("COMEnumerator", "[com][enumerator]")
 #include <winrt/windows.foundation.h>
 #include <windows.foundation.h>
 
-HANDLE g_hangHandle{nullptr};
-HANDLE g_doneHangingHandle{nullptr};
-
 TEST_CASE("com_timeout", "[com][com_timeout]")
 {
     auto init = wil::CoInitializeEx_failfast();
@@ -3083,24 +3080,32 @@ TEST_CASE("com_timeout", "[com][com_timeout]")
     // apartments (MTA -> STA in this case)
     struct COMTimeoutTestObject : winrt::implements<COMTimeoutTestObject, winrt::Windows::Foundation::IStringable, winrt::non_agile>
     {
+        wil::shared_event _hangHandle;
+        wil::shared_event _doneHangingHandle;
+        std::shared_ptr<bool> _shouldHang;
+        COMTimeoutTestObject(wil::shared_event hangHandle, wil::shared_event doneHangingHandle, std::shared_ptr<bool> shouldHang) :
+            _hangHandle(hangHandle), _doneHangingHandle(doneHangingHandle), _shouldHang(shouldHang)
+        {
+        }
+
         winrt::hstring ToString()
         {
-            // If the global handle exists, block on it.  If it doesn't exist then this is a non-hang case
-            // so skip waiting.
-            if (g_hangHandle)
+            // If the test wants to block, then use the hang handles.
+            if (*(_shouldHang.get()))
             {
                 // Pump messages so this STA thread is healthy while we wait.  If this wait fails that means
                 // the cancel did not work.
-                HANDLE handles[1] = {g_hangHandle};
+                HANDLE handles[1] = {_hangHandle.get()};
                 DWORD index;
                 REQUIRE_SUCCEEDED(CoWaitForMultipleObjects(
                     CWMO_DISPATCH_CALLS | CWMO_DISPATCH_WINDOW_MESSAGES, 10000, ARRAYSIZE(handles), handles, &index));
 
-                if (g_doneHangingHandle)
+                if (_doneHangingHandle)
                 {
-                    SetEvent(g_doneHangingHandle);
+                    _doneHangingHandle.SetEvent();
                 }
             }
+
             return L"COMTimeoutTestObject";
         }
     };
@@ -3109,27 +3114,47 @@ TEST_CASE("com_timeout", "[com][com_timeout]")
     // exit.
     wil::shared_event comServerEvent;
     comServerEvent.create();
+
     wil::shared_event agileReferencePopulated;
     agileReferencePopulated.create();
+
+    // These handles are used to coordinate with the COM server thread.  The first one causes it to block.  The
+    // done hanging event lets us know that it is done blocking and we can proceed with a second call that should
+    // avoid reentering.
+    wil::shared_event hangingHandle;
+    hangingHandle.create();
+    wil::shared_event doneHangingHandle;
+    doneHangingHandle.create();
+
+    auto shouldHang = std::make_shared<bool>(false);
+
     wil::com_agile_ref agileStringable;
 
-    auto comServerThread = std::thread([comServerEvent, agileReferencePopulated, &agileStringable] {
-        // This thread must be STA to pull RPC in as mediator between threads.
-        auto init = wil::CoInitializeEx_failfast(COINIT_APARTMENTTHREADED);
+    auto comServerThread =
+        std::thread([comServerEvent, agileReferencePopulated, hangingHandle, doneHangingHandle, &agileStringable, shouldHang] {
+            // This thread must be STA to pull RPC in as mediator between threads.
+            auto init = wil::CoInitializeEx_failfast(COINIT_APARTMENTTHREADED);
 
-        const auto stringable = winrt::make<COMTimeoutTestObject>();
-        agileStringable = wil::com_agile_query(stringable.as<ABI::Windows::Foundation::IStringable>().get());
+            const auto stringable = winrt::make<COMTimeoutTestObject>(hangingHandle, doneHangingHandle, shouldHang);
+            agileStringable = wil::com_agile_query(stringable.as<ABI::Windows::Foundation::IStringable>().get());
 
-        agileReferencePopulated.SetEvent();
+            agileReferencePopulated.SetEvent();
 
-        // Pump messages so this STA thread is healthy.
-        HANDLE handles[1] = {comServerEvent.get()};
-        DWORD index;
-        REQUIRE_SUCCEEDED(CoWaitForMultipleObjects(
-            CWMO_DISPATCH_CALLS | CWMO_DISPATCH_WINDOW_MESSAGES, 10000, ARRAYSIZE(handles), handles, &index));
+            // Pump messages so this STA thread is healthy.
+            HANDLE handles[1] = {comServerEvent.get()};
+            DWORD index;
+            REQUIRE_SUCCEEDED(CoWaitForMultipleObjects(
+                CWMO_DISPATCH_CALLS | CWMO_DISPATCH_WINDOW_MESSAGES, INFINITE, ARRAYSIZE(handles), handles, &index));
+        });
+
+    auto makeSureComServerThreadExits = wil::scope_exit([comServerEvent, &comServerThread] {
+        // We are done testing.  Tell the STA thread to exit and then block until it is done.
+        comServerEvent.SetEvent();
+
+        comServerThread.join();
     });
 
-    agileReferencePopulated.wait(5000);
+    REQUIRE_SUCCEEDED(agileReferencePopulated.wait(5000));
 
     SECTION("Basic construction nothrow")
     {
@@ -3151,42 +3176,33 @@ TEST_CASE("com_timeout", "[com][com_timeout]")
     }
     SECTION("RPC timeout test")
     {
-        // These handles are used to coordinate with the COM server thread.  The first one causes it to block.  The
-        // done hanging event lets us know that it is done blocking and we can proceed with a second call that should
-        // avoid reentering.
-        wil::unique_event hangHandle;
-        hangHandle.create();
-        g_hangHandle = hangHandle.get();
-
-        wil::unique_event doneHangingHandle;
-        doneHangingHandle.create();
-        g_doneHangingHandle = doneHangingHandle.get();
-
         wil::com_timeout timeout{100};
 
+        *(shouldHang.get()) = true;
+
         // The timeout is now in place.  The blocking call should cancel in a timely manner and fail with RPC_E_CALL_CANCELED.
         wil::com_ptr<ABI::Windows::Foundation::IStringable> localServer =
             agileStringable.query<ABI::Windows::Foundation::IStringable>();
         wil::unique_hstring value;
-        REQUIRE(localServer->ToString(&value) == RPC_E_CALL_CANCELED);
+        auto localServerResult = localServer->ToString(&value);
+        REQUIRE(static_cast<bool>(localServerResult == RPC_E_CALL_CANCELED));
         REQUIRE(timeout.timed_out());
 
-        hangHandle.SetEvent();
+        hangingHandle.SetEvent();
         REQUIRE(doneHangingHandle.wait(5000));
 
-        hangHandle.ResetEvent();
+        hangingHandle.ResetEvent();
 
         // Make a second blocking call within the lifetime of the same com_timeout instance.  This second call should also
         // cancel and return.
-        const auto result = localServer->ToString(&value);
-        REQUIRE(result == RPC_E_CALL_CANCELED);
+        localServerResult = localServer->ToString(&value);
+        REQUIRE(static_cast<bool>(localServerResult == RPC_E_CALL_CANCELED));
         REQUIRE(timeout.timed_out());
 
-        hangHandle.SetEvent();
+        hangingHandle.SetEvent();
         REQUIRE(doneHangingHandle.wait(5000));
 
-        g_hangHandle = nullptr;
-        g_doneHangingHandle = nullptr;
+        *(shouldHang.get()) = false;
     }
     SECTION("Non-timeout unaffected test")
     {
@@ -3200,13 +3216,6 @@ TEST_CASE("com_timeout", "[com][com_timeout]")
         REQUIRE(!timeout.timed_out());
         REQUIRE(std::wstring_view{L"COMTimeoutTestObject"} == WindowsGetStringRawBuffer(value.get(), nullptr));
     }
-
-    // We are done testing.  Tell the STA thread to exit and then block until it is done.
-    comServerEvent.SetEvent();
-    comServerThread.join();
-
-    g_hangHandle = nullptr;
-    g_doneHangingHandle = nullptr;
 }
 #endif // defined(__cpp_impl_coroutine) || defined(__cpp_coroutines) || defined(_RESUMABLE_FUNCTIONS_SUPPORTED)
 #endif // (NTDDI_VERSION >= NTDDI_WINBLUE)