1 // Copyright 2021 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 <cstddef>
6 #include <cstring>
7
8 #include "build/build_config.h"
9 #include "build/rust/std/alias.h"
10 #include "build/rust/std/buildflags.h"
11 #include "build/rust/std/immediate_crash.h"
12
13 #if BUILDFLAG(RUST_ALLOCATOR_USES_PARTITION_ALLOC)
14 #include "partition_alloc/partition_alloc_constants.h" // nogncheck
15 #include "partition_alloc/shim/allocator_shim.h" // nogncheck
16 #elif BUILDFLAG(IS_WIN)
17 #include <cstdlib>
18 #endif
19
20 // When linking a final binary, rustc has to pick between either:
21 // * The default Rust allocator
22 // * Any #[global_allocator] defined in *any rlib in its dependency tree*
23 // (https://doc.rust-lang.org/edition-guide/rust-2018/platform-and-target-support/global-allocators.html)
24 //
25 // In this latter case, this fact will be recorded in some of the metadata
26 // within the .rlib file. (An .rlib file is just a .a file, but does have
27 // additional metadata for use by rustc. This is, as far as I know, the only
28 // such metadata we would ideally care about.)
29 //
30 // In all the linked rlibs,
31 // * If 0 crates define a #[global_allocator], rustc uses its default allocator
32 // * If 1 crate defines a #[global_allocator], rustc uses that
33 // * If >1 crates define a #[global_allocator], rustc bombs out.
34 //
35 // Because rustc does these checks, it doesn't just have the __rust_alloc
36 // symbols defined anywhere (neither in the stdlib nor in any of these
37 // crates which have a #[global_allocator] defined.)
38 //
39 // Instead:
40 // Rust's final linking stage invokes dynamic LLVM codegen to create symbols
41 // for the basic heap allocation operations. It literally creates a
42 // __rust_alloc symbol at link time. Unless any crate has specified a
43 // #[global_allocator], it simply calls from __rust_alloc into
44 // __rdl_alloc, which is the default Rust allocator. The same applies to a
45 // few other symbols.
46 //
47 // We're not (always) using rustc for final linking. For cases where we're not
48 // Rustc as the final linker, we'll define those symbols here instead. This
49 // allows us to redirect allocation to PartitionAlloc if clang is doing the
50 // link.
51 //
52 // We use unchecked allocation paths in PartitionAlloc rather than going through
53 // its shims in `malloc()` etc so that we can support fallible allocation paths
54 // such as Vec::try_reserve without crashing on allocation failure.
55 //
56 // In future, we should build a crate with a #[global_allocator] and
57 // redirect these symbols back to Rust in order to use to that crate instead.
58 // This would allow Rust-linked executables to:
59 // 1. Use PartitionAlloc on Windows. The stdlib uses Windows heap functions
60 // directly that PartitionAlloc can not intercept.
61 // 2. Have `Vec::try_reserve` to fail at runtime on Linux instead of crashing in
62 // malloc() where PartitionAlloc replaces that function.
63 //
64 // They're weak symbols, because this file will sometimes end up in targets
65 // which are linked by rustc, and thus we would otherwise get duplicate
66 // definitions. The following definitions will therefore only end up being
67 // used in targets which are linked by our C++ toolchain.
68 //
69 // # On Windows ASAN
70 //
71 // In ASAN builds, PartitionAlloc-Everywhere is disabled, meaning malloc() and
72 // friends in C++ do not go to PartitionAlloc. So we also don't point the Rust
73 // allocation functions at PartitionAlloc. Generally, this means we just direct
74 // them to the Standard Library's allocator.
75 //
76 // However, on Windows the Standard Library uses HeapAlloc() and Windows ASAN
77 // does *not* hook that method, so ASAN does not get to hear about allocations
78 // made in Rust. To resolve this, we redirect allocation to _aligned_malloc
79 // which Windows ASAN *does* hook.
80 //
81 // Note that there is a runtime option to make ASAN hook HeapAlloc() but
82 // enabling it breaks Win32 APIs like CreateProcess:
83 // https://issues.chromium.org/u/1/issues/368070343#comment29
84
85 extern "C" {
86
87 #ifdef COMPONENT_BUILD
88 #if BUILDFLAG(IS_WIN)
89 #define REMAP_ALLOC_ATTRIBUTES __declspec(dllexport) __attribute__((weak))
90 #else
91 #define REMAP_ALLOC_ATTRIBUTES \
92 __attribute__((visibility("default"))) __attribute__((weak))
93 #endif
94 #else
95 #define REMAP_ALLOC_ATTRIBUTES __attribute__((weak))
96 #endif // COMPONENT_BUILD
97
98 #if !BUILDFLAG(RUST_ALLOCATOR_USES_PARTITION_ALLOC) && BUILDFLAG(IS_WIN) && \
99 defined(ADDRESS_SANITIZER)
100 #define USE_WIN_ALIGNED_MALLOC 1
101 #else
102 #define USE_WIN_ALIGNED_MALLOC 0
103 #endif
104
105 // This must exist as the stdlib depends on it to prove that we know the
106 // alloc shims below are unstable. In the future we may be required to replace
107 // them with a #[global_allocator] crate (see file comment above for more).
108 //
109 // Marked as weak as when Rust drives linking it includes this symbol itself,
110 // and we don't want a collision due to C++ being in the same link target, where
111 // C++ causes us to explicitly link in the stdlib and this symbol here.
112 [[maybe_unused]]
113 __attribute__((weak)) unsigned char __rust_no_alloc_shim_is_unstable;
114
__rust_alloc(size_t size,size_t align)115 REMAP_ALLOC_ATTRIBUTES void* __rust_alloc(size_t size, size_t align) {
116 #if BUILDFLAG(RUST_ALLOCATOR_USES_PARTITION_ALLOC)
117 // PartitionAlloc will crash if given an alignment larger than this.
118 if (align > partition_alloc::internal::kMaxSupportedAlignment) {
119 return nullptr;
120 }
121
122 if (align <= alignof(std::max_align_t)) {
123 return allocator_shim::UncheckedAlloc(size);
124 } else {
125 return allocator_shim::UncheckedAlignedAlloc(size, align);
126 }
127 #elif USE_WIN_ALIGNED_MALLOC
128 return _aligned_malloc(size, align);
129 #else
130 extern void* __rdl_alloc(size_t size, size_t align);
131 return __rdl_alloc(size, align);
132 #endif
133 }
134
__rust_dealloc(void * p,size_t size,size_t align)135 REMAP_ALLOC_ATTRIBUTES void __rust_dealloc(void* p, size_t size, size_t align) {
136 #if BUILDFLAG(RUST_ALLOCATOR_USES_PARTITION_ALLOC)
137 if (align <= alignof(std::max_align_t)) {
138 allocator_shim::UncheckedFree(p);
139 } else {
140 allocator_shim::UncheckedAlignedFree(p);
141 }
142 #elif USE_WIN_ALIGNED_MALLOC
143 return _aligned_free(p);
144 #else
145 extern void __rdl_dealloc(void* p, size_t size, size_t align);
146 __rdl_dealloc(p, size, align);
147 #endif
148 }
149
__rust_realloc(void * p,size_t old_size,size_t align,size_t new_size)150 REMAP_ALLOC_ATTRIBUTES void* __rust_realloc(void* p,
151 size_t old_size,
152 size_t align,
153 size_t new_size) {
154 #if BUILDFLAG(RUST_ALLOCATOR_USES_PARTITION_ALLOC)
155 if (align <= alignof(std::max_align_t)) {
156 return allocator_shim::UncheckedRealloc(p, new_size);
157 } else {
158 return allocator_shim::UncheckedAlignedRealloc(p, new_size, align);
159 }
160 #elif USE_WIN_ALIGNED_MALLOC
161 return _aligned_realloc(p, new_size, align);
162 #else
163 extern void* __rdl_realloc(void* p, size_t old_size, size_t align,
164 size_t new_size);
165 return __rdl_realloc(p, old_size, align, new_size);
166 #endif
167 }
168
__rust_alloc_zeroed(size_t size,size_t align)169 REMAP_ALLOC_ATTRIBUTES void* __rust_alloc_zeroed(size_t size, size_t align) {
170 #if BUILDFLAG(RUST_ALLOCATOR_USES_PARTITION_ALLOC) || USE_WIN_ALIGNED_MALLOC
171 // TODO(danakj): When RUST_ALLOCATOR_USES_PARTITION_ALLOC is true, it's
172 // possible that a partition_alloc::UncheckedAllocZeroed() call would perform
173 // better than partition_alloc::UncheckedAlloc() + memset. But there is no
174 // such API today. See b/342251590.
175 void* p = __rust_alloc(size, align);
176 if (p) {
177 memset(p, 0, size);
178 }
179 return p;
180 #else
181 extern void* __rdl_alloc_zeroed(size_t size, size_t align);
182 return __rdl_alloc_zeroed(size, align);
183 #endif
184 }
185
__rust_alloc_error_handler(size_t size,size_t align)186 REMAP_ALLOC_ATTRIBUTES void __rust_alloc_error_handler(size_t size,
187 size_t align) {
188 NO_CODE_FOLDING();
189 IMMEDIATE_CRASH();
190 }
191
192 REMAP_ALLOC_ATTRIBUTES extern const unsigned char
193 __rust_alloc_error_handler_should_panic = 0;
194
195 } // extern "C"
196