1 // Copyright 2022 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/freeslot_bitmap.h"
6 
7 #include <cstdint>
8 #include <limits>
9 
10 #include "partition_alloc/freeslot_bitmap_constants.h"
11 #include "partition_alloc/partition_alloc.h"
12 #include "partition_alloc/partition_alloc_buildflags.h"
13 #include "partition_alloc/partition_alloc_constants.h"
14 #include "partition_alloc/partition_alloc_forward.h"
15 #include "partition_alloc/partition_page.h"
16 #include "partition_alloc/partition_root.h"
17 #include "testing/gtest/include/gtest/gtest.h"
18 
19 // This test is disabled when MEMORY_TOOL_REPLACES_ALLOCATOR is defined because
20 // we cannot locate the freeslot bitmap address in that case.
21 #if BUILDFLAG(USE_FREESLOT_BITMAP) && !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
22 
23 namespace partition_alloc::internal {
24 
25 namespace {
26 
27 class PartitionAllocFreeSlotBitmapTest : public ::testing::Test {
28  protected:
29   static constexpr FreeSlotBitmapCellType kAllUsed = 0u;
30   static constexpr FreeSlotBitmapCellType kAllFree =
31       std::numeric_limits<FreeSlotBitmapCellType>::max();
32 
SetUp()33   void SetUp() override {
34     // Allocates memory and creates a pseudo superpage in it. We need to
35     // allocate |2 * kSuperPageSize| so that a whole superpage is contained in
36     // the allocated region.
37     allocator_.init(PartitionOptions{});
38     allocated_ptr_ = reinterpret_cast<uintptr_t>(
39         allocator_.root()->Alloc(2 * kSuperPageSize));
40     super_page_ = (allocated_ptr_ + kSuperPageSize) & kSuperPageBaseMask;
41 
42     // Checks that the whole superpage is in the allocated region.
43     PA_DCHECK(super_page_ + kSuperPageSize <=
44               allocated_ptr_ + 2 * kSuperPageSize);
45   }
46 
TearDown()47   void TearDown() override {
48     allocator_.root()->Free(reinterpret_cast<void*>(allocated_ptr_));
49   }
50 
51   // Returns the |index|-th slot address in the virtual superpage. It assumes
52   // that there are no slot spans and the superpage is only filled with the slot
53   // of size |kSmallestBucket|.
SlotAddr(size_t index)54   uintptr_t SlotAddr(size_t index) {
55     return SuperPagePayloadBegin(super_page_, false) + index * kSmallestBucket;
56   }
57 
58   // Returns the last slot address in the virtual superpage. It assumes that
59   // there are no slot spans but the superpage is only filled with the slot of
60   // size |kSmallestBucket|.
LastSlotAddr()61   uintptr_t LastSlotAddr() {
62     return super_page_ + kSuperPageSize - PartitionPageSize() - kSmallestBucket;
63   }
64 
65  private:
66   uintptr_t allocated_ptr_;
67   uintptr_t super_page_;
68   PartitionAllocator allocator_;
69 };
70 
71 }  // namespace
72 
TEST_F(PartitionAllocFreeSlotBitmapTest,MarkFirstSlotAsUsed)73 TEST_F(PartitionAllocFreeSlotBitmapTest, MarkFirstSlotAsUsed) {
74   uintptr_t slot_addr = SlotAddr(0);
75   FreeSlotBitmapMarkSlotAsFree(slot_addr);
76   EXPECT_FALSE(FreeSlotBitmapSlotIsUsed(slot_addr));
77 
78   FreeSlotBitmapMarkSlotAsUsed(slot_addr);
79   EXPECT_TRUE(FreeSlotBitmapSlotIsUsed(slot_addr));
80 }
81 
TEST_F(PartitionAllocFreeSlotBitmapTest,MarkFirstSlotAsFree)82 TEST_F(PartitionAllocFreeSlotBitmapTest, MarkFirstSlotAsFree) {
83   uintptr_t slot_addr = SlotAddr(0);
84   // All slots are set to "used" by default.
85   EXPECT_TRUE(FreeSlotBitmapSlotIsUsed(slot_addr));
86 
87   FreeSlotBitmapMarkSlotAsFree(slot_addr);
88   EXPECT_FALSE(FreeSlotBitmapSlotIsUsed(slot_addr));
89 }
90 
TEST_F(PartitionAllocFreeSlotBitmapTest,MarkAllBitsInCellAsUsed)91 TEST_F(PartitionAllocFreeSlotBitmapTest, MarkAllBitsInCellAsUsed) {
92   const size_t kFirstSlotAddr = SlotAddr(0);
93   const size_t kLastSlotAddr = SlotAddr(kFreeSlotBitmapBitsPerCell);
94 
95   auto [cell_first_slot, bit_index_first_slot] =
96       GetFreeSlotBitmapCellPtrAndBitIndex(kFirstSlotAddr);
97   auto [cell_last_slot, bit_index_last_slot] =
98       GetFreeSlotBitmapCellPtrAndBitIndex(kLastSlotAddr);
99 
100   // Check that the bit corresponding to |kFirstSlotAddr| is the first bit in
101   // some cell (= |cell_first_slot|), and the bit for |kLastSlotAddr| is the
102   // first bit in the next cell. This means that we are manipulating all the
103   // bits in |cell_first_slot| in this test.
104   EXPECT_EQ(0u, bit_index_first_slot);
105   EXPECT_EQ(0u, bit_index_last_slot);
106   EXPECT_NE(cell_first_slot, cell_last_slot);
107 
108   for (size_t slot_addr = kFirstSlotAddr; slot_addr < kLastSlotAddr;
109        slot_addr += kSmallestBucket) {
110     FreeSlotBitmapMarkSlotAsFree(slot_addr);
111   }
112 
113   // Check all the bits in |cell_first_slot| are 1 (= free).
114   EXPECT_EQ(kAllFree, *cell_first_slot);
115 
116   for (size_t slot_addr = kFirstSlotAddr; slot_addr < kLastSlotAddr;
117        slot_addr += kSmallestBucket) {
118     FreeSlotBitmapMarkSlotAsUsed(slot_addr);
119   }
120 
121   // Check all the bits in |cell_first_slot| are 0 (= used).
122   EXPECT_EQ(kAllUsed, *cell_first_slot);
123 }
124 
TEST_F(PartitionAllocFreeSlotBitmapTest,MarkLastSlotAsUsed)125 TEST_F(PartitionAllocFreeSlotBitmapTest, MarkLastSlotAsUsed) {
126   uintptr_t last_slot_addr = LastSlotAddr();
127   FreeSlotBitmapMarkSlotAsFree(last_slot_addr);
128   EXPECT_FALSE(FreeSlotBitmapSlotIsUsed(last_slot_addr));
129 
130   FreeSlotBitmapMarkSlotAsUsed(last_slot_addr);
131   EXPECT_TRUE(FreeSlotBitmapSlotIsUsed(last_slot_addr));
132 }
133 
TEST_F(PartitionAllocFreeSlotBitmapTest,ResetBitmap)134 TEST_F(PartitionAllocFreeSlotBitmapTest, ResetBitmap) {
135   const size_t kNumSlots = 3 * kFreeSlotBitmapBitsPerCell;
136   for (size_t i = 0; i < kNumSlots; ++i) {
137     FreeSlotBitmapMarkSlotAsFree(SlotAddr(i));
138   }
139 
140   auto [cell_first_slot, bit_index_first_slot] =
141       GetFreeSlotBitmapCellPtrAndBitIndex(SlotAddr(0));
142   EXPECT_EQ(0u, bit_index_first_slot);
143   EXPECT_EQ(kAllFree, *cell_first_slot);
144   EXPECT_EQ(kAllFree, *(cell_first_slot + 1));
145   EXPECT_EQ(kAllFree, *(cell_first_slot + 2));
146 
147   FreeSlotBitmapReset(SlotAddr(kFreeSlotBitmapBitsPerCell),
148                       SlotAddr(2 * kFreeSlotBitmapBitsPerCell),
149                       kSmallestBucket);
150   EXPECT_EQ(kAllFree, *cell_first_slot);
151   EXPECT_EQ(kAllUsed, *(cell_first_slot + 1));
152   EXPECT_EQ(kAllFree, *(cell_first_slot + 2));
153 }
154 
TEST_F(PartitionAllocFreeSlotBitmapTest,ResetBitmapWithZeroLengthRange)155 TEST_F(PartitionAllocFreeSlotBitmapTest, ResetBitmapWithZeroLengthRange) {
156   const size_t kNumSlots = 3 * kFreeSlotBitmapBitsPerCell;
157   for (size_t i = 0; i < kNumSlots; ++i) {
158     FreeSlotBitmapMarkSlotAsFree(SlotAddr(i));
159   }
160 
161   // Test with an aligned address.
162   uintptr_t aligned_addr = SlotAddr(0);
163   auto [cell_aligned, bit_index_aligned] =
164       GetFreeSlotBitmapCellPtrAndBitIndex(aligned_addr);
165   EXPECT_EQ(0u, bit_index_aligned);
166   EXPECT_EQ(kAllFree, *cell_aligned);
167 
168   FreeSlotBitmapReset(aligned_addr, aligned_addr, kSmallestBucket);
169   EXPECT_EQ(kAllFree, *cell_aligned);
170 
171   // Test with a non-aligned address.
172   uintptr_t non_aligned_addr = SlotAddr(1);
173   auto [cell_non_aligned, bit_index_non_aligned] =
174       GetFreeSlotBitmapCellPtrAndBitIndex(non_aligned_addr);
175   EXPECT_EQ(1u, bit_index_non_aligned);
176   EXPECT_EQ(kAllFree, *cell_non_aligned);
177 
178   FreeSlotBitmapReset(non_aligned_addr, non_aligned_addr, kSmallestBucket);
179   EXPECT_EQ(kAllFree, *cell_non_aligned);
180 }
181 
TEST_F(PartitionAllocFreeSlotBitmapTest,ResetSingleBitInMiddleOfCell)182 TEST_F(PartitionAllocFreeSlotBitmapTest, ResetSingleBitInMiddleOfCell) {
183   const size_t kNumSlots = 3 * kFreeSlotBitmapBitsPerCell;
184   for (size_t i = 0; i < kNumSlots; ++i) {
185     FreeSlotBitmapMarkSlotAsFree(SlotAddr(i));
186   }
187 
188   // Choose a slot address that is in the middle of a cell.
189   uintptr_t mid_cell_slot_addr = SlotAddr(kFreeSlotBitmapBitsPerCell / 2);
190 
191   auto [cell_mid, bit_index_mid] =
192       GetFreeSlotBitmapCellPtrAndBitIndex(mid_cell_slot_addr);
193   EXPECT_NE(0u, bit_index_mid);
194   EXPECT_TRUE(*cell_mid & CellWithAOne(bit_index_mid));
195 
196   // Reset the single bit in the middle of the cell.
197   FreeSlotBitmapReset(mid_cell_slot_addr, mid_cell_slot_addr + kSmallestBucket,
198                       kSmallestBucket);
199 
200   EXPECT_FALSE(*cell_mid & CellWithAOne(bit_index_mid));
201 }
202 
203 }  // namespace partition_alloc::internal
204 
205 #endif  // BUILDFLAG(USE_FREESLOT_BITMAP) &&
206         // !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
207