src/tools/tracing_proto_extensions_unittest.cc

/*
 * Copyright (C) 2026 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "src/tools/tracing_proto_extensions.h"

#include "perfetto/base/status.h"
#include "protos/perfetto/common/descriptor.pbzero.h"
#include "src/base/test/tmp_dir_tree.h"
#include "src/base/test/utils.h"
#include "test/gtest_and_gmock.h"

namespace perfetto {
namespace gen_proto_extensions {
namespace {

TEST(GenProtoExtensionsTest, ParseRegistryFileBasic) {
  const char kJson[] = R"({
    "extensions": [
      {
        "scope": "perfetto.protos.TrackEvent",
        "range": [1000, 2000],
        "allocations": [
          {
            "name": "project_a",
            "range": [1000, 1499],
            "contact": "foo@example.com",
            "description": "Project A",
            "proto": "path/to/a.proto"
          },
          {
            "name": "unallocated",
            "range": [1500, 2000]
          }
        ]
      }
    ]
  })";

  auto result = ParseRegistryFile(kJson, "test.json");
  ASSERT_TRUE(result.ok()) << result.status().message();

  ASSERT_EQ(result->size(), 1u);
  const Registry& reg = (*result)[0];
  EXPECT_EQ(reg.scope, "perfetto.protos.TrackEvent");
  ASSERT_EQ(reg.ranges.size(), 1u);
  EXPECT_EQ(reg.ranges[0], Range(1000, 2000));
  ASSERT_EQ(reg.allocations.size(), 2u);

  EXPECT_EQ(reg.allocations[0].name, "project_a");
  ASSERT_EQ(reg.allocations[0].ranges.size(), 1u);
  EXPECT_EQ(reg.allocations[0].ranges[0], Range(1000, 1499));
  EXPECT_EQ(reg.allocations[0].contact, "foo@example.com");
  EXPECT_EQ(reg.allocations[0].proto, "path/to/a.proto");

  EXPECT_EQ(reg.allocations[1].name, "unallocated");
  ASSERT_EQ(reg.allocations[1].ranges.size(), 1u);
  EXPECT_EQ(reg.allocations[1].ranges[0], Range(1500, 2000));
}

TEST(GenProtoExtensionsTest, ParseRegistryFileWithSubRegistry) {
  const char kJson[] = R"({
    "extensions": [
      {
        "range": [1000, 2000],
        "allocations": [
          {
            "name": "project_a",
            "range": [1000, 1499],
            "registry": "path/to/sub.json"
          },
          {
            "name": "project_b",
            "range": [1500, 1999],
            "repo": "https://example.com/repo",
            "proto": "some/path.proto"
          },
          {
            "name": "unallocated",
            "range": [2000, 2000]
          }
        ]
      }
    ]
  })";

  auto result = ParseRegistryFile(kJson, "test.json");
  ASSERT_TRUE(result.ok()) << result.status().message();
  ASSERT_EQ(result->size(), 1u);
  EXPECT_EQ((*result)[0].allocations[0].registry, "path/to/sub.json");
  EXPECT_EQ((*result)[0].allocations[1].repo, "https://example.com/repo");
}

TEST(GenProtoExtensionsTest, ParseRegistryFileMultipleExtensions) {
  const char kJson[] = R"({
    "extensions": [
      {
        "scope": "perfetto.protos.TrackEvent",
        "range": [1000, 1999],
        "allocations": [
          {"name": "a", "range": [1000, 1999], "proto": "a.proto"}
        ]
      },
      {
        "scope": "perfetto.protos.TrackEvent",
        "range": [2000, 2999],
        "allocations": [
          {"name": "b", "range": [2000, 2999], "proto": "b.proto"}
        ]
      }
    ]
  })";

  auto result = ParseRegistryFile(kJson, "test.json");
  ASSERT_TRUE(result.ok()) << result.status().message();
  ASSERT_EQ(result->size(), 2u);
  EXPECT_EQ((*result)[0].scope, "perfetto.protos.TrackEvent");
  EXPECT_EQ((*result)[1].scope, "perfetto.protos.TrackEvent");
  EXPECT_EQ((*result)[0].allocations[0].name, "a");
  EXPECT_EQ((*result)[1].allocations[0].name, "b");
}

TEST(GenProtoExtensionsTest, ParseRegistryFileWithComment) {
  const char kJson[] = R"({
    "comment": ["Top-level comment"],
    "extensions": [
      {
        "scope": "perfetto.protos.TrackEvent",
        "range": [100, 199],
        "allocations": [
          {"name": "a", "range": [100, 199], "proto": "a.proto"}
        ]
      }
    ]
  })";

  auto result = ParseRegistryFile(kJson, "test.json");
  ASSERT_TRUE(result.ok()) << result.status().message();
  ASSERT_EQ(result->size(), 1u);
}

TEST(GenProtoExtensionsTest, ParseRegistryFileUnknownField) {
  const char kJson[] = R"({
    "extensions": [],
    "unknown_field": 42
  })";
  auto result = ParseRegistryFile(kJson, "test.json");
  EXPECT_FALSE(result.ok());
  EXPECT_THAT(result.status().message(), testing::HasSubstr("Unknown field"));
}

TEST(GenProtoExtensionsTest, ParseRegistryFileInvalidJson) {
  auto result = ParseRegistryFile("{invalid", "test.json");
  EXPECT_FALSE(result.ok());
}

TEST(GenProtoExtensionsTest, ParseRegistryFileWithRanges) {
  const char kJson[] = R"({
    "extensions": [
      {
        "ranges": [[1000, 1499], [2000, 2999]],
        "allocations": [
          {
            "name": "project_a",
            "range": [1000, 1499],
            "proto": "a.proto"
          },
          {
            "name": "project_b",
            "range": [2000, 2999],
            "proto": "b.proto"
          }
        ]
      }
    ]
  })";

  auto result = ParseRegistryFile(kJson, "test.json");
  ASSERT_TRUE(result.ok()) << result.status().message();

  ASSERT_EQ(result->size(), 1u);
  ASSERT_EQ((*result)[0].ranges.size(), 2u);
  EXPECT_EQ((*result)[0].ranges[0], Range(1000, 1499));
  EXPECT_EQ((*result)[0].ranges[1], Range(2000, 2999));
}

TEST(GenProtoExtensionsTest, ParseRegistryFileAllocWithRanges) {
  const char kJson[] = R"({
    "extensions": [
      {
        "range": [1000, 2999],
        "allocations": [
          {
            "name": "project_a",
            "ranges": [[1000, 1499], [2000, 2499]],
            "proto": "a.proto"
          },
          {
            "name": "unallocated",
            "ranges": [[1500, 1999], [2500, 2999]]
          }
        ]
      }
    ]
  })";

  auto result = ParseRegistryFile(kJson, "test.json");
  ASSERT_TRUE(result.ok()) << result.status().message();

  ASSERT_EQ(result->size(), 1u);
  ASSERT_EQ((*result)[0].allocations[0].ranges.size(), 2u);
  EXPECT_EQ((*result)[0].allocations[0].ranges[0], Range(1000, 1499));
  EXPECT_EQ((*result)[0].allocations[0].ranges[1], Range(2000, 2499));
}

TEST(GenProtoExtensionsTest, ParseRegistryFileRangeAndRangesMutuallyExclusive) {
  const char kJson[] = R"({
    "extensions": [
      {
        "range": [1000, 2000],
        "ranges": [[1000, 1500], [1501, 2000]],
        "allocations": []
      }
    ]
  })";

  auto result = ParseRegistryFile(kJson, "test.json");
  EXPECT_FALSE(result.ok());
  EXPECT_THAT(result.status().message(), testing::HasSubstr("both"));
}

TEST(GenProtoExtensionsTest, ParseRegistryFileMissingRange) {
  const char kJson[] = R"({
    "extensions": [
      {
        "allocations": [
          {"name": "a", "range": [1, 10], "proto": "a.proto"}
        ]
      }
    ]
  })";
  auto result = ParseRegistryFile(kJson, "test.json");
  ASSERT_TRUE(result.ok());
  ASSERT_EQ(result->size(), 1u);
  // ranges will be empty (no range specified in the entry).
  // ValidateRegistry should catch this.
  auto status = ValidateRegistry((*result)[0]);
  EXPECT_FALSE(status.ok());
}

TEST(GenProtoExtensionsTest, ValidateRegistryValid) {
  Registry reg;
  reg.source_path = "test.json";
  reg.scope = "perfetto.protos.TrackEvent";
  reg.ranges = {{100, 299}};
  reg.allocations.push_back({"a", {{100, 199}}, "", "", "", "a.proto", ""});
  reg.allocations.push_back({"unallocated", {{200, 299}}, "", "", "", "", ""});

  EXPECT_TRUE(ValidateRegistry(reg).ok());
}

TEST(GenProtoExtensionsTest, ValidateRegistryGap) {
  Registry reg;
  reg.source_path = "test.json";
  reg.scope = "perfetto.protos.TrackEvent";
  reg.ranges = {{100, 299}};
  // Gap between 150 and 200.
  reg.allocations.push_back({"a", {{100, 149}}, "", "", "", "a.proto", ""});
  reg.allocations.push_back({"unallocated", {{200, 299}}, "", "", "", "", ""});

  auto status = ValidateRegistry(reg);
  EXPECT_FALSE(status.ok());
  EXPECT_THAT(status.message(), testing::HasSubstr("gap or overlap"));
}

TEST(GenProtoExtensionsTest, ValidateRegistryOverflow) {
  Registry reg;
  reg.source_path = "test.json";
  reg.scope = "perfetto.protos.TrackEvent";
  reg.ranges = {{100, 199}};
  // Allocation extends beyond registry range.
  reg.allocations.push_back({"a", {{100, 250}}, "", "", "", "a.proto", ""});

  auto status = ValidateRegistry(reg);
  EXPECT_FALSE(status.ok());
  EXPECT_THAT(status.message(), testing::HasSubstr("gap or overlap"));
}

TEST(GenProtoExtensionsTest, ValidateRegistryBothProtoAndRegistry) {
  Registry reg;
  reg.source_path = "test.json";
  reg.scope = "perfetto.protos.TrackEvent";
  reg.ranges = {{100, 199}};
  reg.allocations.push_back(
      {"a", {{100, 199}}, "", "", "", "a.proto", "a.json"});

  auto status = ValidateRegistry(reg);
  EXPECT_FALSE(status.ok());
  EXPECT_THAT(status.message(), testing::HasSubstr("both"));
}

TEST(GenProtoExtensionsTest, ValidateRegistryMissingProtoOrRegistry) {
  Registry reg;
  reg.source_path = "test.json";
  reg.scope = "perfetto.protos.TrackEvent";
  reg.ranges = {{100, 199}};
  // Non-unallocated entry with neither proto nor registry nor repo.
  reg.allocations.push_back({"a", {{100, 199}}, "", "", "", "", ""});

  auto status = ValidateRegistry(reg);
  EXPECT_FALSE(status.ok());
  EXPECT_THAT(status.message(), testing::HasSubstr("must have"));
}

TEST(GenProtoExtensionsTest, ValidateRegistryUnallocatedWithProto) {
  Registry reg;
  reg.source_path = "test.json";
  reg.scope = "perfetto.protos.TrackEvent";
  reg.ranges = {{100, 199}};
  reg.allocations.push_back(
      {"unallocated", {{100, 199}}, "", "", "", "bad.proto", ""});

  auto status = ValidateRegistry(reg);
  EXPECT_FALSE(status.ok());
  EXPECT_THAT(status.message(), testing::HasSubstr("Unallocated"));
}

TEST(GenProtoExtensionsTest, ValidateRegistryRemoteEntrySkipsProtoCheck) {
  Registry reg;
  reg.source_path = "test.json";
  reg.scope = "perfetto.protos.TrackEvent";
  reg.ranges = {{100, 299}};
  // Remote entry: has repo but no local proto or registry - should be fine.
  reg.allocations.push_back(
      {"remote", {{100, 199}}, "", "", "https://example.com", "ext.proto", ""});
  reg.allocations.push_back({"unallocated", {{200, 299}}, "", "", "", "", ""});

  EXPECT_TRUE(ValidateRegistry(reg).ok());
}

TEST(GenProtoExtensionsTest, ValidateRegistryScatteredRangesValid) {
  Registry reg;
  reg.source_path = "test.json";
  reg.scope = "perfetto.protos.TrackEvent";
  reg.ranges = {{100, 199}, {300, 399}};
  reg.allocations.push_back(
      {"a", {{100, 199}, {300, 349}}, "", "", "", "a.proto", ""});
  reg.allocations.push_back({"unallocated", {{350, 399}}, "", "", "", "", ""});

  EXPECT_TRUE(ValidateRegistry(reg).ok());
}

TEST(GenProtoExtensionsTest, ValidateRegistryScatteredRangesGap) {
  Registry reg;
  reg.source_path = "test.json";
  reg.scope = "perfetto.protos.TrackEvent";
  reg.ranges = {{100, 199}, {300, 399}};
  // Missing [150, 199] from the parent range - gap.
  reg.allocations.push_back(
      {"a", {{100, 149}, {300, 399}}, "", "", "", "a.proto", ""});

  auto status = ValidateRegistry(reg);
  EXPECT_FALSE(status.ok());
  EXPECT_THAT(status.message(), testing::HasSubstr("gap or overlap"));
}

TEST(GenProtoExtensionsTest, ValidateRegistryScatteredRangesOverlap) {
  Registry reg;
  reg.source_path = "test.json";
  reg.scope = "perfetto.protos.TrackEvent";
  reg.ranges = {{100, 399}};
  // Two allocations claim overlapping ranges.
  reg.allocations.push_back({"a", {{100, 250}}, "", "", "", "a.proto", ""});
  reg.allocations.push_back({"b", {{200, 399}}, "", "", "", "b.proto", ""});

  auto status = ValidateRegistry(reg);
  EXPECT_FALSE(status.ok());
  EXPECT_THAT(status.message(), testing::HasSubstr("overlap"));
}

TEST(GenProtoExtensionsTest, ValidateRegistryMissingScope) {
  Registry reg;
  reg.source_path = "test.json";
  reg.ranges = {{100, 199}};
  reg.allocations.push_back({"a", {{100, 199}}, "", "", "", "a.proto", ""});

  auto status = ValidateRegistry(reg);
  EXPECT_FALSE(status.ok());
  EXPECT_THAT(status.message(), testing::HasSubstr("scope"));
}

TEST(GenProtoExtensionsTest, ValidateRegistryWrongScope) {
  Registry reg;
  reg.source_path = "test.json";
  // Scopes other than TrackEvent / TracePacket / InternedData are rejected.
  reg.scope = "perfetto.protos.SomeOtherType";
  reg.ranges = {{100, 199}};
  reg.allocations.push_back({"a", {{100, 199}}, "", "", "", "a.proto", ""});

  auto status = ValidateRegistry(reg);
  EXPECT_FALSE(status.ok());
  EXPECT_THAT(status.message(), testing::HasSubstr("scope"));
}

TEST(GenProtoExtensionsTest, ValidateRegistryAcceptsTracePacketScope) {
  Registry reg;
  reg.source_path = "test.json";
  reg.scope = "perfetto.protos.TracePacket";
  reg.ranges = {{100, 199}};
  reg.allocations.push_back({"a", {{100, 199}}, "", "", "", "a.proto", ""});

  EXPECT_TRUE(ValidateRegistry(reg).ok());
}

TEST(GenProtoExtensionsTest, ValidateRegistryAcceptsInternedDataScope) {
  Registry reg;
  reg.source_path = "test.json";
  reg.scope = "perfetto.protos.InternedData";
  reg.ranges = {{100, 199}};
  reg.allocations.push_back({"a", {{100, 199}}, "", "", "", "a.proto", ""});

  EXPECT_TRUE(ValidateRegistry(reg).ok());
}

TEST(GenProtoExtensionsTest, ValidateScopesUniqueAccepts) {
  Registry track_event;
  track_event.source_path = "test.json";
  track_event.scope = "perfetto.protos.TrackEvent";
  Registry trace_packet;
  trace_packet.source_path = "test.json";
  trace_packet.scope = "perfetto.protos.TracePacket";
  Registry interned_data;
  interned_data.source_path = "test.json";
  interned_data.scope = "perfetto.protos.InternedData";

  EXPECT_TRUE(
      ValidateScopesUnique({track_event, trace_packet, interned_data}).ok());
}

TEST(GenProtoExtensionsTest, ValidateScopesUniqueRejectsDuplicates) {
  Registry first;
  first.source_path = "test.json";
  first.scope = "perfetto.protos.TrackEvent";
  Registry second;
  second.source_path = "test.json";
  second.scope = "perfetto.protos.TrackEvent";

  auto status = ValidateScopesUnique({first, second});
  EXPECT_FALSE(status.ok());
  EXPECT_THAT(status.message(),
              testing::HasSubstr("'perfetto.protos.TrackEvent'"));
}

TEST(GenProtoExtensionsTest, GenerateExtensionDescriptorsNoExtend) {
  // A proto file that compiles but has no "extend TrackEvent" should be
  // rejected.
  base::TmpDirTree tmp;
  tmp.AddDir("protos");
  tmp.AddDir("protos/perfetto");
  tmp.AddDir("protos/perfetto/trace");
  tmp.AddDir("protos/perfetto/trace/track_event");
  tmp.AddFile("protos/perfetto/trace/track_event/no_extend.proto", R"(
    syntax = "proto2";
    package test;
    message Foo {
      optional int32 x = 1;
    }
  )");
  tmp.AddFile("registry.json", R"({
    "extensions": [
      {
        "scope": "perfetto.protos.TrackEvent",
        "range": [9900, 9999],
        "allocations": [
          {
            "name": "test",
            "range": [9900, 9999],
            "proto": "protos/perfetto/trace/track_event/no_extend.proto"
          }
        ]
      }
    ]
  })");

  auto result = GenerateExtensionDescriptors(tmp.AbsolutePath("registry.json"),
                                             {tmp.path()}, tmp.path());
  ASSERT_FALSE(result.ok());
  EXPECT_THAT(result.status().message(),
              testing::HasSubstr("no extensions targeting"));
}

TEST(GenProtoExtensionsTest, GenerateExtensionDescriptorsWithUnifiedRegistry) {
  // End-to-end smoke test for GenerateExtensionDescriptors: feed it the
  // real extensions.json and let it compile the in-repo protos it points at.
  // Remote entries (chromium / android-internal) are skipped automatically.
  //
  // Both the -I include path and the registry root_dir need to be the
  // perfetto repo root. Compute it from the JSON's resolved path so the
  // test works regardless of cwd.
  constexpr char kRel[] = "protos/perfetto/trace/extensions.json";
  std::string json_path = base::GetTestDataPath(kRel);
  ASSERT_GE(json_path.size(), std::char_traits<char>::length(kRel));
  std::string repo_root = json_path.substr(
      0, json_path.size() - std::char_traits<char>::length(kRel));
  if (!repo_root.empty() && repo_root.back() == '/') {
    repo_root.pop_back();
  }
  if (repo_root.empty()) {
    repo_root = ".";
  }

  auto result = GenerateExtensionDescriptors(json_path, {repo_root}, repo_root);
  ASSERT_TRUE(result.ok()) << result.status().message();
  // The output should be a non-empty FileDescriptorSet.
  EXPECT_GT(result->size(), 0u);
}

TEST(GenProtoExtensionsTest, GenerateExtensionDescriptorsExcludesCoreProtos) {
  // The output carries the extension files and the types they define, but not
  // the core Perfetto protos (the extendee and what it imports). Here a core
  // extendee references a leaf type in a separate core file; both core files
  // must be excluded from the output.
  base::TmpDirTree tmp;
  tmp.AddDir("protos");
  tmp.AddDir("protos/perfetto");
  tmp.AddDir("protos/perfetto/trace");
  tmp.AddDir("protos/perfetto/trace/track_event");
  tmp.AddDir("ext");
  // A leaf type in its own core file.
  tmp.AddFile("protos/perfetto/trace/track_event/leaf.proto", R"(
    syntax = "proto2";
    package perfetto.protos;
    message CoreLeaf { optional int32 x = 1; }
  )");
  // A core extendee that references the leaf from the separate file.
  tmp.AddFile("protos/perfetto/trace/track_event/track_event.proto", R"(
    syntax = "proto2";
    package perfetto.protos;
    import "protos/perfetto/trace/track_event/leaf.proto";
    message TrackEvent {
      extensions 9900 to 9999;
      optional CoreLeaf core_leaf = 1;
    }
  )");
  // Out-of-tree extension: its own payload, extending the core TrackEvent.
  tmp.AddFile("ext/my_ext.proto", R"(
    syntax = "proto2";
    package com.android.internal;
    import "protos/perfetto/trace/track_event/track_event.proto";
    message MyPayload { optional int32 y = 1; }
    message MyExt {
      extend perfetto.protos.TrackEvent {
        optional MyPayload my_payload = 9900;
      }
    }
  )");
  tmp.AddFile("registry.json", R"({
    "extensions": [
      {
        "scope": "perfetto.protos.TrackEvent",
        "range": [9900, 9999],
        "allocations": [
          {"name": "ext", "range": [9900, 9999], "proto": "ext/my_ext.proto"}
        ]
      }
    ]
  })");

  auto result = GenerateExtensionDescriptors(tmp.AbsolutePath("registry.json"),
                                             {tmp.path()}, tmp.path());
  ASSERT_TRUE(result.ok()) << result.status().message();

  std::vector<std::string> files;
  protos::pbzero::FileDescriptorSet::Decoder fds(result->data(),
                                                 result->size());
  for (auto it = fds.file(); it; ++it) {
    protos::pbzero::FileDescriptorProto::Decoder f(*it);
    files.push_back(f.name().ToStdString());
  }

  // Extension file is present; the core extendee and core leaf are not.
  EXPECT_THAT(files, testing::Contains("ext/my_ext.proto"));
  EXPECT_THAT(files,
              testing::Not(testing::Contains(
                  "protos/perfetto/trace/track_event/track_event.proto")));
  EXPECT_THAT(files, testing::Not(testing::Contains(
                         "protos/perfetto/trace/track_event/leaf.proto")));
}

}  // namespace
}  // namespace gen_proto_extensions
}  // namespace perfetto