From 0c46b3da48a143dc252f3c2dba8c784218b88511 Mon Sep 17 00:00:00 2001
From: Mohamed Elrakad <elrakadm@google.com>
Date: Fri, 16 Aug 2024 16:09:15 +0000
Subject: [PATCH] tp: Fail RegisterAllProtoBuilderFunctions if it attempts to
 register duplicate functions

Bug: 359845211
Change-Id: I0d9f24b2a59f00e4419c706023dfa9a1a3aaecd9
---
 src/trace_processor/trace_processor_impl.cc | 30 +++++++++++++++++----
 src/trace_processor/trace_processor_impl.h  |  1 +
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/src/trace_processor/trace_processor_impl.cc b/src/trace_processor/trace_processor_impl.cc
index f992bf1ec0..9f101bd932 100644
--- a/src/trace_processor/trace_processor_impl.cc
+++ b/src/trace_processor/trace_processor_impl.cc
@@ -144,20 +144,33 @@ void RegisterFunction(PerfettoSqlEngine* engine,
     PERFETTO_ELOG("%s", status.c_message());
 }
 
-void RegisterAllProtoBuilderFunctions(DescriptorPool* pool,
-                                      PerfettoSqlEngine* engine,
-                                      TraceProcessor* tp) {
+base::Status RegisterAllProtoBuilderFunctions(
+    DescriptorPool* pool,
+    std::unordered_map<std::string, std::string>* proto_fn_name_to_path,
+    PerfettoSqlEngine* engine,
+    TraceProcessor* tp) {
   for (uint32_t i = 0; i < pool->descriptors().size(); ++i) {
     // Convert the full name (e.g. .perfetto.protos.TraceMetrics.SubMetric)
     // into a function name of the form (TraceMetrics_SubMetric).
     const auto& desc = pool->descriptors()[i];
     auto fn_name = desc.full_name().substr(desc.package_name().size() + 1);
     std::replace(fn_name.begin(), fn_name.end(), '.', '_');
+    auto registered_fn = proto_fn_name_to_path->find(fn_name);
+    if (registered_fn != proto_fn_name_to_path->end() &&
+        registered_fn->second != desc.full_name()) {
+      return base::ErrStatus(
+          "Attempt to create new metric function '%s' for different descriptor "
+          "'%s' that conflicts with '%s'",
+          fn_name.c_str(), desc.full_name().c_str(),
+          registered_fn->second.c_str());
+    }
     RegisterFunction<metrics::BuildProto>(
         engine, fn_name.c_str(), -1,
         std::make_unique<metrics::BuildProto::Context>(
             metrics::BuildProto::Context{tp, pool, i}));
+    proto_fn_name_to_path->emplace(fn_name, desc.full_name());
   }
+  return base::OkStatus();
 }
 
 void BuildBoundsTable(sqlite3* db, std::pair<int64_t, int64_t> bounds) {
@@ -633,7 +646,8 @@ base::Status TraceProcessorImpl::ExtendMetricsProto(
     size_t size,
     const std::vector<std::string>& skip_prefixes) {
   RETURN_IF_ERROR(pool_.AddFromFileDescriptorSet(data, size, skip_prefixes));
-  RegisterAllProtoBuilderFunctions(&pool_, engine_.get(), this);
+  RETURN_IF_ERROR(RegisterAllProtoBuilderFunctions(
+      &pool_, &proto_fn_name_to_path_, engine_.get(), this));
   return base::OkStatus();
 }
 
@@ -1004,7 +1018,13 @@ void TraceProcessorImpl::InitPerfettoSqlEngine() {
       context_.storage->mutable_string_pool());
 
   // Metrics.
-  RegisterAllProtoBuilderFunctions(&pool_, engine_.get(), this);
+  {
+    auto status = RegisterAllProtoBuilderFunctions(
+        &pool_, &proto_fn_name_to_path_, engine_.get(), this);
+    if (!status.ok()) {
+      PERFETTO_FATAL("%s", status.c_message());
+    }
+  }
 
   // Import prelude module.
   {
diff --git a/src/trace_processor/trace_processor_impl.h b/src/trace_processor/trace_processor_impl.h
index 7b480d206f..68fcafa7e0 100644
--- a/src/trace_processor/trace_processor_impl.h
+++ b/src/trace_processor/trace_processor_impl.h
@@ -118,6 +118,7 @@ class TraceProcessorImpl : public TraceProcessor,
 
   std::vector<metrics::SqlMetricFile> sql_metrics_;
   std::unordered_map<std::string, std::string> proto_field_to_sql_metric_path_;
+  std::unordered_map<std::string, std::string> proto_fn_name_to_path_;
 
   // This is atomic because it is set by the CTRL-C signal handler and we need
   // to prevent single-flow compiler optimizations in ExecuteQuery().