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