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 "partition_alloc/shim/nonscannable_allocator.h"
6 
7 #include "partition_alloc/partition_root.h"
8 
9 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
10 #include "partition_alloc/shim/allocator_shim_default_dispatch_to_partition_alloc.h"
11 
12 #if BUILDFLAG(USE_STARSCAN)
13 #include "partition_alloc/internal_allocator.h"
14 #include "partition_alloc/starscan/pcscan.h"
15 #endif
16 #endif  // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
17 
18 namespace allocator_shim::internal {
19 
20 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
21 template <bool quarantinable>
22 NonScannableAllocatorImpl<quarantinable>::NonScannableAllocatorImpl() = default;
23 template <bool quarantinable>
24 NonScannableAllocatorImpl<quarantinable>::~NonScannableAllocatorImpl() =
25     default;
26 
27 template <bool quarantinable>
28 NonScannableAllocatorImpl<quarantinable>&
Instance()29 NonScannableAllocatorImpl<quarantinable>::Instance() {
30   static partition_alloc::internal::base::NoDestructor<
31       NonScannableAllocatorImpl>
32       instance;
33   return *instance;
34 }
35 
36 template <bool quarantinable>
Alloc(size_t size)37 void* NonScannableAllocatorImpl<quarantinable>::Alloc(size_t size) {
38 #if BUILDFLAG(USE_STARSCAN)
39   // TODO(bikineev): Change to LIKELY once PCScan is enabled by default.
40   if (PA_UNLIKELY(pcscan_enabled_.load(std::memory_order_acquire))) {
41     PA_DCHECK(allocator_.get());
42     return allocator_->root()
43         ->AllocInline<partition_alloc::AllocFlags::kNoHooks>(size);
44   }
45 #endif  // BUILDFLAG(USE_STARSCAN)
46   // Otherwise, dispatch to default partition.
47   return allocator_shim::internal::PartitionAllocMalloc::Allocator()
48       ->AllocInline<partition_alloc::AllocFlags::kNoHooks>(size);
49 }
50 
51 template <bool quarantinable>
Free(void * ptr)52 void NonScannableAllocatorImpl<quarantinable>::Free(void* ptr) {
53 #if BUILDFLAG(USE_STARSCAN)
54   if (PA_UNLIKELY(pcscan_enabled_.load(std::memory_order_acquire))) {
55     allocator_->root()->FreeInline<partition_alloc::FreeFlags::kNoHooks>(ptr);
56     return;
57   }
58 #endif  // BUILDFLAG(USE_STARSCAN)
59   partition_alloc::PartitionRoot::FreeInlineInUnknownRoot<
60       partition_alloc::FreeFlags::kNoHooks>(ptr);
61 }
62 
63 template <bool quarantinable>
NotifyPCScanEnabled()64 void NonScannableAllocatorImpl<quarantinable>::NotifyPCScanEnabled() {
65 #if BUILDFLAG(USE_STARSCAN)
66   partition_alloc::PartitionOptions opts;
67   opts.star_scan_quarantine =
68       quarantinable ? partition_alloc::PartitionOptions::kAllowed
69                     : partition_alloc::PartitionOptions::kDisallowed;
70   opts.backup_ref_ptr = partition_alloc::PartitionOptions::kDisabled;
71   allocator_.reset(partition_alloc::internal::ConstructAtInternalPartition<
72                    partition_alloc::PartitionAllocator>(opts));
73   if constexpr (quarantinable) {
74     partition_alloc::internal::PCScan::RegisterNonScannableRoot(
75         allocator_->root());
76   }
77   pcscan_enabled_.store(true, std::memory_order_release);
78 #endif  // BUILDFLAG(USE_STARSCAN)
79 }
80 
81 template class PA_EXPORT_TEMPLATE_DEFINE(PA_COMPONENT_EXPORT(ALLOCATOR_SHIM))
82     NonScannableAllocatorImpl<true>;
83 template class PA_EXPORT_TEMPLATE_DEFINE(PA_COMPONENT_EXPORT(ALLOCATOR_SHIM))
84     NonScannableAllocatorImpl<false>;
85 
86 #endif  // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
87 
88 }  // namespace allocator_shim::internal
89