From 15773228004f6ab61767bcd26dbcc02a28b72d87 Mon Sep 17 00:00:00 2001
From: Joshua Haberman <haberman@google.com>
Date: Fri, 24 Jul 2026 14:10:57 -0700
Subject: [PATCH] Check maximum extension count limit in upb DefBuilder to
 prevent layout_index overflow.

PiperOrigin-RevId: 953548496
---
 upb/reflection/field_def.c        |  5 +++++
 upb/reflection/reflection_test.cc | 30 ++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/upb/reflection/field_def.c b/upb/reflection/field_def.c
index a03a16ad0e77d..ac4cbc7f5a742 100644
--- a/upb/reflection/field_def.c
+++ b/upb/reflection/field_def.c
@@ -751,6 +751,11 @@ static void _upb_FieldDef_CreateExt(
 
   f->scope.extension_scope = m;
   _upb_DefBuilder_Add(ctx, f->full_name, _upb_DefType_Pack(f, UPB_DEFTYPE_EXT));
+  if (ctx->ext_count > UINT16_MAX) {
+    _upb_DefBuilder_Errf(
+        ctx, "too many extensions in file (%d, max %d) for field %s",
+        ctx->ext_count, UINT16_MAX, f->full_name);
+  }
   f->layout_index = ctx->ext_count++;
 
   if (ctx->layout) {
diff --git a/upb/reflection/reflection_test.cc b/upb/reflection/reflection_test.cc
index b64fd6b10354c..cde3e729514fa 100644
--- a/upb/reflection/reflection_test.cc
+++ b/upb/reflection/reflection_test.cc
@@ -389,5 +389,35 @@ TEST(ReflectionTest, ZeroWeakDependencyIndexWithNoDeps) {
   EXPECT_THAT(std::string(status.message()), HasSubstr("out of range"));
 }
 
+TEST(ReflectionTest, TooManyExtensions) {
+  google::protobuf::FileDescriptorProto proto;
+  proto.set_name("test.proto");
+  proto.set_syntax("proto2");
+
+  auto* msg = proto.add_message_type();
+  msg->set_name("TestMessage");
+  auto* range = msg->add_extension_range();
+  range->set_start(1);
+  range->set_end(100000);
+
+  for (int i = 0; i <= 65536; ++i) {
+    auto* ext = proto.add_extension();
+    ext->set_name(absl::StrCat("ext_", i));
+    ext->set_number(i + 1);
+    ext->set_label(google::protobuf::FieldDescriptorProto::LABEL_OPTIONAL);
+    ext->set_type(google::protobuf::FieldDescriptorProto::TYPE_INT32);
+    ext->set_extendee("TestMessage");
+  }
+
+  google::protobuf::FileDescriptorSet set;
+  *set.add_file() = proto;
+  absl::StatusOr<upb::DefPool> pool = LoadDescriptorSetFromProto(set);
+  EXPECT_FALSE(pool.ok());
+  if (!pool.ok()) {
+    EXPECT_THAT(std::string(pool.status().message()),
+                HasSubstr("too many extensions"));
+  }
+}
+
 }  // namespace
 }  // namespace upb_test