1 // Copyright 2023 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/lightweight_quarantine.h"
6
7 #include "partition_alloc/partition_alloc_for_testing.h"
8 #include "partition_alloc/partition_page.h"
9 #include "partition_alloc/partition_root.h"
10 #include "partition_alloc/partition_stats.h"
11 #include "testing/gtest/include/gtest/gtest.h"
12
13 namespace partition_alloc {
14
15 #if !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
16
17 namespace {
18
19 using QuarantineRoot = internal::LightweightQuarantineRoot;
20 using QuarantineBranch = internal::LightweightQuarantineBranch;
21
22 struct LightweightQuarantineTestParam {
23 size_t capacity_count;
24 size_t capacity_in_bytes;
25 };
26 constexpr LightweightQuarantineTestParam kSmallQuarantineBranch = {
27 .capacity_count = 1024,
28 .capacity_in_bytes = 256};
29 constexpr LightweightQuarantineTestParam kLargeQuarantineBranch = {
30 .capacity_count = 1024,
31 .capacity_in_bytes = 4096};
32
33 class PartitionAllocLightweightQuarantineTest
34 : public testing::TestWithParam<LightweightQuarantineTestParam> {
35 protected:
SetUp()36 void SetUp() override {
37 const auto param = GetParam();
38
39 allocator_ =
40 std::make_unique<PartitionAllocatorForTesting>(PartitionOptions{});
41
42 root_.emplace(*allocator_->root(), param.capacity_in_bytes);
43 branch_.emplace(root_->CreateBranch(param.capacity_count));
44
45 auto stats = GetStats();
46 ASSERT_EQ(0u, stats.size_in_bytes);
47 ASSERT_EQ(0u, stats.count);
48 ASSERT_EQ(0u, stats.cumulative_size_in_bytes);
49 ASSERT_EQ(0u, stats.cumulative_count);
50 }
51
TearDown()52 void TearDown() override {
53 // |Purge()|d here.
54 branch_.reset();
55 root_.reset();
56 allocator_ = nullptr;
57 }
58
GetPartitionRoot() const59 PartitionRoot* GetPartitionRoot() const { return allocator_->root(); }
60
GetQuarantineRoot()61 QuarantineRoot* GetQuarantineRoot() { return &root_.value(); }
GetQuarantineBranch()62 QuarantineBranch* GetQuarantineBranch() { return &branch_.value(); }
63
Quarantine(void * object)64 bool Quarantine(void* object) {
65 auto* slot_span = internal::SlotSpanMetadata::FromObject(object);
66 uintptr_t slot_start = GetPartitionRoot()->ObjectToSlotStart(object);
67 return GetQuarantineBranch()->Quarantine(object, slot_span, slot_start);
68 }
69
GetObjectSize(void * object)70 size_t GetObjectSize(void* object) {
71 auto* entry_slot_span = internal::SlotSpanMetadata::FromObject(object);
72 return GetPartitionRoot()->GetSlotUsableSize(entry_slot_span);
73 }
74
GetStats() const75 LightweightQuarantineStats GetStats() const {
76 LightweightQuarantineStats stats{};
77 root_->AccumulateStats(stats);
78 return stats;
79 }
80
81 std::unique_ptr<PartitionAllocatorForTesting> allocator_;
82 std::optional<QuarantineRoot> root_;
83 std::optional<QuarantineBranch> branch_;
84 };
85 INSTANTIATE_TEST_SUITE_P(
86 PartitionAllocLightweightQuarantineTestMultipleQuarantineSizeInstantiation,
87 PartitionAllocLightweightQuarantineTest,
88 ::testing::Values(kSmallQuarantineBranch, kLargeQuarantineBranch));
89
90 } // namespace
91
TEST_P(PartitionAllocLightweightQuarantineTest,Basic)92 TEST_P(PartitionAllocLightweightQuarantineTest, Basic) {
93 constexpr size_t kObjectSize = 1;
94
95 const size_t capacity_in_bytes =
96 GetQuarantineBranch()->GetRoot().GetCapacityInBytes();
97
98 constexpr size_t kCount = 100;
99 for (size_t i = 1; i <= kCount; i++) {
100 void* object = GetPartitionRoot()->Alloc(kObjectSize);
101 const size_t size = GetObjectSize(object);
102 const size_t max_count = capacity_in_bytes / size;
103
104 const bool success = Quarantine(object);
105
106 ASSERT_TRUE(success);
107 ASSERT_TRUE(GetQuarantineBranch()->IsQuarantinedForTesting(object));
108
109 const auto expected_count = std::min(i, max_count);
110 auto stats = GetStats();
111 ASSERT_EQ(expected_count * size, stats.size_in_bytes);
112 ASSERT_EQ(expected_count, stats.count);
113 ASSERT_EQ(i * size, stats.cumulative_size_in_bytes);
114 ASSERT_EQ(i, stats.cumulative_count);
115 }
116 }
117
TEST_P(PartitionAllocLightweightQuarantineTest,TooLargeAllocation)118 TEST_P(PartitionAllocLightweightQuarantineTest, TooLargeAllocation) {
119 constexpr size_t kObjectSize = 1 << 26; // 64 MiB.
120 const size_t capacity_in_bytes =
121 GetQuarantineBranch()->GetRoot().GetCapacityInBytes();
122
123 void* object = GetPartitionRoot()->Alloc(kObjectSize);
124 const size_t size = GetObjectSize(object);
125 ASSERT_GT(size, capacity_in_bytes);
126
127 const bool success = Quarantine(object);
128
129 ASSERT_FALSE(success);
130 ASSERT_FALSE(GetQuarantineBranch()->IsQuarantinedForTesting(object));
131
132 auto stats = GetStats();
133 ASSERT_EQ(0u, stats.size_in_bytes);
134 ASSERT_EQ(0u, stats.count);
135 ASSERT_EQ(0u, stats.cumulative_size_in_bytes);
136 ASSERT_EQ(0u, stats.cumulative_count);
137 }
138
139 #endif // !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
140
141 } // namespace partition_alloc
142