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