1 // Copyright 2020 The Chromium Authors
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4
5 #include "partition_alloc/partition_tls.h"
6
7 #include <windows.h>
8
9 namespace partition_alloc::internal {
10
11 namespace {
12
13 // Store the key as the thread destruction callback doesn't get it.
14 PartitionTlsKey g_key;
15 void (*g_destructor)(void*) = nullptr;
16 void (*g_on_dll_process_detach)() = nullptr;
17
18 // Static callback function to call with each thread termination.
PartitionTlsOnThreadExit(PVOID module,DWORD reason,PVOID reserved)19 void NTAPI PartitionTlsOnThreadExit(PVOID module,
20 DWORD reason,
21 PVOID reserved) {
22 if (reason != DLL_THREAD_DETACH && reason != DLL_PROCESS_DETACH) {
23 return;
24 }
25
26 if (reason == DLL_PROCESS_DETACH && g_on_dll_process_detach) {
27 g_on_dll_process_detach();
28 }
29
30 if (g_destructor) {
31 void* per_thread_data = PartitionTlsGet(g_key);
32 if (per_thread_data) {
33 g_destructor(per_thread_data);
34 }
35 }
36 }
37
38 } // namespace
39
PartitionTlsCreate(PartitionTlsKey * key,void (* destructor)(void *))40 bool PartitionTlsCreate(PartitionTlsKey* key, void (*destructor)(void*)) {
41 PA_CHECK(g_destructor == nullptr); // Only one TLS key supported at a time.
42 PartitionTlsKey value = TlsAlloc();
43 if (value != TLS_OUT_OF_INDEXES) {
44 *key = value;
45
46 g_key = value;
47 g_destructor = destructor;
48 return true;
49 }
50 return false;
51 }
52
PartitionTlsSetOnDllProcessDetach(void (* callback)())53 void PartitionTlsSetOnDllProcessDetach(void (*callback)()) {
54 g_on_dll_process_detach = callback;
55 }
56
57 } // namespace partition_alloc::internal
58
59 // See thread_local_storage_win.cc for details and reference.
60 //
61 // The callback has to be in any section between .CRT$XLA and .CRT$XLZ, as these
62 // are sentinels used by the TLS code to find the callback array bounds. As we
63 // don't particularly care about where we are called but would prefer to be
64 // deinitialized towards the end (in particular after Chromium's TLS), we locate
65 // ourselves in .CRT$XLY.
66
67 // Force a reference to _tls_used to make the linker create the TLS directory if
68 // it's not already there. (e.g. if __declspec(thread) is not used). Force a
69 // reference to partition_tls_thread_exit_callback to prevent whole program
70 // optimization from discarding the variable.
71 #ifdef _WIN64
72
73 #pragma comment(linker, "/INCLUDE:_tls_used")
74 #pragma comment(linker, "/INCLUDE:partition_tls_thread_exit_callback")
75
76 #else // _WIN64
77
78 #pragma comment(linker, "/INCLUDE:__tls_used")
79 #pragma comment(linker, "/INCLUDE:_partition_tls_thread_exit_callback")
80
81 #endif // _WIN64
82
83 // extern "C" suppresses C++ name mangling so we know the symbol name for the
84 // linker /INCLUDE:symbol pragma above.
85 extern "C" {
86 // The linker must not discard partition_tls_thread_exit_callback. (We force a
87 // reference to this variable with a linker /INCLUDE:symbol pragma to ensure
88 // that.) If this variable is discarded, PartitionTlsOnThreadExit will never be
89 // called.
90 #ifdef _WIN64
91
92 // .CRT section is merged with .rdata on x64 so it must be constant data.
93 #pragma const_seg(".CRT$XLY")
94 // When defining a const variable, it must have external linkage to be sure the
95 // linker doesn't discard it.
96 extern const PIMAGE_TLS_CALLBACK partition_tls_thread_exit_callback;
97 const PIMAGE_TLS_CALLBACK partition_tls_thread_exit_callback =
98 partition_alloc::internal::PartitionTlsOnThreadExit;
99
100 // Reset the default section.
101 #pragma const_seg()
102
103 #else // _WIN64
104
105 #pragma data_seg(".CRT$XLY")
106 PIMAGE_TLS_CALLBACK partition_tls_thread_exit_callback =
107 partition_alloc::internal::PartitionTlsOnThreadExit;
108
109 // Reset the default section.
110 #pragma data_seg()
111
112 #endif // _WIN64
113 } // extern "C"
114