diff --git a/CMakeLists.txt b/CMakeLists.txt index f49447a8a..e62489861 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -379,6 +379,9 @@ if (WIN32) # VirtualAlloc2 is exposed by mincore.lib, not Kernel32.lib (as the # documentation says) target_link_libraries(snmalloc INTERFACE $<$>:mincore>) + # RtlDllShutdownInProgress is exported by ntdll. MSVC links this via a + # pragma in pal_windows.h, but MinGW ignores that pragma. + target_link_libraries(snmalloc INTERFACE ntdll) message(STATUS "snmalloc: Avoiding Windows 10 APIs is ${WIN8COMPAT}") endif() @@ -548,7 +551,8 @@ if(NOT SNMALLOC_HEADER_ONLY_LIBRARY) # These are mitigation-independent and can be compiled once, then linked # against both fast and check testlib variants. set(TESTLIB_ONLY_TESTS - bits first_operation memory memory_usage multi_atexit multi_threadatexit + bits dll_teardown first_operation fls_teardown memory memory_usage + multi_atexit multi_threadatexit redblack statistics teardown contention external_pointer large_alloc lotsofthreads post_teardown singlethread startup @@ -634,6 +638,10 @@ if(NOT SNMALLOC_HEADER_ONLY_LIBRARY) "LLVM_PROFILE_FILE=${CMAKE_BINARY_DIR}/profiles/${TESTNAME}/%p-%m.profraw") endif() if(WIN32) + if (${TEST} STREQUAL "dll_teardown") + add_dependencies(${TESTNAME} + snmalloc-dll-teardown-alloc snmalloc-dll-teardown-touch) + endif() # On Windows these tests use a lot of memory as it doesn't support # lazy commit. if (${TEST} MATCHES "two_alloc_types") @@ -803,6 +811,22 @@ if(NOT SNMALLOC_HEADER_ONLY_LIBRARY) endif() if (SNMALLOC_BUILD_TESTING) + if (WIN32) + # DLLs loaded by the func-dll_teardown tests. + add_library(snmalloc-dll-teardown-alloc SHARED + ${TESTDIR}/func/dll_teardown/dll/alloc.cc) + target_link_libraries(snmalloc-dll-teardown-alloc PRIVATE snmalloc) + target_compile_definitions(snmalloc-dll-teardown-alloc PRIVATE + "SNMALLOC_USE_${TEST_CLEANUP}") + add_library(snmalloc-dll-teardown-touch SHARED + ${TESTDIR}/func/dll_teardown/dll/touch.cc) + foreach(DLL snmalloc-dll-teardown-alloc snmalloc-dll-teardown-touch) + # The test loads the DLLs by name, so avoid MinGW's "lib" prefix. + set_target_properties(${DLL} PROPERTIES PREFIX "") + add_warning_flags(${DLL}) + endforeach() + endif() + set(FLAVOURS fast;check) foreach(FLAVOUR ${FLAVOURS}) diff --git a/snmalloc-rs/snmalloc-sys/build.rs b/snmalloc-rs/snmalloc-sys/build.rs index 9aa525c9d..224c9b6b1 100644 --- a/snmalloc-rs/snmalloc-sys/build.rs +++ b/snmalloc-rs/snmalloc-sys/build.rs @@ -583,6 +583,7 @@ fn configure_linking(config: &BuildConfig) { println!("cargo:rustc-link-lib=ws2_32"); println!("cargo:rustc-link-lib=userenv"); println!("cargo:rustc-link-lib=bcrypt"); + println!("cargo:rustc-link-lib=ntdll"); if config.debug { println!("cargo:rustc-link-lib=msvcrtd"); } else { @@ -592,6 +593,7 @@ fn configure_linking(config: &BuildConfig) { _ if config.is_windows() && config.is_gnu() => { println!("cargo:rustc-link-lib=kernel32"); println!("cargo:rustc-link-lib=bcrypt"); + println!("cargo:rustc-link-lib=ntdll"); println!("cargo:rustc-link-lib=winpthread"); if config.is_clang_msys() { diff --git a/src/snmalloc/pal/pal_windows.h b/src/snmalloc/pal/pal_windows.h index 5199c0604..aa20602fc 100644 --- a/src/snmalloc/pal/pal_windows.h +++ b/src/snmalloc/pal/pal_windows.h @@ -18,6 +18,7 @@ # include # pragma comment(lib, "bcrypt.lib") # include +# pragma comment(lib, "ntdll.lib") // VirtualAlloc2 is exposed in RS5 headers. # ifdef NTDDI_WIN10_RS5 # if (NTDDI_VERSION >= NTDDI_WIN10_RS5) && \ @@ -35,17 +36,39 @@ * allocations have been freed. * * One way to guarantee that the reservations get released - * at the absolute end of the program is to force them to + * at the end of the CRT's teardown is to force them to * be initialized first. Statics and globals get destroyed * in FILO order of when they were initialized. The pragma * init_seg makes sure the statics and globals in this * file are handled first, and thus will be the last to - * be destroyed when the program exits or the DLL is - * unloaded. + * be destroyed by the CRT when the program exits or the + * DLL is unloaded. + * + * Fiber-local storage callbacks (used, for example, by + * Rust's thread-local destructors) and DLL detach + * notifications are run by the loader during process + * exit, which for an executable is after the CRT's + * teardown. Hence, the reservations are only released + * when a DLL is unloaded while the process continues to + * run; see ~VirtualVector. */ # pragma warning(disable : 4075) # pragma init_seg(".CRT$XCB") +/** + * Linker-provided symbol at the base address of the current module. + */ +extern "C" IMAGE_DOS_HEADER __ImageBase; + +/** + * Returns TRUE if the loader has started shutting down the process, and FALSE + * otherwise, including when a DLL is unloaded by FreeLibrary while the process + * continues to run. This is exported by ntdll.dll, but not declared in the + * SDK headers, so must be declared by the caller. See + * https://learn.microsoft.com/en-us/windows/win32/devnotes/rtldllshutdowninprogress + */ +extern "C" __declspec(dllimport) BOOLEAN NTAPI RtlDllShutdownInProgress(); + namespace snmalloc { class PALWindows : public PalTimerDefaultImpl @@ -231,6 +254,18 @@ namespace snmalloc abort(); } + /** + * Returns true if the module containing this code is a DLL, and false if + * it is the executable. + */ + static bool current_module_is_dll() + { + const auto* base = reinterpret_cast(&__ImageBase); + const auto* nt_headers = + reinterpret_cast(base + __ImageBase.e_lfanew); + return (nt_headers->FileHeader.Characteristics & IMAGE_FILE_DLL) != 0; + } + /// Notify platform that we will not be using these pages static void notify_not_using(void* p, size_t size) noexcept { @@ -414,6 +449,18 @@ namespace snmalloc ~VirtualVector() { + // Memory is only returned to the OS if this module is being unloaded + // while the process continues to run. An executable's statics are + // only destroyed as part of process exit, and a DLL may be detached + // as part of process exit. In both cases, code run later by the + // loader, such as fiber-local storage callbacks or other DLLs' detach + // notifications, may still access allocations, and the OS reclaims + // the memory when the process terminates. + if ( + !snmalloc::PALWindows::current_module_is_dll() || + RtlDllShutdownInProgress()) + return; + if (data) { for (size_t i = size; i > 0; i--) @@ -588,7 +635,9 @@ namespace snmalloc /** * This will be destroyed last of all of the - * statics and globals due to init_seg + * statics and globals due to init_seg. See + * ~VirtualVector for when the reservations are + * released. */ static inline VirtualVector reservations; diff --git a/src/test/func/dll_teardown/dll/alloc.cc b/src/test/func/dll_teardown/dll/alloc.cc new file mode 100644 index 000000000..939d2f499 --- /dev/null +++ b/src/test/func/dll_teardown/dll/alloc.cc @@ -0,0 +1,15 @@ +/** + * DLL that uses snmalloc, for the dll_teardown test. + */ + +#include + +extern "C" __declspec(dllexport) void* dll_teardown_alloc(size_t size) +{ + return snmalloc::alloc(size); +} + +extern "C" __declspec(dllexport) void dll_teardown_dealloc(void* p) +{ + snmalloc::dealloc(p); +} diff --git a/src/test/func/dll_teardown/dll/touch.cc b/src/test/func/dll_teardown/dll/touch.cc new file mode 100644 index 000000000..fd3efea66 --- /dev/null +++ b/src/test/func/dll_teardown/dll/touch.cc @@ -0,0 +1,50 @@ +/** + * DLL that accesses an object during its process-exit detach notification, + * for the dll_teardown test. The object is allocated by a different DLL that + * uses snmalloc and was loaded after this one, so it is detached before this + * one. + */ + +#ifndef WIN32_LEAN_AND_MEAN +# define WIN32_LEAN_AND_MEAN +#endif +#include + +namespace +{ + volatile size_t* object = nullptr; +} + +extern "C" __declspec(dllexport) void dll_teardown_set_object(size_t* p) +{ + object = p; +} + +extern "C" BOOL WINAPI DllMain(HINSTANCE, DWORD reason, LPVOID reserved) +{ + // A non-null reserved argument indicates process exit. + if ( + (reason == DLL_PROCESS_DETACH) && (reserved != nullptr) && + (object != nullptr)) + { + // The loader may catch exceptions raised by DllMain, so check that the + // memory is still committed rather than relying on an access violation. + MEMORY_BASIC_INFORMATION info; + if ( + (VirtualQuery(const_cast(object), &info, sizeof(info)) == 0) || + (info.State != MEM_COMMIT)) + { + const char msg[] = "Memory was released before process exit\n"; + DWORD written; + WriteFile( + GetStdHandle(STD_ERROR_HANDLE), + msg, + sizeof(msg) - 1, + &written, + nullptr); + TerminateProcess(GetCurrentProcess(), 1); + } + *object = *object + 1; + } + return TRUE; +} diff --git a/src/test/func/dll_teardown/dll_teardown.cc b/src/test/func/dll_teardown/dll_teardown.cc new file mode 100644 index 000000000..3ed7b988a --- /dev/null +++ b/src/test/func/dll_teardown/dll_teardown.cc @@ -0,0 +1,110 @@ +/** + * Checks when a DLL that uses snmalloc returns its memory to the OS. + * + * When the DLL is unloaded with FreeLibrary while the process continues to + * run, the memory it reserved must be released. + * + * When the process exits, the memory must not be released, as code run later + * by the loader may still access allocations. Here, a second DLL that was + * loaded first, so is detached last, accesses an allocation from the first + * DLL in its process-exit detach notification. + */ + +#include +#include + +#ifdef _WIN32 +# ifndef WIN32_LEAN_AND_MEAN +# define WIN32_LEAN_AND_MEAN +# endif +# include + +using AllocFn = void* (*)(size_t); +using DeallocFn = void (*)(void*); +using SetObjectFn = void (*)(size_t*); + +template +Fn get_function(HMODULE module, const char* name) +{ + FARPROC proc = GetProcAddress(module, name); + if (proc == nullptr) + { + printf("GetProcAddress(%s) failed\n", name); + exit(1); + } + return reinterpret_cast(reinterpret_cast(proc)); +} + +HMODULE load(const char* name) +{ + HMODULE module = LoadLibraryA(name); + if (module == nullptr) + { + printf("LoadLibrary(%s) failed: %lu\n", name, GetLastError()); + exit(1); + } + return module; +} + +DWORD state(void* p) +{ + MEMORY_BASIC_INFORMATION info; + if (VirtualQuery(p, &info, sizeof(info)) == 0) + { + printf("VirtualQuery failed: %lu\n", GetLastError()); + exit(1); + } + return info.State; +} + +constexpr const char* alloc_dll = "snmalloc-dll-teardown-alloc.dll"; +constexpr const char* touch_dll = "snmalloc-dll-teardown-touch.dll"; + +int main() +{ + // The touch DLL is loaded first, so it is detached after the alloc DLL + // during process exit. + HMODULE touch = load(touch_dll); + + // FreeLibrary while the process continues to run releases the memory. + { + HMODULE alloc = load(alloc_dll); + auto p = get_function(alloc, "dll_teardown_alloc")(sizeof(size_t)); + if (state(p) != MEM_COMMIT) + { + printf("Allocation is not committed\n"); + return 1; + } + get_function(alloc, "dll_teardown_dealloc")(p); + + FreeLibrary(alloc); + if (GetModuleHandleA(alloc_dll) != nullptr) + { + // Some runtimes, such as MinGW's, pin DLLs that use thread_local + // destructors, so FreeLibrary does not unload them. + printf("DLL was not unloaded; skipping release check\n"); + } + else if (state(p) != MEM_FREE) + { + printf("Memory was not released by FreeLibrary\n"); + return 1; + } + } + + // Process exit must not release the memory, as the touch DLL accesses it + // after the alloc DLL has been detached. + HMODULE alloc = load(alloc_dll); + auto p = static_cast( + get_function(alloc, "dll_teardown_alloc")(sizeof(size_t))); + *p = 0; + get_function(touch, "dll_teardown_set_object")(p); + + printf("Exiting\n"); + return 0; +} +#else +int main() +{ + return 0; +} +#endif diff --git a/src/test/func/fls_teardown/fls_teardown.cc b/src/test/func/fls_teardown/fls_teardown.cc new file mode 100644 index 000000000..15039de3d --- /dev/null +++ b/src/test/func/fls_teardown/fls_teardown.cc @@ -0,0 +1,64 @@ +/** + * Windows fiber-local storage (FLS) callbacks for the main thread run inside + * ExitProcess, after the CRT has run atexit handlers and static destructors. + * Rust's thread-local destructors are implemented this way, so this test + * checks that memory allocated by snmalloc can still be accessed and freed + * from such a callback. + */ + +#include +#include + +#ifdef _WIN32 +# ifndef WIN32_LEAN_AND_MEAN +# define WIN32_LEAN_AND_MEAN +# endif +# ifndef NOMINMAX +# define NOMINMAX +# endif +# include + +struct Object +{ + volatile size_t count; +}; + +void WINAPI fls_callback(void* p) +{ + if (p == nullptr) + return; + + auto* o = static_cast(p); + // Mirrors dropping a Rust Arc: an access to the object, then a free. + o->count = o->count - 1; + snmalloc::dealloc(o); +} + +int main() +{ + DWORD key = FlsAlloc(&fls_callback); + if (key == FLS_OUT_OF_INDEXES) + { + printf("FlsAlloc failed\n"); + return 1; + } + + auto* o = static_cast(snmalloc::alloc(sizeof(Object))); + o->count = 1; + + if (!FlsSetValue(key, o)) + { + printf("FlsSetValue failed\n"); + return 1; + } + + // Returning from main runs the CRT exit path, then ExitProcess, which runs + // fls_callback for this thread. + return 0; +} +#else +int main() +{ + return 0; +} +#endif