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_alloc_hooks.h"
6 
7 #include <ostream>
8 
9 #include "partition_alloc/partition_alloc_check.h"
10 #include "partition_alloc/partition_lock.h"
11 
12 namespace partition_alloc {
13 
14 namespace {
15 
16 internal::Lock g_hook_lock;
17 
GetHooksLock()18 internal::Lock& GetHooksLock() {
19   return g_hook_lock;
20 }
21 
22 }  // namespace
23 
24 std::atomic<bool> PartitionAllocHooks::hooks_enabled_(false);
25 std::atomic<PartitionAllocHooks::AllocationObserverHook*>
26     PartitionAllocHooks::allocation_observer_hook_(nullptr);
27 std::atomic<PartitionAllocHooks::FreeObserverHook*>
28     PartitionAllocHooks::free_observer_hook_(nullptr);
29 std::atomic<PartitionAllocHooks::AllocationOverrideHook*>
30     PartitionAllocHooks::allocation_override_hook_(nullptr);
31 std::atomic<PartitionAllocHooks::FreeOverrideHook*>
32     PartitionAllocHooks::free_override_hook_(nullptr);
33 std::atomic<PartitionAllocHooks::ReallocOverrideHook*>
34     PartitionAllocHooks::realloc_override_hook_(nullptr);
35 std::atomic<PartitionAllocHooks::QuarantineOverrideHook*>
36     PartitionAllocHooks::quarantine_override_hook_(nullptr);
37 
SetObserverHooks(AllocationObserverHook * alloc_hook,FreeObserverHook * free_hook)38 void PartitionAllocHooks::SetObserverHooks(AllocationObserverHook* alloc_hook,
39                                            FreeObserverHook* free_hook) {
40   internal::ScopedGuard guard(GetHooksLock());
41 
42   // Chained hooks are not supported. Registering a non-null hook when a
43   // non-null hook is already registered indicates somebody is trying to
44   // overwrite a hook.
45   PA_CHECK((!allocation_observer_hook_ && !free_observer_hook_) ||
46            (!alloc_hook && !free_hook))
47       << "Overwriting already set observer hooks";
48   allocation_observer_hook_ = alloc_hook;
49   free_observer_hook_ = free_hook;
50 
51   hooks_enabled_ = allocation_observer_hook_ || allocation_override_hook_;
52 }
53 
SetOverrideHooks(AllocationOverrideHook * alloc_hook,FreeOverrideHook * free_hook,ReallocOverrideHook realloc_hook)54 void PartitionAllocHooks::SetOverrideHooks(AllocationOverrideHook* alloc_hook,
55                                            FreeOverrideHook* free_hook,
56                                            ReallocOverrideHook realloc_hook) {
57   internal::ScopedGuard guard(GetHooksLock());
58 
59   PA_CHECK((!allocation_override_hook_ && !free_override_hook_ &&
60             !realloc_override_hook_) ||
61            (!alloc_hook && !free_hook && !realloc_hook))
62       << "Overwriting already set override hooks";
63   allocation_override_hook_ = alloc_hook;
64   free_override_hook_ = free_hook;
65   realloc_override_hook_ = realloc_hook;
66 
67   hooks_enabled_ = allocation_observer_hook_ || allocation_override_hook_;
68 }
69 
AllocationObserverHookIfEnabled(const partition_alloc::AllocationNotificationData & notification_data)70 void PartitionAllocHooks::AllocationObserverHookIfEnabled(
71     const partition_alloc::AllocationNotificationData& notification_data) {
72   if (auto* hook = allocation_observer_hook_.load(std::memory_order_relaxed)) {
73     hook(notification_data);
74   }
75 }
76 
AllocationOverrideHookIfEnabled(void ** out,AllocFlags flags,size_t size,const char * type_name)77 bool PartitionAllocHooks::AllocationOverrideHookIfEnabled(
78     void** out,
79     AllocFlags flags,
80     size_t size,
81     const char* type_name) {
82   if (auto* hook = allocation_override_hook_.load(std::memory_order_relaxed)) {
83     return hook(out, flags, size, type_name);
84   }
85   return false;
86 }
87 
FreeObserverHookIfEnabled(const FreeNotificationData & notification_data)88 void PartitionAllocHooks::FreeObserverHookIfEnabled(
89     const FreeNotificationData& notification_data) {
90   if (auto* hook = free_observer_hook_.load(std::memory_order_relaxed)) {
91     hook(notification_data);
92   }
93 }
94 
FreeOverrideHookIfEnabled(void * address)95 bool PartitionAllocHooks::FreeOverrideHookIfEnabled(void* address) {
96   if (auto* hook = free_override_hook_.load(std::memory_order_relaxed)) {
97     return hook(address);
98   }
99   return false;
100 }
101 
ReallocObserverHookIfEnabled(const FreeNotificationData & free_notification_data,const AllocationNotificationData & allocation_notification_data)102 void PartitionAllocHooks::ReallocObserverHookIfEnabled(
103     const FreeNotificationData& free_notification_data,
104     const AllocationNotificationData& allocation_notification_data) {
105   // Report a reallocation as a free followed by an allocation.
106   AllocationObserverHook* allocation_hook =
107       allocation_observer_hook_.load(std::memory_order_relaxed);
108   FreeObserverHook* free_hook =
109       free_observer_hook_.load(std::memory_order_relaxed);
110   if (allocation_hook && free_hook) {
111     free_hook(free_notification_data);
112     allocation_hook(allocation_notification_data);
113   }
114 }
115 
ReallocOverrideHookIfEnabled(size_t * out,void * address)116 bool PartitionAllocHooks::ReallocOverrideHookIfEnabled(size_t* out,
117                                                        void* address) {
118   if (ReallocOverrideHook* hook =
119           realloc_override_hook_.load(std::memory_order_relaxed)) {
120     return hook(out, address);
121   }
122   return false;
123 }
124 
125 // Do not unset the hook if there are remaining quarantined slots
126 // not to break checks on unquarantining.
SetQuarantineOverrideHook(QuarantineOverrideHook * hook)127 void PartitionAllocHooks::SetQuarantineOverrideHook(
128     QuarantineOverrideHook* hook) {
129   quarantine_override_hook_.store(hook, std::memory_order_release);
130 }
131 
132 }  // namespace partition_alloc
133