[PA] Fix PartitionRoot::CheckMetadataIntegrity

Calling `PartitionRoot::FromAddrInFirstSuperpage` on a pointer to
middle of a direct-mapped allocation is not allowed.

Bug: 435448631
Change-Id: Ia3a6ecf37350184458c8a18c6a0782338ae75d37
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/6821095
Commit-Queue: Mikihito Matsuura <[email protected]>
Reviewed-by: Keishi Hattori <[email protected]>
Cr-Commit-Position: refs/heads/main@{#1496874}
NOKEYCHECK=True
GitOrigin-RevId: d6a376f7bc36e3b7a9052d8a4170df09e0cd1fe1

Author: mikt

src/partition_alloc/partition_alloc_unittest.cc

@@ -2495,6 +2495,29 @@ TEST_P(PartitionAllocTest, LostFreeSlotSpansBug) {
   EXPECT_TRUE(bucket->decommitted_slot_spans_head);
 }
 
+TEST_P(PartitionAllocTest, CheckMetadataIntegrityPass) {
+  char* const small_ptr =
+      static_cast<char*>(allocator.root()->Alloc(kTestAllocSize));
+  ASSERT_TRUE(small_ptr);
+
+  // Should not crash.
+  PartitionRoot::CheckMetadataIntegrity(small_ptr);
+  PartitionRoot::CheckMetadataIntegrity(small_ptr + kTestAllocSize - 1);
+
+  allocator.root()->Free(small_ptr);
+
+  constexpr size_t kDirectMapSize = BucketIndexLookup::kMaxBucketSize + 1;
+  char* const large_ptr =
+      static_cast<char*>(allocator.root()->Alloc(kDirectMapSize));
+  ASSERT_TRUE(large_ptr);
+
+  // Should not crash.
+  PartitionRoot::CheckMetadataIntegrity(large_ptr);
+  PartitionRoot::CheckMetadataIntegrity(large_ptr + kDirectMapSize - 1);
+
+  allocator.root()->Free(large_ptr);
+}
+
 #if PA_USE_DEATH_TESTS()
 
 // Unit tests that check if an allocation fails in "return null" mode,

src/partition_alloc/partition_root.cc

@@ -1842,17 +1842,17 @@ void PartitionRoot::CheckMetadataIntegrity(const void* ptr) {
     return;
   }
 
-  auto* root = FromAddrInFirstSuperpage(address);
-
   const internal::ReservationOffsetTable& reservation_offset =
-      root->GetReservationOffsetTable();
+      internal::ReservationOffsetTable::Get(address);
   if (reservation_offset.IsManagedByDirectMap(address)) {
     // OOB for direct-mapped allocations is likely immediate crash.
     // No extra benefit from additional checks.
     return;
   }
 
   PA_CHECK(reservation_offset.IsManagedByNormalBuckets(address));
+
+  auto* root = FromAddrInFirstSuperpage(address);
   SlotSpanMetadata* slot_span = SlotSpanMetadata::FromAddr(address, root);
   PA_CHECK(PartitionRoot::FromSlotSpanMetadata(slot_span) == root);