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