From 23479ae0d3332f5525cfd9491137fc6c0ffcb46a Mon Sep 17 00:00:00 2001
From: Roger McFarlane <rogerm@chromium.org>
Date: Tue, 10 Dec 2024 20:31:23 +0000
Subject: [PATCH] Remove PersistentMemoryAllocator::GetAllocSize()

This CL removes PersistentMemoryAllocator::GetAllocSize() in favor
of allowing various other API entry points to return the alloc size.
This mitigates potential TOCTOU errors where the size of an alloc
is validated by one API then separately fetched in another call. The
size could otherwise be manipulated in between initial validation and
the subsequent fetch.

Bug: 378623799
Change-Id: I8021cf4c07f1a96172deb2a252326e9ffa525798
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/6025612
Reviewed-by: Alexei Svitkine <asvitkine@chromium.org>
Commit-Queue: Roger McFarlane <rogerm@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1394492}
---

diff --git a/base/metrics/field_trial.cc b/base/metrics/field_trial.cc
index 829a5c4b..bf2ffe6 100644
--- a/base/metrics/field_trial.cc
+++ b/base/metrics/field_trial.cc
@@ -124,7 +124,7 @@
 }
 
 // Returns the boundary value for comparing against the FieldTrial's added
-// groups for a given |divisor| (total probability) and |entropy_value|.
+// groups for a given `divisor` (total probability) and `entropy_value`.
 FieldTrial::Probability GetGroupBoundaryValue(
     FieldTrial::Probability divisor,
     double entropy_value) {
@@ -138,7 +138,7 @@
   const double kEpsilon = 1e-8;
   const FieldTrial::Probability result =
       static_cast<FieldTrial::Probability>(divisor * entropy_value + kEpsilon);
-  // Ensure that adding the epsilon still results in a value < |divisor|.
+  // Ensure that adding the epsilon still results in a value < `divisor`.
   return std::min(result, divisor - 1);
 }
 
@@ -259,7 +259,7 @@
   if (forced_) {
     DCHECK(!group_name_.empty());
     if (name == group_name_) {
-      // Note that while |group_| may be equal to |kDefaultGroupNumber| on the
+      // Note that while `group_` may be equal to `kDefaultGroupNumber` on the
       // forced trial, it will not have the same value as the default group
       // number returned from the non-forced |FactoryGetFieldTrial()| call,
       // which takes care to ensure that this does not happen.
@@ -326,7 +326,7 @@
 void FieldTrial::EnableBenchmarking() {
   // We don't need to see field trials created via CreateFieldTrial() for
   // benchmarking, because such field trials have only a single group and are
-  // not affected by randomization that |enable_benchmarking_| would disable.
+  // not affected by randomization that `enable_benchmarking_` would disable.
   DCHECK_EQ(0u, FieldTrialList::GetRandomizedFieldTrialCount());
   enable_benchmarking_ = true;
 }
@@ -453,7 +453,7 @@
   if (group_ != kNotFinalized)
     return;
   accumulated_group_probability_ = divisor_;
-  // Here it's OK to use |kDefaultGroupNumber| since we can't be forced and not
+  // Here it's OK to use `kDefaultGroupNumber` since we can't be forced and not
   // finalized.
   DCHECK(!forced_);
   SetGroupChoice(default_group_name_, kDefaultGroupNumber);
@@ -807,7 +807,7 @@
   field_trial = new FieldTrial(name, kTotalProbability, group_name, 0,
                                is_low_anonymity, is_overridden);
   // The group choice will be finalized in this method. So
-  // |is_randomized_trial| should be false.
+  // `is_randomized_trial` should be false.
   FieldTrialList::Register(field_trial, /*is_randomized_trial=*/false);
   // Force the trial, which will also finalize the group choice.
   field_trial->SetForced();
@@ -910,12 +910,12 @@
   if (!field_trial->ref_)
     return false;
 
+  size_t allocated_size = 0;
   const FieldTrial::FieldTrialEntry* entry =
       global_->field_trial_allocator_->GetAsObject<FieldTrial::FieldTrialEntry>(
-          field_trial->ref_);
+          field_trial->ref_, &allocated_size);
+  CHECK(entry);
 
-  size_t allocated_size =
-      global_->field_trial_allocator_->GetAllocSize(field_trial->ref_);
   uint64_t actual_size =
       sizeof(FieldTrial::FieldTrialEntry) + entry->pickle_size;
   if (allocated_size < actual_size)
diff --git a/base/metrics/persistent_histogram_allocator.cc b/base/metrics/persistent_histogram_allocator.cc
index 825c117..dcc2b2a 100644
--- a/base/metrics/persistent_histogram_allocator.cc
+++ b/base/metrics/persistent_histogram_allocator.cc
@@ -89,13 +89,13 @@
 }
 
 // Calculate the number of bytes required to store all of a histogram's
-// "counts". This will return zero (0) if |bucket_count| is not valid.
+// "counts". This will return zero (0) if `bucket_count` is not valid.
 size_t CalculateRequiredCountsBytes(size_t bucket_count) {
   // 2 because each "sample count" also requires a backup "logged count"
   // used for calculating the delta during snapshot operations.
   const size_t kBytesPerBucket = 2 * sizeof(HistogramBase::AtomicCount);
 
-  // If the |bucket_count| is such that it would overflow the return type,
+  // If the `bucket_count` is such that it would overflow the return type,
   // perhaps as the result of a malicious actor, then return zero to
   // indicate the problem to the caller.
   if (bucket_count > std::numeric_limits<size_t>::max() / kBytesPerBucket)
@@ -176,7 +176,7 @@
 PersistentSparseHistogramDataManager::LoadRecords(
     PersistentSampleMapRecords* sample_map_records,
     std::optional<HistogramBase::Sample> until_value) {
-  // DataManager must be locked in order to access the |sample_records_|
+  // DataManager must be locked in order to access the `sample_records_`
   // vectors.
   base::AutoLock auto_lock(lock_);
 
@@ -222,7 +222,7 @@
   }
 
   // Return all references found that have not yet been seen by
-  // |sample_map_records|, up until |until_value| (if applicable).
+  // `sample_map_records`, up until `until_value` (if applicable).
   std::vector<PersistentMemoryAllocator::Reference> new_references;
   CHECK_GE(found_records.size(), sample_map_records->seen_);
   auto new_found_records =
@@ -230,9 +230,9 @@
   new_references.reserve(new_found_records.size());
   for (const auto& new_record : new_found_records) {
     new_references.push_back(new_record.reference);
-    // Maybe references after |until_value| were found. Stop here immediately in
+    // Maybe references after `until_value` were found. Stop here immediately in
     // such a case, since the caller will not expect any more samples after
-    // |until_value|.
+    // `until_value`.
     if (until_value.has_value() && new_record.value == until_value.value()) {
       break;
     }
@@ -321,9 +321,9 @@
   // count data (while these must reference the persistent counts) and always
   // add it to the local list of known histograms (while these may be simple
   // references to histograms in other processes).
+  size_t length = 0;
   PersistentHistogramData* data =
-      memory_allocator_->GetAsObject<PersistentHistogramData>(ref);
-  const size_t length = memory_allocator_->GetAllocSize(ref);
+      memory_allocator_->GetAsObject<PersistentHistogramData>(ref, &length);
 
   // Check that metadata is reasonable: name is null-terminated and non-empty,
   // ID fields have been loaded with a hash of the name (0 is considered
@@ -331,7 +331,7 @@
   if (!data || data->name[0] == '\0' ||
       reinterpret_cast<char*>(data)[length - 1] != '\0' ||
       data->samples_metadata.id == 0 || data->logged_metadata.id == 0 ||
-      // Note: Sparse histograms use |id + 1| in |logged_metadata|.
+      // Note: Sparse histograms use `id + 1` in `logged_metadata`.
       (data->logged_metadata.id != data->samples_metadata.id &&
        data->logged_metadata.id != data->samples_metadata.id + 1) ||
       // Most non-matching values happen due to truncated names. Ideally, we
@@ -374,7 +374,7 @@
     histogram_data->histogram_type = histogram_type;
     histogram_data->flags = flags | HistogramBase::kIsPersistent;
 
-    // |counts_ref| relies on being zero'd out initially. Even though this
+    // `counts_ref` relies on being zero'd out initially. Even though this
     // should always be the case, manually zero it out again here in case there
     // was memory corruption (e.g. if the memory was mapped from a corrupted
     // spare file).
@@ -388,7 +388,7 @@
     size_t bucket_count = bucket_ranges->bucket_count();
     size_t counts_bytes = CalculateRequiredCountsBytes(bucket_count);
     if (counts_bytes == 0) {
-      // |bucket_count| was out-of-range.
+      // `bucket_count` was out-of-range.
       return nullptr;
     }
 
@@ -396,8 +396,8 @@
     // objects for re-use, it would be dangerous for one to hold a reference
     // from a persistent allocator that is not the global one (which is
     // permanent once set). If this stops being the case, this check can
-    // become an "if" condition beside "!ranges_ref" below and before
-    // set_persistent_reference() farther down.
+    // become an `if` condition beside `!ranges_ref` below and before
+    // `set_persistent_reference()` farther down.
     DCHECK_EQ(this, GlobalHistogramAllocator::Get());
 
     // Re-use an existing BucketRanges persistent allocation if one is known;
@@ -434,7 +434,7 @@
     if (ranges_ref && histogram_data) {
       histogram_data->minimum = minimum;
       histogram_data->maximum = maximum;
-      // |bucket_count| must fit within 32-bits or the allocation of the counts
+      // `bucket_count` must fit within 32-bits or the allocation of the counts

diff --git a/base/metrics/persistent_memory_allocator_unittest.cc b/base/metrics/persistent_memory_allocator_unittest.cc
index 01ac173..a0b96616 100644
--- a/base/metrics/persistent_memory_allocator_unittest.cc
+++ b/base/metrics/persistent_memory_allocator_unittest.cc
@@ -141,11 +141,12 @@
   ASSERT_TRUE(obj1);
   Reference block1 = allocator_->GetAsReference(obj1);
   ASSERT_NE(0U, block1);
-  EXPECT_NE(nullptr, allocator_->GetAsObject<TestObject1>(block1));
   EXPECT_EQ(nullptr, allocator_->GetAsObject<TestObject2>(block1));
-  EXPECT_LE(sizeof(TestObject1), allocator_->GetAllocSize(block1));
-  EXPECT_GT(sizeof(TestObject1) + kAllocAlignment,
-            allocator_->GetAllocSize(block1));
+  size_t alloc_size_1 = 0;
+  EXPECT_NE(nullptr,
+            allocator_->GetAsObject<TestObject1>(block1, &alloc_size_1));
+  EXPECT_LE(sizeof(TestObject1), alloc_size_1);
+  EXPECT_GT(sizeof(TestObject1) + kAllocAlignment, alloc_size_1);
   PersistentMemoryAllocator::MemoryInfo meminfo1;
   allocator_->GetMemoryInfo(&meminfo1);
   EXPECT_EQ(meminfo0.total, meminfo1.total);
@@ -181,11 +182,12 @@
   ASSERT_TRUE(obj2);
   Reference block2 = allocator_->GetAsReference(obj2);
   ASSERT_NE(0U, block2);
-  EXPECT_NE(nullptr, allocator_->GetAsObject<TestObject2>(block2));
   EXPECT_EQ(nullptr, allocator_->GetAsObject<TestObject1>(block2));
-  EXPECT_LE(sizeof(TestObject2), allocator_->GetAllocSize(block2));
-  EXPECT_GT(sizeof(TestObject2) + kAllocAlignment,
-            allocator_->GetAllocSize(block2));
+  size_t alloc_size_2 = 0;
+  EXPECT_NE(nullptr,
+            allocator_->GetAsObject<TestObject2>(block2, &alloc_size_2));
+  EXPECT_LE(sizeof(TestObject2), alloc_size_2);
+  EXPECT_GT(sizeof(TestObject2) + kAllocAlignment, alloc_size_2);
   PersistentMemoryAllocator::MemoryInfo meminfo2;
   allocator_->GetMemoryInfo(&meminfo2);
   EXPECT_EQ(meminfo1.total, meminfo2.total);
@@ -966,10 +968,10 @@
       uint32_t type_id;
       Reference ref;
       while ((ref = iter.GetNext(&type_id)) != 0) {
+        size_t size = 0;
         const char* data = allocator.GetAsArray<char>(
-            ref, 0, PersistentMemoryAllocator::kSizeAny);
+            ref, 0, PersistentMemoryAllocator::kSizeAny, &size);
         uint32_t type = allocator.GetType(ref);
-        size_t size = allocator.GetAllocSize(ref);
         // Ensure compiler can't optimize-out above variables.
         (void)data;
         (void)type;